commit 3165c4acf86e61a4c5cd90e271a81f5f3ab5857a Author: wehub-resource-sync Date: Mon Jul 13 12:49:10 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.github/.codedev.yml b/.github/.codedev.yml new file mode 100644 index 0000000..ccbcfbd --- /dev/null +++ b/.github/.codedev.yml @@ -0,0 +1,31 @@ +coverage: + status: + project: #add everything under here, more options at https://docs.codecov.com/docs/commit-status + default: # default is the status check's name, not default settings + target: auto #default + threshold: 1% #allow coverage to drop by 1% + base: auto + if_ci_failed: error #success, failure, error, ignore + patch: + default: + target: 82% #default + threshold: 1% #allow coverage to drop by 1% + base: auto + if_ci_failed: error #success, failure, error, ignore + +comment: #this is a top-level key + layout: " diff, flags, files" + behavior: default + require_changes: false # if true: only post the comment if coverage changes + require_base: false # [true :: must have a base report to post] + require_head: true # [true :: must have a head report to post] + hide_project_coverage: false # [true :: only show coverage on the git diff aka patch coverage] + +# sample regex patterns +ignore: + - "tests" + - "examples/" + - "mock/" + - "callbacks/interface.go" + - "utils/safe" + - "components/tool/utils/create_options.go" \ No newline at end of file diff --git a/.github/.commit-rules.json b/.github/.commit-rules.json new file mode 100644 index 0000000..3bc6176 --- /dev/null +++ b/.github/.commit-rules.json @@ -0,0 +1,41 @@ +{ + "allowedTypes": [ + "feat", + "fix", + "docs", + "style", + "refactor", + "perf", + "test", + "build", + "ci", + "chore", + "revert" + ], + "allowedScopes": [ + "adk", + "adk/filesystem", + "callbacks", + "components", + "compose", + "deep", + "dynamictool", + "filesystem", + "flow", + "internal", + "middlewares", + "planexecute", + "plantask", + "prebuilt", + "reduction", + "schema", + "skill", + "summarization", + "supervisor", + "toolsearch", + "utils", + "docs", + "ci", + "serialization" + ] +} diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..28bb067 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,40 @@ +--- +name: Bug report +about: Create a report to help us improve +title: '' +labels: '' +assignees: '' + +--- + +**Describe the bug** + +A clear and concise description of what the bug is. + +**To Reproduce** + +Steps to reproduce the behavior: +1. Go to '...' +2. Click on '....' +3. Scroll down to '....' +4. See error + +**Expected behavior** + +A clear and concise description of what you expected to happen. + +**Screenshots** + +If applicable, add screenshots to help explain your problem. + +**Version:** + +Please provide the version of {project_name} you are using. + +**Environment:** + +The output of `go env`. + +**Additional context** + +Add any other context about the problem here. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..f0291e0 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,24 @@ +--- +name: Feature request +about: Suggest an idea for this project +title: '' +labels: '' +assignees: '' + +--- + +**Is your feature request related to a problem? Please describe.** + +A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] + +**Describe the solution you'd like** + +A clear and concise description of what you want to happen. + +**Describe alternatives you've considered** + +A clear and concise description of any alternative solutions or features you've considered. + +**Additional context** + +Add any other context or screenshots about the feature request here. diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..3f6e529 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,49 @@ +#### What type of PR is this? + + +#### Check the PR title. + +- [ ] This PR title match the format: \(optional scope): \ +- [ ] The description of this PR title is user-oriented and clear enough for others to understand. +- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io) + + +#### (Optional) Translate the PR title into Chinese. + + +#### (Optional) More detailed description for this PR(en: English/zh: Chinese). + +en: +zh(optional): + + +#### (Optional) Which issue(s) this PR fixes: + + +#### (optional) The PR that updates user documentation: + diff --git a/.github/static/img/eino/chain.png b/.github/static/img/eino/chain.png new file mode 100644 index 0000000..91632f3 Binary files /dev/null and b/.github/static/img/eino/chain.png differ diff --git a/.github/static/img/eino/chain_append_branch.png b/.github/static/img/eino/chain_append_branch.png new file mode 100644 index 0000000..070ab14 Binary files /dev/null and b/.github/static/img/eino/chain_append_branch.png differ diff --git a/.github/static/img/eino/chain_append_parallel.png b/.github/static/img/eino/chain_append_parallel.png new file mode 100644 index 0000000..0405f01 Binary files /dev/null and b/.github/static/img/eino/chain_append_parallel.png differ diff --git a/.github/static/img/eino/eino_architecture_overview.png b/.github/static/img/eino/eino_architecture_overview.png new file mode 100644 index 0000000..f2667c9 Binary files /dev/null and b/.github/static/img/eino/eino_architecture_overview.png differ diff --git a/.github/static/img/eino/eino_project_structure_and_modules.png b/.github/static/img/eino/eino_project_structure_and_modules.png new file mode 100644 index 0000000..78a94cd Binary files /dev/null and b/.github/static/img/eino/eino_project_structure_and_modules.png differ diff --git a/.github/static/img/eino/graph.gif b/.github/static/img/eino/graph.gif new file mode 100644 index 0000000..768f769 Binary files /dev/null and b/.github/static/img/eino/graph.gif differ diff --git a/.github/static/img/eino/graph.png b/.github/static/img/eino/graph.png new file mode 100644 index 0000000..fe2f05c Binary files /dev/null and b/.github/static/img/eino/graph.png differ diff --git a/.github/static/img/eino/graph_add_branch.png b/.github/static/img/eino/graph_add_branch.png new file mode 100644 index 0000000..cd21a45 Binary files /dev/null and b/.github/static/img/eino/graph_add_branch.png differ diff --git a/.github/static/img/eino/graph_add_edge.png b/.github/static/img/eino/graph_add_edge.png new file mode 100644 index 0000000..3c56661 Binary files /dev/null and b/.github/static/img/eino/graph_add_edge.png differ diff --git a/.github/static/img/eino/graph_parallel.png b/.github/static/img/eino/graph_parallel.png new file mode 100644 index 0000000..0dae772 Binary files /dev/null and b/.github/static/img/eino/graph_parallel.png differ diff --git a/.github/static/img/eino/lark_group_en.png b/.github/static/img/eino/lark_group_en.png new file mode 100644 index 0000000..994eeb3 Binary files /dev/null and b/.github/static/img/eino/lark_group_en.png differ diff --git a/.github/static/img/eino/lark_group_zh.png b/.github/static/img/eino/lark_group_zh.png new file mode 100644 index 0000000..4c9b5ee Binary files /dev/null and b/.github/static/img/eino/lark_group_zh.png differ diff --git a/.github/static/img/eino/react.png b/.github/static/img/eino/react.png new file mode 100644 index 0000000..a92ada1 Binary files /dev/null and b/.github/static/img/eino/react.png differ diff --git a/.github/static/img/eino/simple_chain.png b/.github/static/img/eino/simple_chain.png new file mode 100644 index 0000000..cc49a7b Binary files /dev/null and b/.github/static/img/eino/simple_chain.png differ diff --git a/.github/static/img/eino/simple_workflow.png b/.github/static/img/eino/simple_workflow.png new file mode 100644 index 0000000..e74afd2 Binary files /dev/null and b/.github/static/img/eino/simple_workflow.png differ diff --git a/.github/static/img/eino/stream.png b/.github/static/img/eino/stream.png new file mode 100644 index 0000000..17f718d Binary files /dev/null and b/.github/static/img/eino/stream.png differ diff --git a/.github/static/img/eino/tool_call_graph.png b/.github/static/img/eino/tool_call_graph.png new file mode 100644 index 0000000..2fd8c9a Binary files /dev/null and b/.github/static/img/eino/tool_call_graph.png differ diff --git a/.github/workflows/pr-check.yml b/.github/workflows/pr-check.yml new file mode 100644 index 0000000..3f5bbac --- /dev/null +++ b/.github/workflows/pr-check.yml @@ -0,0 +1,156 @@ +name: Pull Request Check + +on: [ pull_request ] + +jobs: + compliant: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Check License Header + uses: apache/skywalking-eyes/header@v0.4.0 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + - name: Check Spell + uses: crate-ci/typos@v1.42.3 + + golangci-lint: + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + repository-projects: write + steps: + - uses: actions/checkout@v4 + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: 1.18 + # for self-hosted, the cache path is shared across projects, + # and it works well without the cache of GitHub actions + # Enable it if we're going to use GitHub only + cache: true + + - name: Golang CI Lint + # https://golangci-lint.run/ + uses: golangci/golangci-lint-action@v9.2.0 + with: + version: v2.8.0 + args: --timeout 5m + + commit-msg-check: + name: Commit Message Check + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: read + steps: + - uses: actions/checkout@v4 + - name: Validate commit messages format and scope + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const fs = require('fs'); + const pr = context.payload.pull_request; + if (!pr) { + core.setFailed('This workflow must run on pull_request events.'); + return; + } + + let allowedTypes = []; + let allowedScopes = []; + try { + const raw = fs.readFileSync('.github/.commit-rules.json', 'utf8'); + const cfg = JSON.parse(raw); + allowedTypes = Array.isArray(cfg.allowedTypes) ? cfg.allowedTypes : []; + allowedScopes = Array.isArray(cfg.allowedScopes) ? cfg.allowedScopes : []; + } catch (e) { + core.setFailed('Cannot read .github/.commit-rules.json: ' + e.message); + return; + } + if (!allowedTypes.length) { + core.setFailed('allowedTypes is empty in .github/.commit-rules.json'); + return; + } + + const { owner, repo } = context.repo; + const pull_number = pr.number; + const commits = await github.paginate( + github.rest.pulls.listCommits, + { owner, repo, pull_number, per_page: 100 } + ); + + let errors = []; + + for (const c of commits) { + const sha = c.sha.slice(0, 7); + const subject = (c.commit.message || '').split('\n')[0]; + const m = subject.match(/^([a-z]+)(\(([a-z0-9\-\/]+)\))?:\s(.+)$/); + if (!m) { + errors.push(`(${sha}) invalid format: "${subject}"`); + continue; + } + + const type = m[1]; + const scope = m[3]; // may be undefined + const desc = m[4]; + + if (!allowedTypes.includes(type)) { + errors.push(`(${sha}) invalid type "${type}"`); + } + + if (!desc || !desc.trim()) { + errors.push(`(${sha}) description must be non-empty`); + } + + if (scope) { + const topScope = scope.split('/')[0]; + if (allowedScopes.length && !allowedScopes.includes(topScope)) { + errors.push(`(${sha}) invalid scope "${scope}"`); + } + } + } + + if (errors.length) { + core.setFailed('Commit message check failed:\n' + errors.join('\n')); + } else { + core.info('All commit messages conform to "(optional scope): " and scope rules.'); + } + + pr-title-check: + name: PR Title Check + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: read + steps: + - uses: actions/checkout@v4 + - name: Read commit rules + id: rules + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const fs = require('fs'); + let cfg; + try { + const raw = fs.readFileSync('.github/.commit-rules.json', 'utf8'); + cfg = JSON.parse(raw); + } catch (e) { + core.setFailed('Cannot read .github/.commit-rules.json: ' + e.message); + return; + } + const toMultiline = (list) => Array.isArray(list) ? list.join('\n') : ''; + core.setOutput('types', toMultiline(cfg.allowedTypes)); + core.setOutput('scopes', toMultiline(cfg.allowedScopes)); + - name: Validate PR title + uses: amannn/action-semantic-pull-request@v6.1.1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + types: ${{ steps.rules.outputs.types }} + scopes: ${{ steps.rules.outputs.scopes }} + requireScope: false diff --git a/.github/workflows/tag-notification.yml b/.github/workflows/tag-notification.yml new file mode 100644 index 0000000..2b7ea3e --- /dev/null +++ b/.github/workflows/tag-notification.yml @@ -0,0 +1,120 @@ +name: Tag Notification + +on: + push: + tags: + - 'v*' + +jobs: + notify: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + fetch-tags: true + + - name: Fetch tag info + run: | + git fetch --tags -f + + - name: Get tag info and send notification + run: | + # Get the tag name + TAG_NAME="${{ github.ref_name }}" + echo "Processing tag: $TAG_NAME" + + # Get tag message + echo "Getting tag message..." + TAG_MESSAGE=$(git tag -l --format='%(contents)' "$TAG_NAME") + echo "Tag message:" + echo "$TAG_MESSAGE" + echo "---" + + # Create base content parts + HEADER="### 🏷️ Eino New Tag Created: \`$TAG_NAME\`" + VERSION_INFO="📦 Version: \`$TAG_NAME\`" + + # Prepare the message parts for jq + if [ ! -z "$TAG_MESSAGE" ]; then + # Pass all parts to jq and let it handle the formatting + jq -n \ + --arg header "$HEADER" \ + --arg version "$VERSION_INFO" \ + --arg notes "$TAG_MESSAGE" \ + --arg repo_url "https://github.com/${{ github.repository }}/releases/tag/$TAG_NAME" \ + '{ + "msg_type": "interactive", + "card": { + "elements": [ + { + "tag": "markdown", + "content": ($header + "\n\n" + $version + "\n\n### 📝 Release Notes:\n" + $notes) + }, + { + "tag": "action", + "actions": [ + { + "tag": "button", + "text": { + "tag": "plain_text", + "content": "🔗 View Tag" + }, + "url": $repo_url, + "type": "default" + } + ] + } + ], + "header": { + "title": { + "tag": "plain_text", + "content": "🏷️ Eino New Tag Created" + } + } + } + }' > webhook_payload.json + else + # Without release notes + jq -n \ + --arg header "$HEADER" \ + --arg version "$VERSION_INFO" \ + --arg repo_url "https://github.com/${{ github.repository }}/releases/tag/$TAG_NAME" \ + '{ + "msg_type": "interactive", + "card": { + "elements": [ + { + "tag": "markdown", + "content": ($header + "\n\n" + $version) + }, + { + "tag": "action", + "actions": [ + { + "tag": "button", + "text": { + "tag": "plain_text", + "content": "🔗 View Tag" + }, + "url": $repo_url, + "type": "default" + } + ] + } + ], + "header": { + "title": { + "tag": "plain_text", + "content": "🏷️ Eino New Tag Created" + } + } + } + }' > webhook_payload.json + fi + + # Send webhook + curl -X POST \ + -H "Content-Type: application/json" \ + -d @webhook_payload.json \ + "${{ secrets.FEISHU_WEBHOOK_URL }}" \ No newline at end of file diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..7748257 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,262 @@ +name: Eino Tests + +on: + pull_request: + push: + branches: + - main + +env: + DEFAULT_GO_VERSION: "1.18" + +jobs: + unit-test: + name: eino-unit-test + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + repository-projects: write + env: + COVERAGE_FILE: coverage.out + BREAKDOWN_FILE: main.breakdown + + steps: + - uses: actions/checkout@v4 + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: ${{ env.DEFAULT_GO_VERSION }} + + - name: Exec Go Test + run: | + modules=`find . -name "go.mod" -exec dirname {} \;` + echo $modules + list="" + coverpkg="" + if [[ ! -f "go.work" ]];then go work init;fi + for module in $modules; do go work use $module; list=$module"/... "$list; coverpkg=$module"/...,"$coverpkg; done + go work sync + go test -race -v -coverprofile=${{ env.COVERAGE_FILE }} -gcflags="all=-l -N" -coverpkg=$coverpkg $list + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v5 + with: + name: eino-unit-test + env_vars: GOLANG,EINO + files: ${{ env.COVERAGE_FILE }} + token: ${{ secrets.CODECOV_TOKEN }} + codecov_yml_path: ./github/.codecov.yml + + benchmark-test: + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + repository-projects: write + steps: + - uses: actions/checkout@v4 + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: ${{ env.DEFAULT_GO_VERSION }} + + - name: Run Benchmark Tests + run: go test -bench=. -benchmem -run=none ./... + + compatibility-test: + strategy: + matrix: + go: [ "1.19", "1.20", "1.21", "1.22", "1.23", "1.24" ] + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + repository-projects: write + steps: + - uses: actions/checkout@v4 + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + cache: true + + - name: Compatibility Test + run: | + # just basic unit test, no coverage report + go test -race ./... + + api-compatibility: + name: api-compatibility-check + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + repository-projects: write + if: github.event_name == 'pull_request' + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: "1.22" + + - name: Install go-apidiff + run: go install github.com/joelanford/go-apidiff@v0.8.2 + + - name: Check API compatibility + id: apidiff + run: | + BASE_SHA=${{ github.event.pull_request.base.sha }} + HEAD_SHA=${{ github.event.pull_request.head.sha }} + + echo "Checking API compatibility between $BASE_SHA and $HEAD_SHA" + + go mod tidy + + if ! DIFF_OUTPUT=$(go-apidiff $BASE_SHA $HEAD_SHA 2>&1); then + echo "go-apidiff output: $DIFF_OUTPUT" + fi + + echo "diff_output<> $GITHUB_ENV + echo "$DIFF_OUTPUT" >> $GITHUB_ENV + echo "EOF" >> $GITHUB_ENV + + if echo "$DIFF_OUTPUT" | grep -q "Incompatible changes:"; then + echo "has_breaking_changes=true" >> $GITHUB_OUTPUT + else + echo "has_breaking_changes=false" >> $GITHUB_OUTPUT + fi + + - name: Create Review Thread + if: steps.apidiff.outputs.has_breaking_changes == 'true' + continue-on-error: true + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const reviewComments = await github.rest.pulls.listReviewComments({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: context.issue.number + }); + + const existingPackageComments = new Map(); + + for (const comment of reviewComments.data) { + if (comment.body.includes('Breaking API Changes Detected')) { + const packageMatch = comment.body.match(/Package: `([^`]+)`/); + if (packageMatch) { + const pkg = packageMatch[1]; + if (!existingPackageComments.has(pkg)) { + existingPackageComments.set(pkg, new Set()); + } + existingPackageComments.get(pkg).add(comment.path); + } + } + } + + const files = await github.rest.pulls.listFiles({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: context.issue.number + }); + + const diffOutput = process.env.diff_output || ''; + + const breakingChanges = new Map(); + + let currentPackage = ''; + let isInIncompatibleSection = false; + const lines = diffOutput.split('\n'); + + for (let i = 0; i < lines.length; i++) { + const line = lines[i].trim(); + + if (line.startsWith('github.com/')) { + currentPackage = line; + if (!breakingChanges.has(currentPackage)) { + breakingChanges.set(currentPackage, []); + } + continue; + } + + if (line === 'Incompatible changes:') { + isInIncompatibleSection = true; + continue; + } + + if (line === '') { + isInIncompatibleSection = false; + continue; + } + + if (isInIncompatibleSection && line.startsWith('- ')) { + const change = line.substring(2); + if (currentPackage) { + breakingChanges.get(currentPackage).push(change); + } + } + } + + const changedFiles = files.data; + + for (const [pkg, changes] of breakingChanges) { + if (changes.length === 0) continue; + + const pkgPath = pkg.split('/').slice(3).join('/'); + const matchingFile = changedFiles.find(file => + file.filename.includes(pkgPath) + ) || changedFiles[0]; + + const hasCommentForPackage = existingPackageComments.has(pkg) && + existingPackageComments.get(pkg).has(matchingFile.filename); + + if (matchingFile && !hasCommentForPackage) { + const changesList = changes.map(change => { + const [name, desc] = change.split(':').map(s => s.trim()); + return `- **${name}:** ${desc}`; + }).join('\n'); + + const commentBody = [ + '🚨 **Breaking API Changes Detected**', + '', + `Package: \`${pkg}\``, + '', + 'Incompatible changes:', + changesList, + '', + '
', + 'Review Guidelines', + '', + 'Please ensure that:', + '- The changes are absolutely necessary', + '- They are properly documented', + '- Migration guides are provided if needed', + '
', + '', + '⚠️ Please resolve this thread after reviewing the breaking changes.' + ].join('\n'); + + await github.rest.pulls.createReview({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: context.issue.number, + event: 'COMMENT', + comments: [{ + path: matchingFile.filename, + position: matchingFile.patch ? matchingFile.patch.split('\n').findIndex(line => line.startsWith('+')) + 1 : 1, + body: commentBody + }] + }); + + if (!existingPackageComments.has(pkg)) { + existingPackageComments.set(pkg, new Set()); + } + existingPackageComments.get(pkg).add(matchingFile.filename); + } + } \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..65ff0af --- /dev/null +++ b/.gitignore @@ -0,0 +1,66 @@ +# Binaries for programs and plugins +*.exe +*.exe~ +*.dll +*.so +*.dylib + +# Test binary, built with `go test -c` +*.test + +# Output of the go coverage tool, specifically when used with LiteIDE +*.out + +# Dependency directories (remove the comment below to include it) +# vendor/ + +# Go workspace file +go.work +go.work.sum + +# env file +.env + +# the result of the go build +output* +output/* + +# Files generated by IDEs +.idea/ +*.iml + +# Vim swap files +*.swp + +# Vscode files +.vscode + +/patches + +/vendor + +# Trae files +.trae + +# Specs files (internal documentation) +**/specs/ + +# Reports (generated analysis files) +reports/ +/todos + +.DS_Store +*.log* +.claude +CLAUDE.md +*.jsonl +*.txt +.agents/ + +# Specs directories +*/specs +/todos +/.claude/ + +# Internal dev setup (not for public repo) +/scripts/dev_setup_internal.sh diff --git a/.golangci.yaml b/.golangci.yaml new file mode 100644 index 0000000..d84b83a --- /dev/null +++ b/.golangci.yaml @@ -0,0 +1,148 @@ +# output configuration options +version: "2" +# All available settings of specific linters. +# Refer to https://golangci-lint.run/usage/linters +linters: + default: standard + enable: + - revive + - godoclint + - funlen + - cyclop + disable: + - errcheck + - staticcheck + - unused + - ineffassign + exclusions: + generated: lax + paths: + - ".*_test.go" + - ".*_mock.go" + rules: + - path: "^internal/.*" + linters: + - revive + - text: "var-naming: don't use underscores in Go names" + linters: + - revive + - path: "/utils/" + text: "var-naming: avoid meaningless package names" + linters: + - revive + - text: "exported: type name will be used as agent.AgentOption by other packages" + linters: + - revive + - path: "adk/prebuilt/deep/task_tool.go" + text: "argument-limit: maximum number of arguments per function exceeded" + linters: + - revive + - path: "compose/component_to_graph_node.go" + text: "argument-limit: maximum number of arguments per function exceeded" + linters: + - revive + - path: "compose/graph_run.go" + text: "argument-limit: maximum number of arguments per function exceeded" + linters: + - revive + - path: "adk/workflow.go" + text: "argument-limit: maximum number of arguments per function exceeded" + linters: + - revive + - path: "compose/graph.go" + linters: + - cyclop + text: "calculated cyclomatic complexity for function compile" + - path: "schema/message.go" + linters: + - cyclop + text: "calculated cyclomatic complexity for function ConcatMessages" + - path: "compose/graph_run.go" + linters: + - cyclop + text: "calculated cyclomatic complexity for function run" + - path: "compose/graph.go" + linters: + - funlen + text: "Function 'compile' is too long" + - path: "compose/graph_run.go" + linters: + - funlen + text: "Function 'run' is too long" + settings: + govet: + enable-all: true + # Disable analyzers by name. + # Run `go tool vet help` to see all analyzers. + disable: + - fieldalignment + revive: + # Sets the default failure confidence. + # This means that linting errors with less than 0.8 confidence will be ignored. + # Default: 0.8 + confidence: 0.8 + rules: + # Exported function and methods should have comments. + - name: exported + severity: error + exclude: + - "^internal/.*" + arguments: + - "disable-checks-on-constants" + - "disable-checks-on-variables" + - "disable-checks-on-types" + - "disable-checks-on-methods" + - name: package-comments + disabled: false + - name: var-naming + disabled: false + arguments: + # AllowList + - [ "utils", "s_", "err_", "err__", "plan_", "userInput_", "executedSteps_", "executedStep_", "iterator_", "in_", "out_" ] + # DenyList + - [ ] + - - extra-bad-package-names: + - helpers + - models + - name: argument-limit + arguments: [ 6 ] + - name: function-length + arguments: [ 120, 0 ] + godoclint: + check-exported: true + require-package-documentation: true + funlen: + lines: 200 + statements: 120 + cyclop: + max-complexity: 40 + package-average: 20 + +formatters: + enable: + - gci + - gofmt + settings: + gofmt: + # Simplify code: gofmt with `-s` option. + # Default: true + simplify: true + # Apply the rewrite rules to the source before reformatting. + # https://pkg.go.dev/cmd/gofmt + # Default: [] + rewrite-rules: + - pattern: 'interface{}' + replacement: 'any' + - pattern: 'a[b:len(a)]' + replacement: 'a[b:]' + gci: + # Section configuration to compare against. + # Section names are case-insensitive and may contain parameters in (). + # The default order of sections is `standard > default > custom > blank > dot > alias > localmodule`. + # If `custom-order` is `true`, it follows the order of `sections` option. + # Default: ["standard", "default"] + sections: + - standard + - default + - localmodule + custom-order: true diff --git a/.licenserc.yaml b/.licenserc.yaml new file mode 100644 index 0000000..9a13260 --- /dev/null +++ b/.licenserc.yaml @@ -0,0 +1,27 @@ +header: + license: + spdx-id: Apache-2.0 + copyright-owner: CloudWeGo Authors + + template: | + /* + * Copyright {{ .Year }} CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + + paths: + - '**/*.go' + - '**/*.s' + + comment: on-failure \ No newline at end of file diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..25ceb89 --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,128 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our +community a harassment-free experience for everyone, regardless of age, body +size, visible or invisible disability, ethnicity, sex characteristics, gender +identity and expression, level of experience, education, socio-economic status, +nationality, personal appearance, race, religion, or sexual identity +and orientation. + +We pledge to act and interact in ways that contribute to an open, welcoming, +diverse, inclusive, and healthy community. + +## Our Standards + +Examples of behavior that contributes to a positive environment for our +community include: + +* Demonstrating empathy and kindness toward other people +* Being respectful of differing opinions, viewpoints, and experiences +* Giving and gracefully accepting constructive feedback +* Accepting responsibility and apologizing to those affected by our mistakes, + and learning from the experience +* Focusing on what is best not just for us as individuals, but for the + overall community + +Examples of unacceptable behavior include: + +* The use of sexualized language or imagery, and sexual attention or + advances of any kind +* Trolling, insulting or derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or email + address, without their explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Enforcement Responsibilities + +Community leaders are responsible for clarifying and enforcing our standards of +acceptable behavior and will take appropriate and fair corrective action in +response to any behavior that they deem inappropriate, threatening, offensive, +or harmful. + +Community leaders have the right and responsibility to remove, edit, or reject +comments, commits, code, wiki edits, issues, and other contributions that are +not aligned to this Code of Conduct, and will communicate reasons for moderation +decisions when appropriate. + +## Scope + +This Code of Conduct applies within all community spaces, and also applies when +an individual is officially representing the community in public spaces. +Examples of representing our community include using an official e-mail address, +posting via an official social media account, or acting as an appointed +representative at an online or offline event. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported to the community leaders responsible for enforcement at +conduct@cloudwego.io. +All complaints will be reviewed and investigated promptly and fairly. + +All community leaders are obligated to respect the privacy and security of the +reporter of any incident. + +## Enforcement Guidelines + +Community leaders will follow these Community Impact Guidelines in determining +the consequences for any action they deem in violation of this Code of Conduct: + +### 1. Correction + +**Community Impact**: Use of inappropriate language or other behavior deemed +unprofessional or unwelcome in the community. + +**Consequence**: A private, written warning from community leaders, providing +clarity around the nature of the violation and an explanation of why the +behavior was inappropriate. A public apology may be requested. + +### 2. Warning + +**Community Impact**: A violation through a single incident or series +of actions. + +**Consequence**: A warning with consequences for continued behavior. No +interaction with the people involved, including unsolicited interaction with +those enforcing the Code of Conduct, for a specified period of time. This +includes avoiding interactions in community spaces as well as external channels +like social media. Violating these terms may lead to a temporary or +permanent ban. + +### 3. Temporary Ban + +**Community Impact**: A serious violation of community standards, including +sustained inappropriate behavior. + +**Consequence**: A temporary ban from any sort of interaction or public +communication with the community for a specified period of time. No public or +private interaction with the people involved, including unsolicited interaction +with those enforcing the Code of Conduct, is allowed during this period. +Violating these terms may lead to a permanent ban. + +### 4. Permanent Ban + +**Community Impact**: Demonstrating a pattern of violation of community +standards, including sustained inappropriate behavior, harassment of an +individual, or aggression toward or disparagement of classes of individuals. + +**Consequence**: A permanent ban from any sort of public interaction within +the community. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], +version 2.0, available at +https://www.contributor-covenant.org/version/2/0/code_of_conduct.html. + +Community Impact Guidelines were inspired by [Mozilla's code of conduct +enforcement ladder](https://github.com/mozilla/diversity). + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see the FAQ at +https://www.contributor-covenant.org/faq. Translations are available at +https://www.contributor-covenant.org/translations. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..db3d580 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,51 @@ +# How to Contribute + +## Your First Pull Request +We use GitHub for our codebase. You can start by reading [How To Pull Request](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests/about-pull-requests). + +## Branch Organization +We use [git-flow](https://nvie.com/posts/a-successful-git-branching-model/) as our branch organization, as known as [FDD](https://en.wikipedia.org/wiki/Feature-driven_development) + +## Bugs +### 1. How to Find Known Issues +We are using [Github Issues](https://github.com/cloudwego/{project_name}/issues) for our public bugs. We keep a close eye on this and try to make it clear when we have an internal fix in progress. Before filing a new task, try to make sure your problem doesn’t already exist. + +### 2. Reporting New Issues +Providing a reduced test code is a recommended way for reporting issues. Then can place in: +- Just in issues +- [Golang Playground](https://play.golang.org/) + +### 3. Security Bugs +Please do not report the safe disclosure of bugs to public issues. Contact us by [Support Email](mailto:conduct@cloudwego.io) + +## How to Get in Touch +- [Email](mailto:conduct@cloudwego.io) + +## Submit a Pull Request +Before you submit your Pull Request (PR) consider the following guidelines: +1. Search [GitHub](https://github.com/cloudwego/{project_name}/pulls) for an open or closed PR that relates to your submission. You don't want to duplicate existing efforts. +2. Be sure that an issue describes the problem you're fixing, or documents the design for the feature you'd like to add. Discussing the design upfront helps to ensure that we're ready to accept your work. +3. [Fork](https://docs.github.com/en/github/getting-started-with-github/fork-a-repo) the cloudwego {project_name} repo. +4. In your forked repository, make your changes in a new git branch: + ``` + git checkout -b my-fix-branch develop + ``` +5. Create your patch, including appropriate test cases. +6. Follow our [Style Guides](#code-style-guides). +7. Commit your changes using a descriptive commit message that follows [AngularJS Git Commit Message Conventions](https://docs.google.com/document/d/1QrDFcIiPjSLDn3EL15IJygNPiHORgU1_OOAqWjiDU5Y/edit). + Adherence to these conventions is necessary because release notes are automatically generated from these messages. +8. Push your branch to GitHub: + ``` + git push origin my-fix-branch + ``` +9. In GitHub, send a pull request to `{project_name}:develop` + +## Contribution Prerequisites +- Our development environment keeps up with [Go Official](https://golang.org/project/). +- You need fully checking with lint tools before submit your pull request. [gofmt](https://golang.org/pkg/cmd/gofmt/) and [golangci-lint](https://github.com/golangci/golangci-lint) +- You are familiar with [GitHub](https://github.com) +- Maybe you need familiar with [Actions](https://github.com/features/actions)(our default workflow tool). + +## Code Style Guides +- [Effective Go](https://golang.org/doc/effective_go) +- [Go Code Review Comments](https://github.com/golang/go/wiki/CodeReviewComments) diff --git a/LICENSE-APACHE b/LICENSE-APACHE new file mode 100644 index 0000000..66a27ec --- /dev/null +++ b/LICENSE-APACHE @@ -0,0 +1,177 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + diff --git a/README.md b/README.md new file mode 100644 index 0000000..a74ba79 --- /dev/null +++ b/README.md @@ -0,0 +1,209 @@ +# Eino + +![coverage](https://raw.githubusercontent.com/cloudwego/eino/badges/.badges/main/coverage.svg) +[![Release](https://img.shields.io/github/v/release/cloudwego/eino)](https://github.com/cloudwego/eino/releases) +[![WebSite](https://img.shields.io/website?up_message=cloudwego&url=https%3A%2F%2Fwww.cloudwego.io%2F)](https://www.cloudwego.io/) +[![License](https://img.shields.io/github/license/cloudwego/eino)](https://github.com/cloudwego/eino/blob/main/LICENSE) +[![Go Report Card](https://goreportcard.com/badge/github.com/cloudwego/eino)](https://goreportcard.com/report/github.com/cloudwego/eino) +[![OpenIssue](https://img.shields.io/github/issues/cloudwego/eino)](https://github.com/cloudwego/kitex/eino) +[![ClosedIssue](https://img.shields.io/github/issues-closed/cloudwego/eino)](https://github.com/cloudwego/eino/issues?q=is%3Aissue+is%3Aclosed) +![Stars](https://img.shields.io/github/stars/cloudwego/eino) +![Forks](https://img.shields.io/github/forks/cloudwego/eino) + +English | [中文](README.zh_CN.md) + +# Overview + +**Eino['aino]** is an LLM application development framework in Golang. It draws from LangChain, Google ADK, and other open-source frameworks, and is designed to follow Golang conventions. + +Eino provides: +- **[Components](https://github.com/cloudwego/eino-ext)**: reusable building blocks like `ChatModel`, `Tool`, `Retriever`, and `ChatTemplate`, with official implementations for OpenAI, Ollama, and more. +- **Agent Development Kit (ADK)**: build AI agents with tool use, multi-agent coordination, context management, interrupt/resume for human-in-the-loop, and ready-to-use agent patterns. +- **Composition**: connect components into graphs and workflows that can run standalone or be exposed as tools for agents. +- **[Examples](https://github.com/cloudwego/eino-examples)**: working code for common patterns and real-world use cases. + +![](.github/static/img/eino/eino_concept.jpeg) + +# Quick Start + +## ChatModelAgent + +Configure a ChatModel, optionally add tools, and you have a working agent: + +```Go +chatModel, _ := openai.NewChatModel(ctx, &openai.ChatModelConfig{ + Model: "gpt-4o", + APIKey: os.Getenv("OPENAI_API_KEY"), +}) + +agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Model: chatModel, +}) + +runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent}) +iter := runner.Query(ctx, "Hello, who are you?") +for { + event, ok := iter.Next() + if !ok { + break + } + fmt.Println(event.Message.Content) +} +``` + +Add tools to give the agent capabilities: + +```Go +agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Model: chatModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{weatherTool, calculatorTool}, + }, + }, +}) +``` + +The agent handles the ReAct loop internally — it decides when to call tools and when to respond. + +→ [ChatModelAgent examples](https://github.com/cloudwego/eino-examples/tree/main/adk/intro) · [docs](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/chat_model/) + +## DeepAgent + +For complex tasks, use DeepAgent. It breaks down problems into steps, delegates to sub-agents, and tracks progress: + +```Go +deepAgent, _ := deep.New(ctx, &deep.Config{ + ChatModel: chatModel, + SubAgents: []adk.Agent{researchAgent, codeAgent}, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{shellTool, pythonTool, webSearchTool}, + }, + }, +}) + +runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent}) +iter := runner.Query(ctx, "Analyze the sales data in report.csv and generate a summary chart") +``` + +DeepAgent can be configured to coordinate multiple specialized agents, run shell commands, execute Python code, and search the web. + +→ [DeepAgent example](https://github.com/cloudwego/eino-examples/tree/main/adk/multiagent/deep) · [docs](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_implementation/deepagents/) + +## Composition + +When you need precise control over execution flow, use `compose` to build graphs and workflows: + +```Go +graph := compose.NewGraph[*Input, *Output]() +graph.AddLambdaNode("validate", validateFn) +graph.AddChatModelNode("generate", chatModel) +graph.AddLambdaNode("format", formatFn) + +graph.AddEdge(compose.START, "validate") +graph.AddEdge("validate", "generate") +graph.AddEdge("generate", "format") +graph.AddEdge("format", compose.END) + +runnable, _ := graph.Compile(ctx) +result, _ := runnable.Invoke(ctx, input) +``` + +Compositions can be exposed as tools for agents, bridging deterministic workflows with autonomous behavior: + +```Go +tool, _ := graphtool.NewInvokableGraphTool(graph, "data_pipeline", "Process and validate data") + +agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Model: chatModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool}, + }, + }, +}) +``` + +This lets you build domain-specific pipelines with exact control, then let agents decide when to use them. + +→ [GraphTool examples](https://github.com/cloudwego/eino-examples/tree/main/adk/common/tool/graphtool) · [compose docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/) + +# Key Features + +## Component Ecosystem + +Eino defines component abstractions (ChatModel, Tool, Retriever, Embedding, etc.) with official implementations for OpenAI, Claude, Gemini, Ark, Ollama, Elasticsearch, and more. + +→ [eino-ext](https://github.com/cloudwego/eino-ext) + +## Stream Processing + +Eino automatically handles streaming throughout orchestration: concatenating, boxing, merging, and copying streams as data flows between nodes. Components only implement the streaming paradigms that make sense for them; the framework handles the rest. + +→ [docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/) + +## Callback Aspects + +Inject logging, tracing, and metrics at fixed points (OnStart, OnEnd, OnError, OnStartWithStreamInput, OnEndWithStreamOutput) across components, graphs, and agents. + +→ [docs](https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/) + +## Interrupt/Resume + +Any agent or tool can pause execution for human input and resume from checkpoint. The framework handles state persistence and routing. + +→ [docs](https://www.cloudwego.io/docs/eino/core_modules/eino_adk/agent_hitl/) · [examples](https://github.com/cloudwego/eino-examples/tree/main/adk/human-in-the-loop) + +# Framework Structure + +![](.github/static/img/eino/eino_framework.jpeg) + +The Eino framework consists of: + +- Eino (this repo): Type definitions, streaming mechanism, component abstractions, orchestration, agent implementations, aspect mechanisms + +- [EinoExt](https://github.com/cloudwego/eino-ext): Component implementations, callback handlers, usage examples, evaluators, prompt optimizers + +- [Eino Devops](https://github.com/cloudwego/eino-ext/tree/main/devops): Visualized development and debugging + +- [EinoExamples](https://github.com/cloudwego/eino-examples): Example applications and best practices + +## Documentation + +- [Eino User Manual](https://www.cloudwego.io/zh/docs/eino/) +- [Eino: Quick Start](https://www.cloudwego.io/zh/docs/eino/quick_start/) + +## Dependencies +- Go 1.18 and above. + +## Code Style + +This repo uses `golangci-lint`. Check locally with: + +```bash +golangci-lint run ./... +``` + +Rules enforced: +- Exported functions, interfaces, packages, etc. should have GoDoc comments +- Code should be formatted with `gofmt -s` +- Import order should follow `goimports` (std -> third party -> local) + +## Security + +If you discover a potential security issue in this project, or think you may +have discovered a security issue, we ask that you notify Bytedance Security via our [security center](https://security.bytedance.com/src) or [vulnerability reporting email](mailto:sec@bytedance.com?subject=Feedback%20On%20Eino). + +Do **not** create a public GitHub issue. + +## Contact +- Membership: [COMMUNITY MEMBERSHIP](https://github.com/cloudwego/community/blob/main/COMMUNITY_MEMBERSHIP.md) +- Issues: [Issues](https://github.com/cloudwego/eino/issues) +- Lark: Scan the QR code below with [Feishu](https://www.feishu.cn/en/) to join the CloudWeGo/eino user group. + +    LarkGroup + +## License + +This project is licensed under the [Apache-2.0 License](LICENSE-APACHE). diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..c9b8ef7 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`cloudwego/eino` +- 原始仓库：https://github.com/cloudwego/eino +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/README.zh_CN.md b/README.zh_CN.md new file mode 100644 index 0000000..d21e7a1 --- /dev/null +++ b/README.zh_CN.md @@ -0,0 +1,210 @@ +# Eino + +![coverage](https://raw.githubusercontent.com/cloudwego/eino/badges/.badges/main/coverage.svg) +[![Release](https://img.shields.io/github/v/release/cloudwego/eino)](https://github.com/cloudwego/eino/releases) +[![WebSite](https://img.shields.io/website?up_message=cloudwego&url=https%3A%2F%2Fwww.cloudwego.io%2F)](https://www.cloudwego.io/) +[![License](https://img.shields.io/github/license/cloudwego/eino)](https://github.com/cloudwego/eino/blob/main/LICENSE) +[![Go Report Card](https://goreportcard.com/badge/github.com/cloudwego/eino)](https://goreportcard.com/report/github.com/cloudwego/eino) +[![OpenIssue](https://img.shields.io/github/issues/cloudwego/eino)](https://github.com/cloudwego/kitex/eino) +[![ClosedIssue](https://img.shields.io/github/issues-closed/cloudwego/eino)](https://github.com/cloudwego/eino/issues?q=is%3Aissue+is%3Aclosed) +![Stars](https://img.shields.io/github/stars/cloudwego/eino) +![Forks](https://img.shields.io/github/forks/cloudwego/eino) + +[English](README.md) | 中文 + +# 简介 + +**Eino['aino]** 是一个 Go 语言的 LLM 应用开发框架，借鉴了 LangChain、Google ADK 等开源项目，按照 Go 的惯例设计。 + +Eino 提供： +- **[组件](https://github.com/cloudwego/eino-ext)**：`ChatModel`、`Tool`、`Retriever`、`ChatTemplate` 等可复用模块，官方实现覆盖 OpenAI、Ollama 等 +- **智能体开发套件（ADK）**：支持工具调用、多智能体协同、上下文管理、中断/恢复等人机交互，以及开箱即用的智能体模式 +- **编排**：把组件组装成图或工作流，既能独立运行，也能作为工具给智能体调用 +- **[示例](https://github.com/cloudwego/eino-examples)**：常见模式和实际场景的可运行代码 + +![](.github/static/img/eino/eino_concept.jpeg) + +# 快速上手 + +## ChatModelAgent + +配置好 ChatModel，加上工具（可选），就能跑起来： + +```Go +chatModel, _ := openai.NewChatModel(ctx, &openai.ChatModelConfig{ + Model: "gpt-4o", + APIKey: os.Getenv("OPENAI_API_KEY"), +}) + +agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Model: chatModel, +}) + +runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent}) +iter := runner.Query(ctx, "Hello, who are you?") +for { + event, ok := iter.Next() + if !ok { + break + } + fmt.Println(event.Message.Content) +} +``` + +加工具让智能体有更多能力： + +```Go +agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Model: chatModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{weatherTool, calculatorTool}, + }, + }, +}) +``` + +智能体内部自动处理 ReAct 循环，自己判断什么时候调工具、什么时候回复。 + +→ [ChatModelAgent 示例](https://github.com/cloudwego/eino-examples/tree/main/adk/intro) · [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/agent_implementation/chat_model/) + +## DeepAgent + +复杂任务用 DeepAgent，它会把问题拆成步骤，分派给子智能体，并追踪进度： + +```Go +deepAgent, _ := deep.New(ctx, &deep.Config{ + ChatModel: chatModel, + SubAgents: []adk.Agent{researchAgent, codeAgent}, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{shellTool, pythonTool, webSearchTool}, + }, + }, +}) + +runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent}) +iter := runner.Query(ctx, "Analyze the sales data in report.csv and generate a summary chart") +``` + +DeepAgent 可以配置成：协调多个专业智能体、跑 shell 命令、执行 Python、搜索网络。 + +→ [DeepAgent 示例](https://github.com/cloudwego/eino-examples/tree/main/adk/multiagent/deep) · [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/agent_implementation/deepagents/) + +## 编排 + +需要精确控制执行流程时，用 `compose` 搭图或工作流： + +```Go +graph := compose.NewGraph[*Input, *Output]() +graph.AddLambdaNode("validate", validateFn) +graph.AddChatModelNode("generate", chatModel) +graph.AddLambdaNode("format", formatFn) + +graph.AddEdge(compose.START, "validate") +graph.AddEdge("validate", "generate") +graph.AddEdge("generate", "format") +graph.AddEdge("format", compose.END) + +runnable, _ := graph.Compile(ctx) +result, _ := runnable.Invoke(ctx, input) +``` + +编排出来的流程可以包装成工具给智能体用，把确定性流程和自主决策结合起来： + +```Go +tool, _ := graphtool.NewInvokableGraphTool(graph, "data_pipeline", "Process and validate data") + +agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Model: chatModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool}, + }, + }, +}) +``` + +这样你可以写出精确可控的业务流程，再让智能体决定什么时候调用。 + +→ [GraphTool 示例](https://github.com/cloudwego/eino-examples/tree/main/adk/common/tool/graphtool) · [编排文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/) + +# 主要特性 + +## 组件生态 + +Eino 定义了组件抽象（ChatModel、Tool、Retriever、Embedding 等），官方实现覆盖 OpenAI、Claude、Gemini、Ark、Ollama、Elasticsearch 等。 + +→ [eino-ext](https://github.com/cloudwego/eino-ext) + +## 流式处理 + +Eino 在编排中自动处理流式：拼接、装箱、合并、复制。组件只需实现有业务意义的流式范式，框架处理剩下的。 + +→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/) + +## 回调切面 + +在固定切点（OnStart、OnEnd、OnError、OnStartWithStreamInput、OnEndWithStreamOutput）注入日志、追踪、指标，适用于组件、图、智能体。 + +→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/chain_and_graph_orchestration/callback_manual/) + +## 中断/恢复 + +任何智能体或工具都能暂停等待人工输入，从检查点恢复。框架处理状态持久化和路由。 + +→ [文档](https://www.cloudwego.io/zh/docs/eino/core_modules/eino_adk/agent_hitl/) · [示例](https://github.com/cloudwego/eino-examples/tree/main/adk/human-in-the-loop) + +# 框架结构 + +![](.github/static/img/eino/eino_framework.jpeg) + +Eino 框架包含： + +- Eino（本仓库）：类型定义、流处理机制、组件抽象、编排、智能体实现、切面机制 + +- [EinoExt](https://github.com/cloudwego/eino-ext)：组件实现、回调处理器、使用示例、评估器、提示优化器 + +- [Eino Devops](https://github.com/cloudwego/eino-ext/tree/main/devops)：可视化开发和调试 + +- [EinoExamples](https://github.com/cloudwego/eino-examples)：示例应用和最佳实践 + +## 文档 + +- [Eino 用户手册](https://www.cloudwego.io/zh/docs/eino/) +- [Eino: 快速开始](https://www.cloudwego.io/zh/docs/eino/quick_start/) + +## 依赖 +- Go 1.18 及以上 + +## 代码规范 + +本仓库使用 `golangci-lint`，本地检查： + +```bash +golangci-lint run ./... +``` + +规则： +- 导出的函数、接口、package 等需要 GoDoc 注释 +- 代码格式符合 `gofmt -s` +- import 顺序符合 `goimports`（std -> third party -> local） + +## 安全 + +如果你在该项目中发现潜在的安全问题，或你认为可能发现了安全问题，请通过我们的[安全中心](https://security.bytedance.com/src) +或[漏洞报告邮箱](mailto:sec@bytedance.com?subject=关于Eino的反馈)通知字节跳动安全团队。 + + +请**不要**创建公开的 GitHub Issue。 + +## 联系我们 +- 成为 member：[COMMUNITY MEMBERSHIP](https://github.com/cloudwego/community/blob/main/COMMUNITY_MEMBERSHIP.md) +- Issues：[Issues](https://github.com/cloudwego/eino/issues) +- 飞书：扫码加入 CloudWeGo/eino 用户群 + +    LarkGroup + +## 开源许可证 + +本项目基于 [Apache-2.0 许可证](LICENSE-APACHE) 开源。 diff --git a/_typos.toml b/_typos.toml new file mode 100644 index 0000000..c3b335d --- /dev/null +++ b/_typos.toml @@ -0,0 +1,14 @@ +# Typo check: https://github.com/crate-ci/typos +[default] + +[default.extend-words] +Invokable = "Invokable" +invokable = "invokable" +InvokableLambda = "InvokableLambda" +InvokableRun = "InvokableRun" +typ = "typ" +byted = "byted" +cpy = "cpy" +mak = "mak" +[files] +extend-exclude = ["go.mod", "go.sum", "check_branch_name.sh"] diff --git a/adk/agent_tool.go b/adk/agent_tool.go new file mode 100644 index 0000000..3f120c2 --- /dev/null +++ b/adk/agent_tool.go @@ -0,0 +1,417 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package adk provides core agent development kit utilities and types. +package adk + +import ( + "context" + "errors" + "fmt" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +var ( + defaultAgentToolParam = schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "request": { + Desc: "request to be processed", + Required: true, + Type: schema.String, + }, + }) +) + +type AgentToolOptions struct { + fullChatHistoryAsInput bool + agentInputSchema *schema.ParamsOneOf +} + +type AgentToolOption func(*AgentToolOptions) + +// WithFullChatHistoryAsInput enables using the full chat history as input. +func WithFullChatHistoryAsInput() AgentToolOption { + return func(options *AgentToolOptions) { + options.fullChatHistoryAsInput = true + } +} + +// WithAgentInputSchema sets a custom input schema for the agent tool. +func WithAgentInputSchema(schema *schema.ParamsOneOf) AgentToolOption { + return func(options *AgentToolOptions) { + options.agentInputSchema = schema + } +} + +func withAgentToolEnableStreaming(enabled bool) tool.Option { + return tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) { + opt.enableStreaming = enabled + }) +} + +// NewAgentTool creates a tool that wraps an agent for invocation. +// +// The agent must have a non-empty Name and Description, as they are used as +// the tool's name and description respectively. This is validated when Info() +// is called during tool setup. +// +// Event Streaming: +// When EmitInternalEvents is enabled in ToolsConfig, the agent tool will emit AgentEvent +// from the inner agent to the parent agent's AsyncGenerator, allowing real-time streaming +// of the inner agent's output to the end-user via Runner. +// +// Note that these forwarded events are NOT recorded in the parent agent's runSession. +// They are only emitted to the end-user and have no effect on the parent agent's state +// or checkpoint. The only exception is Interrupted action, which is propagated via +// CompositeInterrupt to enable proper interrupt/resume across agent boundaries. +// +// Action Scoping: +// Actions emitted by the inner agent are scoped to the agent tool boundary: +// - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume across boundaries +// - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool; these actions only affect +// the inner agent's execution and do not propagate to the parent agent +// +// This scoping ensures that nested agents cannot unexpectedly terminate or transfer control +// of their parent agent's execution flow. +func NewAgentTool(_ context.Context, agent Agent, options ...AgentToolOption) tool.BaseTool { + opts := &AgentToolOptions{} + for _, opt := range options { + opt(opts) + } + + return &agentTool{ + agent: agent, + fullChatHistoryAsInput: opts.fullChatHistoryAsInput, + inputSchema: opts.agentInputSchema, + } +} + +// NewTypedAgentTool creates a new agent tool that wraps a TypedAgent as a tool.BaseTool. +func NewTypedAgentTool[M MessageType](_ context.Context, agent TypedAgent[M], options ...AgentToolOption) tool.BaseTool { + opts := &AgentToolOptions{} + for _, opt := range options { + opt(opts) + } + + return &typedAgentTool[M]{ + agent: agent, + fullChatHistoryAsInput: opts.fullChatHistoryAsInput, + inputSchema: opts.agentInputSchema, + } +} + +type typedAgentTool[M MessageType] struct { + agent TypedAgent[M] + + fullChatHistoryAsInput bool + inputSchema *schema.ParamsOneOf +} + +type agentTool = typedAgentTool[*schema.Message] + +type agentToolRequest struct { + Request string `json:"request"` +} + +func (at *typedAgentTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) { + name := at.agent.Name(ctx) + if name == "" { + return nil, errors.New("agent tool requires a non-empty Name") + } + desc := at.agent.Description(ctx) + if desc == "" { + return nil, errors.New("agent tool requires a non-empty Description") + } + param := at.inputSchema + if param == nil { + param = defaultAgentToolParam + } + + return &schema.ToolInfo{ + Name: name, + Desc: desc, + ParamsOneOf: param, + }, nil +} + +func (at *typedAgentTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + if cancelCtx := getCancelContext(ctx); cancelCtx != nil { + cancelCtx.markAgentToolDescendant() + } + + gen, enableStreaming := getEmitGeneratorAndEnableStreaming[M](opts) + var ms *bridgeStore + var iter *AsyncIterator[*TypedAgentEvent[M]] + var err error + + wasInterrupted, hasState, state := tool.GetInterruptState[[]byte](ctx) + if !wasInterrupted { + ms = newBridgeStore() + + var input []M + if at.fullChatHistoryAsInput { + var zero M + if _, ok := any(zero).(*schema.Message); !ok { + // fullChatHistoryAsInput is only supported for *schema.Message agents and will not + // be extended to *schema.AgenticMessage. The chat history format and role semantics + // differ fundamentally between Message and AgenticMessage, and the history rewriting + // logic (role attribution, system message filtering, transfer messages) is specific + // to the Message model. + return "", fmt.Errorf("fullChatHistoryAsInput is only supported for *schema.Message agents") + } + msgInput, histErr := getReactChatHistory(ctx, at.agent.Name(ctx)) + if histErr != nil { + return "", histErr + } + input = any(msgInput).([]M) + } else { + if at.inputSchema == nil { + req := &agentToolRequest{} + err = sonic.UnmarshalString(argumentsInJSON, req) + if err != nil { + return "", err + } + argumentsInJSON = req.Request + } + input = newTypedUserMessages[M](argumentsInJSON) + } + + runner := newTypedInvokableAgentToolRunner(at.agent, ms, enableStreaming) + iter = runner.Run(ctx, input, + append(extractAndDeriveAgentToolCancelCtx(ctx, at.agent.Name(ctx), opts), WithCheckPointID(bridgeCheckpointID), withSharedParentSession())...) + } else { + if !hasState { + return "", fmt.Errorf("agent tool '%s' interrupt has happened, but cannot find interrupt state", at.agent.Name(ctx)) + } + + ms = newResumeBridgeStore(bridgeCheckpointID, state) + + agentOpts := extractAndDeriveAgentToolCancelCtx(ctx, at.agent.Name(ctx), opts) + agentOpts = append(agentOpts, withSharedParentSession()) + + runner := newTypedInvokableAgentToolRunner(at.agent, ms, enableStreaming) + iter, err = runner.Resume(ctx, bridgeCheckpointID, agentOpts...) + if err != nil { + return "", err + } + } + + var lastEvent *TypedAgentEvent[M] + for { + event, ok := iter.Next() + if !ok { + break + } + + if lastEvent != nil && + lastEvent.Output != nil && + lastEvent.Output.MessageOutput != nil && + lastEvent.Output.MessageOutput.MessageStream != nil { + lastEvent.Output.MessageOutput.MessageStream.Close() + } + + if event.Err != nil { + return "", event.Err + } + + if gen != nil { + if event.Action == nil || event.Action.Interrupted == nil { + if parentRunCtx := getRunCtx(ctx); parentRunCtx != nil && len(parentRunCtx.RunPath) > 0 { + rp := make([]RunStep, 0, len(parentRunCtx.RunPath)+len(event.RunPath)) + rp = append(rp, parentRunCtx.RunPath...) + rp = append(rp, event.RunPath...) + event.RunPath = rp + } + tmp := copyTypedAgentEvent(event) + gen.Send(event) + event = tmp + } + } + + lastEvent = event + } + + if lastEvent != nil && lastEvent.Action != nil && lastEvent.Action.Interrupted != nil { + data, existed, err_ := ms.Get(ctx, bridgeCheckpointID) + if err_ != nil { + return "", fmt.Errorf("failed to get interrupt info: %w", err_) + } + if !existed { + return "", fmt.Errorf("interrupt has happened, but cannot find interrupt info") + } + + return "", tool.CompositeInterrupt(ctx, "agent tool interrupt", data, + lastEvent.Action.internalInterrupted) + } + + if lastEvent == nil { + return "", errors.New("no event returned") + } + + var ret string + if lastEvent.Output != nil { + if output := lastEvent.Output.MessageOutput; output != nil { + msg, err := output.GetMessage() + if err != nil { + return "", err + } + ret = extractTextContent(msg) + } + } + + return ret, nil +} + +// agentToolOptions is a wrapper structure used to convert AgentRunOption slices to tool.Option. +// It stores the agent name and corresponding run options for tool-specific processing. +type agentToolOptions struct { + agentName string + opts []AgentRunOption + enableStreaming bool +} + +// typedAgentToolEventOptions carries the parent runner's event generator for a +// specific message type. This keeps forwarded internal events type-compatible +// with the parent event stream. +type typedAgentToolEventOptions[M MessageType] struct { + generator *AsyncGenerator[*TypedAgentEvent[M]] +} + +func withAgentToolOptions(agentName string, opts []AgentRunOption) tool.Option { + return tool.WrapImplSpecificOptFn(func(opt *agentToolOptions) { + opt.agentName = agentName + opt.opts = opts + }) +} + +func withAgentToolEventGenerator(gen *AsyncGenerator[*AgentEvent]) tool.Option { + return withTypedAgentToolEventGenerator(gen) +} + +func withTypedAgentToolEventGenerator[M MessageType](gen *AsyncGenerator[*TypedAgentEvent[M]]) tool.Option { + return tool.WrapImplSpecificOptFn(func(o *typedAgentToolEventOptions[M]) { + o.generator = gen + }) +} + +func getOptionsByAgentName(agentName string, opts []tool.Option) []AgentRunOption { + var ret []AgentRunOption + for _, opt := range opts { + o := tool.GetImplSpecificOptions[agentToolOptions](nil, opt) + if o != nil && o.agentName == agentName { + ret = append(ret, o.opts...) + } + } + return ret +} + +func extractAndDeriveAgentToolCancelCtx(ctx context.Context, agentName string, opts []tool.Option) []AgentRunOption { + agentOpts := getOptionsByAgentName(agentName, opts) + baseOpts := getCommonOptions(nil, agentOpts...) + parentCtx := baseOpts.cancelCtx + if parentCtx == nil { + parentCtx = getCancelContext(ctx) + } + if parentCtx != nil { + parentCtx.markAgentToolDescendant() + childCtx := parentCtx.deriveAgentToolCancelContext(ctx) + agentOpts = append(agentOpts, WrapImplSpecificOptFn(func(o *options) { + o.cancelCtx = childCtx + })) + } + return agentOpts +} + +func getEmitGeneratorAndEnableStreaming[M MessageType](opts []tool.Option) (*AsyncGenerator[*TypedAgentEvent[M]], bool) { + o := tool.GetImplSpecificOptions[agentToolOptions](nil, opts...) + eventOptions := tool.GetImplSpecificOptions[typedAgentToolEventOptions[M]](nil, opts...) + if o == nil && eventOptions == nil { + return nil, false + } + + var gen *AsyncGenerator[*TypedAgentEvent[M]] + if eventOptions != nil { + gen = eventOptions.generator + } + + var enableStreaming bool + if o != nil { + enableStreaming = o.enableStreaming + } + + return gen, enableStreaming +} + +func getReactChatHistory(ctx context.Context, destAgentName string) ([]Message, error) { + var messages []Message + err := compose.ProcessState(ctx, func(ctx context.Context, st *State) error { + if len(st.Messages) == 0 { + return nil + } + messages = make([]Message, len(st.Messages)-1) + copy(messages, st.Messages[:len(st.Messages)-1]) + return nil + }) + if err != nil { + return nil, fmt.Errorf("failed to get chat history from state: %w", err) + } + + var agentName string + if runCtx := getRunCtx(ctx); runCtx != nil && len(runCtx.RunPath) > 0 { + agentName = runCtx.RunPath[len(runCtx.RunPath)-1].agentName + } + + a, t := GenTransferMessages(ctx, destAgentName) + messages = append(messages, a, t) + history := make([]Message, 0, len(messages)) + for _, msg := range messages { + if msg.Role == schema.System { + continue + } + + if msg.Role == schema.Assistant || msg.Role == schema.Tool { + msg = rewriteMessage(msg, agentName) + } + + history = append(history, msg) + } + + return history, nil +} + +func newTypedUserMessages[M MessageType](text string) []M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any([]Message{schema.UserMessage(text)}).([]M) + case *schema.AgenticMessage: + return any([]*schema.AgenticMessage{schema.UserAgenticMessage(text)}).([]M) + default: + return nil + } +} + +func newTypedInvokableAgentToolRunner[M MessageType](agent TypedAgent[M], store compose.CheckPointStore, enableStreaming bool) *TypedRunner[M] { + return &TypedRunner[M]{ + a: agent, + enableStreaming: enableStreaming, + store: store, + } +} diff --git a/adk/agent_tool_test.go b/adk/agent_tool_test.go new file mode 100644 index 0000000..785ad99 --- /dev/null +++ b/adk/agent_tool_test.go @@ -0,0 +1,1241 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "fmt" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type mockChatModelForAttack struct { + generateFn func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) +} + +func (m *mockChatModelForAttack) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return m.generateFn(ctx, input, opts...) +} + +func (m *mockChatModelForAttack) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + result, err := m.generateFn(ctx, input, opts...) + if err != nil { + return nil, err + } + r, w := schema.Pipe[*schema.Message](1) + go func() { defer w.Close(); w.Send(result, nil) }() + return r, nil +} + +// mockAgent implements the Agent interface for testing +type mockAgentForTool struct { + name string + description string + responses []*AgentEvent +} + +func (a *mockAgentForTool) Name(_ context.Context) string { + return a.name +} + +func (a *mockAgentForTool) Description(_ context.Context) string { + return a.description +} + +func (a *mockAgentForTool) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + defer generator.Close() + + for _, event := range a.responses { + generator.Send(event) + + // If the event has an Exit action, stop sending events + if event.Action != nil && event.Action.Exit { + break + } + } + }() + + return iterator +} + +func newMockAgentForTool(name, description string, responses []*AgentEvent) *mockAgentForTool { + return &mockAgentForTool{ + name: name, + description: description, + responses: responses, + } +} + +func TestAgentTool_Info(t *testing.T) { + // Create a mock agent + mockAgent_ := newMockAgentForTool("TestAgent", "Test agent description", nil) + + // Create an agentTool with the mock agent + agentTool_ := NewAgentTool(context.Background(), mockAgent_) + + // Test the Info method + ctx := context.Background() + info, err := agentTool_.Info(ctx) + + // Verify results + assert.NoError(t, err) + assert.NotNil(t, info) + assert.Equal(t, "TestAgent", info.Name) + assert.Equal(t, "Test agent description", info.Desc) + assert.NotNil(t, info.ParamsOneOf) +} + +func TestAgentTool_Info_EmptyName(t *testing.T) { + agentTool_ := NewAgentTool(context.Background(), newMockAgentForTool("", "desc", nil)) + _, err := agentTool_.Info(context.Background()) + assert.ErrorContains(t, err, "non-empty Name") +} + +func TestAgentTool_Info_EmptyDescription(t *testing.T) { + agentTool_ := NewAgentTool(context.Background(), newMockAgentForTool("name", "", nil)) + _, err := agentTool_.Info(context.Background()) + assert.ErrorContains(t, err, "non-empty Description") +} + +func TestAgentTool_SharedParentSessionValues(t *testing.T) { + ctx := context.Background() + + inner := &sessionValuesAgent{name: "inner"} + innerTool := NewAgentTool(ctx, inner).(tool.InvokableTool) + + input := &AgentInput{Messages: []Message{schema.UserMessage("q")}} + ctx, _ = initRunCtx(ctx, "outer", input) + AddSessionValue(ctx, "parent_key", "parent_val") + parentSession := getRunCtx(ctx).Session + + _, err := innerTool.InvokableRun(ctx, `{"request":"hello"}`) + assert.NoError(t, err) + + assert.Equal(t, "parent_val", inner.seenParentValue) + assert.NotNil(t, inner.capturedSession) + assert.NotSame(t, parentSession, inner.capturedSession) + assert.NotNil(t, parentSession.valuesMtx) + assert.Same(t, parentSession.valuesMtx, inner.capturedSession.valuesMtx) + + mtx := parentSession.valuesMtx + mtx.Lock() + inner.capturedSession.Values["direct_child_key"] = "direct_child_val" + mtx.Unlock() + + mtx.Lock() + v2, ok2 := parentSession.Values["direct_child_key"] + mtx.Unlock() + assert.True(t, ok2) + assert.Equal(t, "direct_child_val", v2) + + mtx.Lock() + parentSession.Values["direct_parent_key"] = "direct_parent_val" + mtx.Unlock() + + mtx.Lock() + v3, ok3 := inner.capturedSession.Values["direct_parent_key"] + mtx.Unlock() + assert.True(t, ok3) + assert.Equal(t, "direct_parent_val", v3) + + v, ok := GetSessionValue(ctx, "child_key") + assert.True(t, ok) + assert.Equal(t, "child_val", v) +} + +type sessionValuesAgent struct { + name string + seenParentValue any + capturedSession *runSession +} + +func (a *sessionValuesAgent) Name(context.Context) string { return a.name } +func (a *sessionValuesAgent) Description(context.Context) string { return "test" } +func (a *sessionValuesAgent) Run(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + if rc := getRunCtx(ctx); rc != nil { + a.capturedSession = rc.Session + } + a.seenParentValue, _ = GetSessionValue(ctx, "parent_key") + AddSessionValue(ctx, "child_key", "child_val") + + it, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Send(&AgentEvent{ + AgentName: a.name, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("ok", nil), + Role: schema.Assistant, + }, + }, + }) + gen.Close() + return it +} + +func TestAgentTool_InvokableRun(t *testing.T) { + // Create a context + ctx := context.Background() + + // Test cases + tests := []struct { + name string + agentResponses []*AgentEvent + request string + expectedOutput string + expectError bool + }{ + { + name: "successful model response", + agentResponses: []*AgentEvent{ + { + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Test response", nil), + Role: schema.Assistant, + }, + }, + }, + }, + request: `{"request":"Test request"}`, + expectedOutput: "Test response", + expectError: false, + }, + { + name: "successful tool call response", + agentResponses: []*AgentEvent{ + { + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.ToolMessage("Tool response", "test-id"), + Role: schema.Tool, + }, + }, + }, + }, + request: `{"request":"Test tool request"}`, + expectedOutput: "Tool response", + expectError: false, + }, + { + name: "invalid request JSON", + agentResponses: nil, + request: `invalid json`, + expectedOutput: "", + expectError: true, + }, + { + name: "no events returned", + agentResponses: []*AgentEvent{}, + request: `{"request":"Test request"}`, + expectedOutput: "", + expectError: true, + }, + { + name: "error in event", + agentResponses: []*AgentEvent{ + { + AgentName: "TestAgent", + Err: assert.AnError, + }, + }, + request: `{"request":"Test request"}`, + expectedOutput: "", + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a mock agent with the test responses + mockAgent_ := newMockAgentForTool("TestAgent", "Test agent description", tt.agentResponses) + + // Create an agentTool with the mock agent + agentTool_ := NewAgentTool(ctx, mockAgent_) + + // Call InvokableRun + output, err := agentTool_.(tool.InvokableTool).InvokableRun(ctx, tt.request) + + // Verify results + if tt.expectError { + assert.Error(t, err) + } else { + assert.NoError(t, err) + assert.Equal(t, tt.expectedOutput, output) + } + }) + } +} + +func TestGetReactHistory(t *testing.T) { + g := compose.NewGraph[string, []Message](compose.WithGenLocalState(func(ctx context.Context) (state *State) { + return &State{ + Messages: []Message{ + schema.UserMessage("user query"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "tool call id 1", Function: schema.FunctionCall{Name: "tool1", Arguments: "arguments1"}}}), + schema.ToolMessage("tool result 1", "tool call id 1", schema.WithToolName("tool1")), + schema.AssistantMessage("", []schema.ToolCall{{ID: "tool call id 2", Function: schema.FunctionCall{Name: "tool2", Arguments: "arguments2"}}}), + }, + } + })) + assert.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output []Message, err error) { + return getReactChatHistory(ctx, "DestAgentName") + }))) + assert.NoError(t, g.AddEdge(compose.START, "1")) + assert.NoError(t, g.AddEdge("1", compose.END)) + + ctx := context.Background() + ctx, _ = initRunCtx(ctx, "MyAgent", nil) + runner, err := g.Compile(ctx) + assert.NoError(t, err) + result, err := runner.Invoke(ctx, "") + assert.NoError(t, err) + assert.Equal(t, []Message{ + schema.UserMessage("user query"), + schema.UserMessage("For context: [MyAgent] called tool: `tool1` with arguments: arguments1."), + schema.UserMessage("For context: [MyAgent] `tool1` tool returned result: tool result 1."), + schema.UserMessage("For context: [MyAgent] called tool: `transfer_to_agent` with arguments: DestAgentName."), + schema.UserMessage("For context: [MyAgent] `transfer_to_agent` tool returned result: successfully transferred to agent [DestAgentName]."), + }, result) +} + +// mockAgentWithInputCapture implements the Agent interface for testing and captures the input it receives +type mockAgentWithInputCapture struct { + name string + description string + capturedInput []Message + responses []*AgentEvent +} + +func (a *mockAgentWithInputCapture) Name(_ context.Context) string { + return a.name +} + +func (a *mockAgentWithInputCapture) Description(_ context.Context) string { + return a.description +} + +func (a *mockAgentWithInputCapture) Run(_ context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + a.capturedInput = input.Messages + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + defer generator.Close() + + for _, event := range a.responses { + generator.Send(event) + + // If the event has an Exit action, stop sending events + if event.Action != nil && event.Action.Exit { + break + } + } + }() + + return iterator +} + +func newMockAgentWithInputCapture(name, description string, responses []*AgentEvent) *mockAgentWithInputCapture { + return &mockAgentWithInputCapture{ + name: name, + description: description, + responses: responses, + } +} + +func TestAgentToolWithOptions(t *testing.T) { + // Test Case 1: WithFullChatHistoryAsInput + t.Run("WithFullChatHistoryAsInput", func(t *testing.T) { + ctx := context.Background() + + // 1. Set up a mock agent that will capture the input it receives + mockAgent := newMockAgentWithInputCapture("test-agent", "a test agent", []*AgentEvent{ + { + AgentName: "test-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("done", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // 2. Create an agentTool with the option + agentTool := NewAgentTool(ctx, mockAgent, WithFullChatHistoryAsInput()) + + // 3. Set up a context with a chat history using a graph + history := []Message{ + schema.UserMessage("first user message"), + schema.AssistantMessage("first assistant response", nil), + } + + g := compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) (state *State) { + return &State{ + Messages: append(history, schema.AssistantMessage("tool call msg", nil)), + } + })) + + assert.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + // Run the tool within the graph context that has the state + _, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"request":"some ignored input"}`) + return "done", err + }))) + assert.NoError(t, g.AddEdge(compose.START, "1")) + assert.NoError(t, g.AddEdge("1", compose.END)) + + ctx, _ = initRunCtx(ctx, "react-agent", nil) + runner, err := g.Compile(ctx) + assert.NoError(t, err) + + // 4. Run the graph which will execute the tool with the state + _, err = runner.Invoke(ctx, "") + assert.NoError(t, err) + + // 5. Assert that the agent received the full history + // The agent should receive: history (minus last assistant message) + transfer messages + assert.Len(t, mockAgent.capturedInput, 4) // 2 from history + 2 transfer messages + assert.Equal(t, "first user message", mockAgent.capturedInput[0].Content) + assert.Equal(t, "For context: [react-agent] said: first assistant response.", mockAgent.capturedInput[1].Content) + assert.Equal(t, "For context: [react-agent] called tool: `transfer_to_agent` with arguments: test-agent.", mockAgent.capturedInput[2].Content) + assert.Equal(t, "For context: [react-agent] `transfer_to_agent` tool returned result: successfully transferred to agent [test-agent].", mockAgent.capturedInput[3].Content) + }) + + // Test Case 2: WithAgentInputSchema + t.Run("WithAgentInputSchema", func(t *testing.T) { + ctx := context.Background() + + // 1. Define a custom schema + customSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "custom_arg": { + Desc: "a custom argument", + Required: true, + Type: schema.String, + }, + }) + + // 2. Set up a mock agent to capture input + mockAgent := newMockAgentWithInputCapture("schema-agent", "agent with custom schema", []*AgentEvent{ + { + AgentName: "schema-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("schema processed", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // 3. Create agentTool with the custom schema option + agentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(customSchema)) + + // 4. Verify the Info() method returns the custom schema + info, err := agentTool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, customSchema, info.ParamsOneOf) + + // 5. Run the tool with arguments matching the custom schema + _, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"custom_arg":"hello world"}`) + assert.NoError(t, err) + + // 6. Assert that the agent received the correctly parsed argument + // With custom schema, the agent should receive the raw JSON as input + assert.Len(t, mockAgent.capturedInput, 1) + assert.Equal(t, `{"custom_arg":"hello world"}`, mockAgent.capturedInput[0].Content) + }) + + // Test Case 3: WithAgentInputSchema with complex schema + t.Run("WithAgentInputSchema_ComplexSchema", func(t *testing.T) { + ctx := context.Background() + + // 1. Define a complex custom schema with multiple parameters + complexSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "name": { + Desc: "user name", + Required: true, + Type: schema.String, + }, + "age": { + Desc: "user age", + Required: false, + Type: schema.Integer, + }, + "active": { + Desc: "user status", + Required: false, + Type: schema.Boolean, + }, + }) + + // 2. Set up a mock agent + mockAgent := newMockAgentWithInputCapture("complex-agent", "agent with complex schema", []*AgentEvent{ + { + AgentName: "complex-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("complex processed", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // 3. Create agentTool with the complex schema option + agentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(complexSchema)) + + // 4. Verify the Info() method returns the complex schema + info, err := agentTool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, complexSchema, info.ParamsOneOf) + + // 5. Run the tool with complex arguments + _, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"name":"John","age":30,"active":true}`) + assert.NoError(t, err) + + // 6. Assert that the agent received the complex JSON + assert.Len(t, mockAgent.capturedInput, 1) + assert.Equal(t, `{"name":"John","age":30,"active":true}`, mockAgent.capturedInput[0].Content) + }) + + // Test Case 4: Both options together + t.Run("BothOptionsTogether", func(t *testing.T) { + ctx := context.Background() + + // 1. Define a custom schema + customSchema := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": { + Desc: "search query", + Required: true, + Type: schema.String, + }, + }) + + // 2. Set up a mock agent + mockAgent := newMockAgentWithInputCapture("combined-agent", "agent with both options", []*AgentEvent{ + { + AgentName: "combined-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("combined processed", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // 3. Create agentTool with both options + agentTool := NewAgentTool(ctx, mockAgent, WithAgentInputSchema(customSchema), WithFullChatHistoryAsInput()) + + // 4. Set up a context with chat history using a graph + history := []Message{ + schema.UserMessage("previous conversation"), + schema.AssistantMessage("previous response", nil), + } + + g := compose.NewGraph[string, string](compose.WithGenLocalState(func(ctx context.Context) (state *State) { + return &State{ + Messages: append(history, schema.AssistantMessage("tool call", nil)), + } + })) + + assert.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + // Run the tool within the graph context that has the state + _, err = agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"query":"current query"}`) + return "done", err + }))) + assert.NoError(t, g.AddEdge(compose.START, "1")) + assert.NoError(t, g.AddEdge("1", compose.END)) + + ctx, _ = initRunCtx(ctx, "react-agent", nil) + runner, err := g.Compile(ctx) + assert.NoError(t, err) + + // 5. Run the graph which will execute the tool with the state + _, err = runner.Invoke(ctx, "") + assert.NoError(t, err) + + // 6. Verify both options work together + info, err := agentTool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, customSchema, info.ParamsOneOf) + + // The agent should receive full history + the custom query + assert.Len(t, mockAgent.capturedInput, 4) // 2 history + 2 transfer messages + assert.Equal(t, "previous conversation", mockAgent.capturedInput[0].Content) + assert.Equal(t, "For context: [react-agent] said: previous response.", mockAgent.capturedInput[1].Content) + assert.Equal(t, "For context: [react-agent] called tool: `transfer_to_agent` with arguments: combined-agent.", mockAgent.capturedInput[2].Content) + assert.Equal(t, "For context: [react-agent] `transfer_to_agent` tool returned result: successfully transferred to agent [combined-agent].", mockAgent.capturedInput[3].Content) + }) +} + +type fakeTCM struct{} + +func (f *fakeTCM) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + o := model.GetCommonOptions(&model.Options{}, opts...) + tc := schema.ToolCall{ID: "id-1", Type: "function"} + if len(o.Tools) > 0 { + tc.Function.Name = o.Tools[0].Name + } + tc.Function.Arguments = `{"request":"hello"}` + return schema.AssistantMessage("", []schema.ToolCall{tc}), nil +} +func (f *fakeTCM) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, _ := f.Generate(ctx, input, opts...) + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} +func (f *fakeTCM) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return f, nil +} + +type emitOnceModel struct{} + +func (e *emitOnceModel) Generate(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("inner2", nil), nil +} +func (e *emitOnceModel) Stream(ctx context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + m, _ := e.Generate(ctx, input) + return schema.StreamReaderFromArray([]*schema.Message{m}), nil +} +func (e *emitOnceModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return e, nil +} + +type emitEventsAgent struct{ events []*AgentEvent } + +func (e *emitEventsAgent) Name(context.Context) string { return "emit" } +func (e *emitEventsAgent) Description(context.Context) string { return "test" } +func (e *emitEventsAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] { + it, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + for _, ev := range e.events { + gen.Send(ev) + } + gen.Close() + }() + return it +} + +// spyAgent captures runSession from ctx in a single nested run +type spyAgent struct { + a Agent + mu sync.Mutex + captured *runSession +} + +func (s *spyAgent) Name(ctx context.Context) string { return s.a.Name(ctx) } +func (s *spyAgent) Description(ctx context.Context) string { return s.a.Description(ctx) } +func (s *spyAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + if rc := getRunCtx(ctx); rc != nil { + s.mu.Lock() + s.captured = rc.Session + s.mu.Unlock() + } + return s.a.Run(ctx, input, options...) +} + +func (s *spyAgent) getCaptured() *runSession { + s.mu.Lock() + defer s.mu.Unlock() + return s.captured +} + +func TestNestedAgentTool_RunPath(t *testing.T) { + ctx := context.Background() + + inner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner2", + Description: "leaf", + Model: &emitOnceModel{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: true}, + }) + inner2Spy := &spyAgent{a: inner2} + inner2Tool := NewAgentTool(ctx, inner2Spy) + + inner, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner", + Description: "mid", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: true, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{inner2Tool}}}, + }) + innerSpy := &spyAgent{a: inner} + innerTool := NewAgentTool(ctx, innerSpy) + + outer, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "outer", + Description: "top", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: true, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerTool}}}, + }) + + input := &AgentInput{Messages: []Message{schema.UserMessage("q")}} + ctx, outerRunCtx := initRunCtx(ctx, "outer", input) + r := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()}) + it := r.Run(ctx, []Message{schema.UserMessage("q")}) + + var target *AgentEvent + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming { + if ev.Output.MessageOutput.Message != nil && ev.Output.MessageOutput.Message.Content == "inner2" { + target = ev + break + } + } + } + if target == nil { + t.Fatalf("no inner2 event found in ephemerals") + } + + got := make([]string, len(target.RunPath)) + for i := range target.RunPath { + got[i] = target.RunPath[i].agentName + } + want := []string{"outer", "inner", "inner2"} + if len(got) != len(want) { + t.Fatalf("unexpected runPath len: got %d want %d: %+v", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("runPath mismatch at %d: got %s want %s; full: %+v", i, got[i], want[i], got) + } + } + + for _, w := range outerRunCtx.Session.getEvents() { + if w.AgentName != "outer" { + t.Fatalf("outer session contains non-outer event: %s", w.AgentName) + } + } + if innerSpy.getCaptured() == nil { + t.Fatalf("inner spy did not capture session") + } + for _, w := range innerSpy.getCaptured().getEvents() { + if w.AgentName != "inner" { + t.Fatalf("inner session contains non-inner event: %s", w.AgentName) + } + } + if inner2Spy.getCaptured() == nil { + t.Fatalf("inner2 spy did not capture session") + } + for _, w := range inner2Spy.getCaptured().getEvents() { + if w.AgentName != "inner2" { + t.Fatalf("inner2 session contains non-inner2 event: %s", w.AgentName) + } + } +} + +func TestNestedAgentTool_NoInternalEventsWhenDisabled(t *testing.T) { + ctx := context.Background() + + inner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner2", + Description: "leaf", + Model: &emitOnceModel{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: false}, + }) + inner2Tool := NewAgentTool(ctx, inner2) + + inner, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner", + Description: "mid", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: false, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{inner2Tool}}}, + }) + innerTool := NewAgentTool(ctx, inner) + + outer, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "outer", + Description: "top", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: false, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerTool}}}, + }) + + r := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()}) + it := r.Run(ctx, []Message{schema.UserMessage("q")}) + + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.AgentName == "inner2" { + t.Fatalf("inner2 internal event should not be emitted when disabled") + } + } +} + +func TestNestedAgentTool_InnerToolResultNotEmittedToOuter(t *testing.T) { + ctx := context.Background() + + innerTool := &simpleTool{name: "inner_tool", result: "inner_tool_result"} + inner, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner", + Description: "inner agent with tool", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerTool}}}, + }) + innerAgentTool := NewAgentTool(ctx, inner) + + outer, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "outer", + Description: "outer agent", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{innerAgentTool}}}, + }) + + r := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()}) + it := r.Run(ctx, []Message{schema.UserMessage("q")}) + + var allEvents []*AgentEvent + for { + ev, ok := it.Next() + if !ok { + break + } + allEvents = append(allEvents, ev) + } + + for _, ev := range allEvents { + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.Message != nil && + ev.Output.MessageOutput.Message.Role == schema.Tool && + ev.AgentName == "outer" && + ev.Output.MessageOutput.Message.Content == "inner_tool_result" { + t.Fatalf("inner agent's tool result (inner_tool_result) should not be emitted as outer agent's event, but got event with AgentName=%s, Content=%s", + ev.AgentName, ev.Output.MessageOutput.Message.Content) + } + } +} + +type simpleTool struct { + name string + result string +} + +func (s *simpleTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: s.name, Desc: "simple tool"}, nil +} + +func (s *simpleTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + return s.result, nil +} + +func TestAgentTool_InterruptWithoutCheckpoint(t *testing.T) { + ctx := context.Background() + ctx, _ = initRunCtx(ctx, "TestAgent", &AgentInput{Messages: []Message{}}) + + interrupted := &AgentEvent{AgentName: "TestAgent"} + interrupted.Action = StatefulInterrupt(ctx, "info", "state").Action + + err := compositeInterruptFromLast(ctx, &bridgeStore{}, interrupted) + if err == nil { + t.Fatalf("expected error for interrupt without checkpoint") + } + if !strings.Contains(err.Error(), "interrupt occurred but checkpoint data is missing") { + t.Fatalf("unexpected error: %v", err) + } +} + +func compositeInterruptFromLast(ctx context.Context, ms *bridgeStore, lastEvent *AgentEvent) error { + if lastEvent == nil || lastEvent.Action == nil || lastEvent.Action.Interrupted == nil { + return nil + } + data, existed, err := ms.Get(ctx, bridgeCheckpointID) + if err != nil { + return fmt.Errorf("failed to get interrupt info: %w", err) + } + if !existed { + return fmt.Errorf("interrupt occurred but checkpoint data is missing") + } + return tool.CompositeInterrupt(ctx, "agent tool interrupt", data, lastEvent.Action.internalInterrupted) +} + +func TestAgentTool_InvokableRun_FinalOnly(t *testing.T) { + ctx := context.Background() + + inner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner2", + Description: "leaf", + Model: &emitOnceModel{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: true}, + }) + invTool := NewAgentTool(ctx, inner2) + out, err := invTool.(tool.InvokableTool).InvokableRun(ctx, `{"request":"q"}`) + if err != nil { + t.Fatalf("invokable run error: %v", err) + } + if out != "inner2" { + t.Fatalf("unexpected output: %s", out) + } +} + +type streamingAgent struct{} + +func (s *streamingAgent) Name(context.Context) string { return "stream" } +func (s *streamingAgent) Description(context.Context) string { return "test" } +func (s *streamingAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] { + it, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + mv := &MessageVariant{IsStreaming: true, MessageStream: schema.StreamReaderFromArray([]Message{schema.AssistantMessage("1", nil), schema.AssistantMessage("2", nil)})} + gen.Send(&AgentEvent{AgentName: "stream", Output: &AgentOutput{MessageOutput: mv}}) + mv = &MessageVariant{IsStreaming: true, MessageStream: schema.StreamReaderFromArray([]Message{schema.AssistantMessage("a", nil), schema.AssistantMessage("b", nil)})} + gen.Send(&AgentEvent{AgentName: "stream", Output: &AgentOutput{MessageOutput: mv}}) + gen.Close() + }() + return it +} + +func TestAgentTool_InvokableRun_StreamingVariant(t *testing.T) { + ctx := context.Background() + agent := &streamingAgent{} + it := NewAgentTool(ctx, agent) + out, err := it.(tool.InvokableTool).InvokableRun(ctx, `{"request":"q"}`) + if err != nil { + t.Fatalf("invokable run error: %v", err) + } + if out != "ab" { + t.Fatalf("unexpected output: %s", out) + } +} + +func TestSequentialWorkflow_WithChatModelAgentTool_NestedRunPathAndSessions(t *testing.T) { + ctx := context.Background() + + inner2, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner2", + Description: "leaf", + Model: &emitOnceModel{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: true}, + }) + inner2Spy := &spyAgent{a: inner2} + inner2ToolSpy := NewAgentTool(ctx, inner2Spy) + + innerWithSpy, _ := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner", + Description: "mid", + Model: &fakeTCM{}, + ToolsConfig: ToolsConfig{EmitInternalEvents: true, ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{inner2ToolSpy}}}, + }) + innerSpy := &spyAgent{a: innerWithSpy} + + outer, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "outer-seq", + Description: "workflow", + SubAgents: []Agent{innerSpy}, + }) + if err != nil { + t.Fatalf("new sequential agent err: %v", err) + } + + input := &AgentInput{Messages: []Message{schema.UserMessage("q")}} + ctx, outerRunCtx := initRunCtx(ctx, "outer-seq", input) + r := NewRunner(ctx, RunnerConfig{Agent: outer, EnableStreaming: false, CheckPointStore: newBridgeStore()}) + it := r.Run(ctx, []Message{schema.UserMessage("q")}) + + var target *AgentEvent + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming { + if ev.Output.MessageOutput.Message != nil && ev.Output.MessageOutput.Message.Content == "inner2" { + target = ev + break + } + } + } + if target == nil { + t.Fatalf("no inner2 event found") + } + + got := make([]string, len(target.RunPath)) + for i := range target.RunPath { + got[i] = target.RunPath[i].agentName + } + want := []string{"outer-seq", "inner", "inner2"} + if len(got) != len(want) { + t.Fatalf("unexpected runPath len: got %d want %d: %+v", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("runPath mismatch at %d: got %s want %s; full: %+v", i, got[i], want[i], got) + } + } + + for _, w := range outerRunCtx.Session.getEvents() { + if w.AgentName != "outer-seq" { + t.Fatalf("outer session contains non-outer event: %s", w.AgentName) + } + } + if innerSpy.getCaptured() == nil { + t.Fatalf("inner spy did not capture session") + } + for _, w := range innerSpy.getCaptured().getEvents() { + if w.AgentName != "inner" { + t.Fatalf("inner session contains non-inner event: %s", w.AgentName) + } + } + if inner2Spy.getCaptured() == nil { + t.Fatalf("inner2 spy did not capture session") + } + for _, w := range inner2Spy.getCaptured().getEvents() { + if w.AgentName != "inner2" { + t.Fatalf("inner2 session contains non-inner2 event: %s", w.AgentName) + } + } +} + +func TestRunPathGating_IgnoresInnerExitAndAllowsOutput(t *testing.T) { + ctx := context.Background() + + innerExit := &AgentEvent{Action: &AgentAction{Exit: true}, RunPath: []RunStep{{agentName: "inner"}}} + finalOut := EventFromMessage(schema.AssistantMessage("ok", nil), nil, schema.Assistant, "") + + sub := &emitEventsAgent{events: []*AgentEvent{innerExit, finalOut}} + fa := toFlowAgent(ctx, sub) + + it := fa.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("q")}}) + + var sawFinal bool + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming { + if ev.Output.MessageOutput.Message != nil && ev.Output.MessageOutput.Message.Content == "ok" { + sawFinal = true + } + } + } + if !sawFinal { + t.Fatalf("final output not observed; parent may have exited on inner Exit action") + } +} + +func TestRunPathGating_IgnoresInnerTransfer(t *testing.T) { + ctx := context.Background() + + innerTransfer := &AgentEvent{Action: NewTransferToAgentAction("ghost"), RunPath: []RunStep{{agentName: "inner"}}} + finalOut := EventFromMessage(schema.AssistantMessage("done", nil), nil, schema.Assistant, "") + + sub := &emitEventsAgent{events: []*AgentEvent{innerTransfer, finalOut}} + fa := toFlowAgent(ctx, sub) + + it := fa.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("q")}}) + + var outputs int + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil && !ev.Output.MessageOutput.IsStreaming { + if ev.Output.MessageOutput.Message != nil { + outputs++ + } + } + } + if outputs == 0 { + t.Fatalf("no outputs observed; parent may have transferred on inner transfer action") + } +} + +type streamAgent struct{} + +func (s *streamAgent) Name(context.Context) string { return "s" } +func (s *streamAgent) Description(context.Context) string { return "s" } +func (s *streamAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] { + it, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + frames := []*schema.Message{ + schema.AssistantMessage("hello ", nil), + schema.AssistantMessage("world", nil), + } + stream := schema.StreamReaderFromArray(frames) + gen.Send(EventFromMessage(nil, stream, schema.Assistant, "")) + gen.Close() + }() + return it +} + +func TestInvokableAgentTool_InfoAndRun(t *testing.T) { + ctx := context.Background() + + at := NewAgentTool(ctx, &streamAgent{}) + info, err := at.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "s", info.Name) + assert.Equal(t, "s", info.Desc) + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + found := false + for _, r := range js.Required { + if r == "request" { + found = true + break + } + } + assert.True(t, found) + prop, ok := js.Properties.Get("request") + assert.True(t, ok) + assert.Equal(t, string(schema.String), prop.Type) + + custom := schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "x": {Desc: "arg", Required: true, Type: schema.String}, + }) + at2 := NewAgentTool(ctx, &streamAgent{}, WithAgentInputSchema(custom)) + info2, err := at2.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, custom, info2.ParamsOneOf) + out, err := at.(tool.InvokableTool).InvokableRun(ctx, `{"request":"x"}`) + assert.NoError(t, err) + assert.Equal(t, "hello world", out) +} + +type emptyAgent struct{} + +func (e *emptyAgent) Name(context.Context) string { return "empty" } +func (e *emptyAgent) Description(context.Context) string { return "empty" } +func (e *emptyAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] { + it, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { gen.Close() }() + return it +} + +type noOutputAgent struct{} + +func (n *noOutputAgent) Name(context.Context) string { return "no" } +func (n *noOutputAgent) Description(context.Context) string { return "no" } +func (n *noOutputAgent) Run(context.Context, *AgentInput, ...AgentRunOption) *AsyncIterator[*AgentEvent] { + it, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { gen.Send(&AgentEvent{}); gen.Close() }() + return it +} + +func TestInvokableAgentTool_ErrorCases(t *testing.T) { + ctx := context.Background() + + atEmpty := NewAgentTool(ctx, &emptyAgent{}) + out, err := atEmpty.(tool.InvokableTool).InvokableRun(ctx, `{"request":"x"}`) + assert.Equal(t, "", out) + assert.Error(t, err) + + atNo := NewAgentTool(ctx, &noOutputAgent{}) + out2, err := atNo.(tool.InvokableTool).InvokableRun(ctx, `{"request":"x"}`) + assert.NoError(t, err) + assert.Equal(t, "", out2) +} + +func TestCrossTypeAgentToolGracefulError(t *testing.T) { + ctx := context.Background() + + innerModel := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("inner result"), nil + }, + } + + innerAgent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticInner", + Description: "An agentic agent used as a tool", + Model: innerModel, + }) + require.NoError(t, err) + + agenticAgentTool := NewTypedAgentTool(ctx, TypedAgent[*schema.AgenticMessage](innerAgent)) + + var outerCallCount int32 + outerModel := &mockChatModelForAttack{ + generateFn: func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&outerCallCount, 1) + if count == 1 { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "c1", Function: schema.FunctionCall{Name: "AgenticInner", Arguments: `{"request":"test"}`}}, + }, + }, nil + } + return schema.AssistantMessage("done", nil), nil + }, + } + + outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterMessageAgent", + Description: "A Message agent using an AgenticMessage sub-agent tool", + Model: outerModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{agenticAgentTool}, + }, + }, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: outerAgent, EnableStreaming: true}) + iter := runner.Query(ctx, "test cross-type") + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + t.Logf("Cross-type error message: %v", event.Err) + } + } + + if capturedErr == nil { + t.Log("DESIGN CONCERN: Cross-type agent tool (AgenticMessage sub-agent in Message agent) " + + "only errors at event forwarding time when streaming is enabled. " + + "The error check happens in the gen.Send path, which is only exercised " + + "when the outer agent actually calls the tool AND streaming is enabled. " + + "Without streaming, the tool result is returned as a string, so no type mismatch occurs.") + } else { + assert.Contains(t, capturedErr.Error(), "cross-message-type", + "Error should mention cross-message-type incompatibility") + } +} diff --git a/adk/agentic_callback_integration_test.go b/adk/agentic_callback_integration_test.go new file mode 100644 index 0000000..2473091 --- /dev/null +++ b/adk/agentic_callback_integration_test.go @@ -0,0 +1,268 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type agenticCallbackRecorder struct { + mu sync.Mutex + onStartCalled bool + onEndCalled bool + runInfo *callbacks.RunInfo + inputReceived *TypedAgentCallbackInput[*schema.AgenticMessage] + eventsReceived []*TypedAgentEvent[*schema.AgenticMessage] + eventsDone chan struct{} + closeOnce sync.Once +} + +func (r *agenticCallbackRecorder) getOnStartCalled() bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.onStartCalled +} + +func (r *agenticCallbackRecorder) getOnEndCalled() bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.onEndCalled +} + +func (r *agenticCallbackRecorder) getEventsReceived() []*TypedAgentEvent[*schema.AgenticMessage] { + r.mu.Lock() + defer r.mu.Unlock() + result := make([]*TypedAgentEvent[*schema.AgenticMessage], len(r.eventsReceived)) + copy(result, r.eventsReceived) + return result +} + +func newAgenticRecordingHandler(recorder *agenticCallbackRecorder) callbacks.Handler { + recorder.eventsDone = make(chan struct{}) + return callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgenticAgent { + return ctx + } + recorder.mu.Lock() + defer recorder.mu.Unlock() + recorder.onStartCalled = true + recorder.runInfo = info + if agentInput := ConvTypedCallbackInput[*schema.AgenticMessage](input); agentInput != nil { + recorder.inputReceived = agentInput + } + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgenticAgent { + return ctx + } + recorder.mu.Lock() + recorder.onEndCalled = true + recorder.runInfo = info + recorder.mu.Unlock() + + if agentOutput := ConvTypedCallbackOutput[*schema.AgenticMessage](output); agentOutput != nil { + if agentOutput.Events != nil { + go func() { + defer recorder.closeOnce.Do(func() { close(recorder.eventsDone) }) + for { + event, ok := agentOutput.Events.Next() + if !ok { + break + } + recorder.mu.Lock() + recorder.eventsReceived = append(recorder.eventsReceived, event) + recorder.mu.Unlock() + } + }() + return ctx + } + } + recorder.closeOnce.Do(func() { close(recorder.eventsDone) }) + return ctx + }). + Build() +} + +func TestAgenticCallback(t *testing.T) { + ctx := context.Background() + + expectedContent := "This is the test response content" + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg(expectedContent), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestChatAgent", + Description: "Test chat agent", + Instruction: "You are a test agent", + Model: m, + }) + require.NoError(t, err) + + recorder := &agenticCallbackRecorder{} + handler := newAgenticRecordingHandler(recorder) + + var agentEvents []*TypedAgentEvent[*schema.AgenticMessage] + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + event, ok := iter.Next() + if !ok { + break + } + agentEvents = append(agentEvents, event) + } + + <-recorder.eventsDone + assertAgenticEventRoleFields(t, agentEvents) + + t.Run("OnStart_Invocation", func(t *testing.T) { + assert.True(t, recorder.getOnStartCalled(), "OnStart should be called") + require.NotNil(t, recorder.inputReceived, "Input should be received") + require.NotNil(t, recorder.inputReceived.Input, "AgentInput should be set") + assert.Len(t, recorder.inputReceived.Input.Messages, 1) + }) + + t.Run("OnEnd_Invocation", func(t *testing.T) { + assert.True(t, recorder.getOnEndCalled(), "OnEnd should be called") + assert.Len(t, recorder.getEventsReceived(), 1) + }) + + t.Run("RunInfo_Fields", func(t *testing.T) { + require.NotNil(t, recorder.runInfo) + assert.Equal(t, "TestChatAgent", recorder.runInfo.Name) + assert.Equal(t, ComponentOfAgenticAgent, recorder.runInfo.Component) + }) + + t.Run("Events_MatchAgentOutput", func(t *testing.T) { + require.NotEmpty(t, agentEvents, "Agent should emit events") + received := recorder.getEventsReceived() + require.NotEmpty(t, received, "Callback should receive events") + + require.Len(t, received, 1, "Callback should receive exactly 1 event") + require.NotNil(t, received[0].Output) + require.NotNil(t, received[0].Output.MessageOutput) + require.NotNil(t, received[0].Output.MessageOutput.Message) + assert.Equal(t, expectedContent, agenticTextContent(received[0].Output.MessageOutput.Message)) + }) +} + +func TestAgenticCallbackMultipleHandlers(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("test response"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test agent", + Model: m, + }) + require.NoError(t, err) + + recorder1 := &agenticCallbackRecorder{} + recorder2 := &agenticCallbackRecorder{} + handler1 := newAgenticRecordingHandler(recorder1) + handler2 := newAgenticRecordingHandler(recorder2) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler1, handler2)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + <-recorder1.eventsDone + <-recorder2.eventsDone + + assert.True(t, recorder1.getOnStartCalled(), "Handler1 OnStart should be called") + assert.True(t, recorder2.getOnStartCalled(), "Handler2 OnStart should be called") + assert.True(t, recorder1.getOnEndCalled(), "Handler1 OnEnd should be called") + assert.True(t, recorder2.getOnEndCalled(), "Handler2 OnEnd should be called") + + assert.NotEmpty(t, recorder1.getEventsReceived(), "Handler1 should receive events") + assert.NotEmpty(t, recorder2.getEventsReceived(), "Handler2 should receive events") +} + +func TestCoverage_WrapAgenticIterWithOnEnd(t *testing.T) { + ctx := context.Background() + + var onEndCalled bool + handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component == ComponentOfAgenticAgent { + onEndCalled = true + } + return ctx + }). + Build() + + ctx = initAgenticCallbacks(ctx, "test-agent", "ChatModel", + WithCallbacks(handler)) + + cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{ + Input: &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + }, + } + ctx = callbacks.OnStart(ctx, cbInput) + + origIter, origGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer origGen.Close() + origGen.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("done"), + }, + }, + }) + }() + + wrappedIter := wrapAgenticIterWithOnEnd(ctx, origIter) + + for { + _, ok := wrappedIter.Next() + if !ok { + break + } + } + + assert.True(t, onEndCalled, "OnEnd callback should have been called") +} diff --git a/adk/agentic_integration_test.go b/adk/agentic_integration_test.go new file mode 100644 index 0000000..2c6df1e --- /dev/null +++ b/adk/agentic_integration_test.go @@ -0,0 +1,665 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "encoding/json" + "sync/atomic" + "testing" + "time" + + "github.com/eino-contrib/jsonschema" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +func agenticMsg(text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: text}), + }, + } +} + +func agenticTextContent(msg *schema.AgenticMessage) string { + for _, b := range msg.ContentBlocks { + if b.AssistantGenText != nil { + return b.AssistantGenText.Text + } + } + return "" +} + +func TestAgenticIntegration_ChatModelSingleShot(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("Handled internally with tool result: 42"), nil + }, + } + + dummyTool := newSlowTool("calculator", 0, "42") + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "ToolCallAgent", + Description: "Agent with tools for agentic model", + Instruction: "You are a calculator.", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{dummyTool}, + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + + iter := runner.Query(ctx, "What is 6*7?") + + var events []*TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + require.Len(t, events, 1) + assertAgenticEventRoleFields(t, events) + lastEvent := events[len(events)-1] + require.Nil(t, lastEvent.Err) + require.NotNil(t, lastEvent.Output) + require.NotNil(t, lastEvent.Output.MessageOutput) + assert.Equal(t, "Handled internally with tool result: 42", + agenticTextContent(lastEvent.Output.MessageOutput.Message)) +} + +func TestAgenticIntegration_ChatModelToolsPassedViaOptions(t *testing.T) { + ctx := context.Background() + + var receivedTools []*schema.ToolInfo + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + o := model.GetCommonOptions(&model.Options{}, opts...) + receivedTools = o.Tools + return agenticMsg("done"), nil + }, + } + + dummyTool := newSlowTool("my_tool", 0, "result") + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "ToolOptAgent", + Description: "Agent verifying tools are passed via options", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{dummyTool}, + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + iter := runner.Query(ctx, "test tools") + for { + _, ok := iter.Next() + if !ok { + break + } + } + + require.NotNil(t, receivedTools, "tools should be passed via model.Options") + require.Len(t, receivedTools, 1) + assert.Equal(t, "my_tool", receivedTools[0].Name) +} + +func TestAgenticIntegration_StreamingWithRunner(t *testing.T) { + ctx := context.Background() + + chunk1 := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}), + }, + } + chunk2 := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + } + + m := &mockAgenticModel{ + streamFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(chunk1, nil) + w.Send(chunk2, nil) + }() + return r, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "StreamRunner", + Description: "Streaming runner agent", + Model: m, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + iter := runner.Query(ctx, "stream me") + + event, ok := iter.Next() + require.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + if event.Output.MessageOutput.IsStreaming { + require.NotNil(t, event.Output.MessageOutput.MessageStream) + var chunks []*schema.AgenticMessage + for { + chunk, err := event.Output.MessageOutput.MessageStream.Recv() + if err != nil { + break + } + chunks = append(chunks, chunk) + } + assert.Equal(t, 2, len(chunks)) + } else { + assert.NotNil(t, event.Output.MessageOutput.Message) + } + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestAgenticIntegration_CancelDuringExecution(t *testing.T) { + ctx := context.Background() + + modelStarted := make(chan struct{}, 1) + modelBlocked := make(chan struct{}) + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + select { + case modelStarted <- struct{}{}: + default: + } + select { + case <-modelBlocked: + return agenticMsg("should not reach"), nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "CancelAgent", + Description: "cancel test", + Model: m, + }) + require.NoError(t, err) + + cancelCtx, cancel := context.WithCancel(ctx) + defer cancel() + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + iter := runner.Run(cancelCtx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }) + + <-modelStarted + cancel() + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should propagate cancel error") + assert.ErrorIs(t, capturedErr, context.Canceled) +} + +func TestAgenticIntegration_CancelWithTimeout(t *testing.T) { + ctx := context.Background() + + sa := &myAgenticAgent{ + name: "slow-agent", + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + select { + case <-time.After(10 * time.Second): + generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("slow response"), + }, + }, + }) + case <-ctx.Done(): + generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + Err: ctx.Err(), + }) + } + }() + return iter + }, + } + + timeoutCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: sa, + }) + iter := runner.Run(timeoutCtx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("slow request"), + }) + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + + require.Error(t, capturedErr, "should get timeout/cancel error") + assert.ErrorIs(t, capturedErr, context.DeadlineExceeded) +} +func TestAgenticIntegration_AgentTool(t *testing.T) { + ctx := context.Background() + + innerModel := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("inner tool result"), nil + }, + } + + innerAgent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "InnerAgent", + Description: "An agent used as a tool", + Model: innerModel, + }) + require.NoError(t, err) + + agentTool := NewTypedAgentTool(ctx, TypedAgent[*schema.AgenticMessage](innerAgent)) + require.NotNil(t, agentTool) + + info, err := agentTool.Info(ctx) + require.NoError(t, err) + assert.Equal(t, "InnerAgent", info.Name) + assert.Equal(t, "An agent used as a tool", info.Desc) + + outerModel := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("outer response after inner tool"), nil + }, + } + + outerAgent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "OuterAgent", + Description: "Outer agent with agent tool", + Model: outerModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{agentTool}, + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: outerAgent, + }) + iter := runner.Query(ctx, "delegate to inner") + + var events []*TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + require.NotEmpty(t, events) + assertAgenticEventRoleFields(t, events) + lastEvent := events[len(events)-1] + assert.Nil(t, lastEvent.Err) + assert.NotNil(t, lastEvent.Output) +} +func TestAgenticIntegration_InterruptEventFormation(t *testing.T) { + ctx := context.Background() + + t.Run("simple interrupt", func(t *testing.T) { + agent := &myAgenticAgent{ + name: "int-agent", + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, "approval needed") + intEvent.Action.Interrupted.Data = "approval data" + generator.Send(intEvent) + }() + return iter + }, + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + iter := runner.Query(ctx, "interrupt test") + + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + } + + require.NotNil(t, interruptEvent) + assert.Equal(t, "approval data", interruptEvent.Action.Interrupted.Data) + require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts) + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID) + assert.Equal(t, "approval needed", interruptEvent.Action.Interrupted.InterruptContexts[0].Info) + assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause) + }) + + t.Run("stateful interrupt", func(t *testing.T) { + agent := &myAgenticAgent{ + name: "st-agent", + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + intEvent := TypedStatefulInterrupt[*schema.AgenticMessage](ctx, "state interrupt", "my-state") + intEvent.Action.Interrupted.Data = "stateful data" + generator.Send(intEvent) + }() + return iter + }, + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + iter := runner.Query(ctx, "stateful test") + + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + } + + require.NotNil(t, interruptEvent) + assert.Equal(t, "stateful data", interruptEvent.Action.Interrupted.Data) + require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts) + assert.Equal(t, "state interrupt", interruptEvent.Action.Interrupted.InterruptContexts[0].Info) + }) +} +func TestAgenticIntegration_CheckpointInterruptResume(t *testing.T) { + ctx := context.Background() + + var resumeCalled int32 + agent := &myAgenticAgent{ + name: "ckpt-agent", + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "ckpt-agent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("before interrupt"), + }, + }, + }) + intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, "need approval") + intEvent.Action.Interrupted.Data = "approval data" + generator.Send(intEvent) + }() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + atomic.StoreInt32(&resumeCalled, 1) + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "ckpt-agent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("after resume"), + }, + }, + }) + }() + return iter + }, + } + + store := newMyStore() + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "checkpoint test", WithCheckPointID("ckpt-1")) + + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + var preInterruptOutputs []string + for { + event, ok := iter.Next() + if !ok { + break + } + require.Nil(t, event.Err) + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil { + preInterruptOutputs = append(preInterruptOutputs, agenticTextContent(event.Output.MessageOutput.Message)) + } + } + + require.NotNil(t, interruptEvent, "should receive interrupt event") + assert.Contains(t, preInterruptOutputs, "before interrupt") + require.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts) + + interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID + require.NotEmpty(t, interruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, "ckpt-1", &ResumeParams{ + Targets: map[string]any{ + interruptID: nil, + }, + }) + require.NoError(t, err) + + var postResumeOutputs []string + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Err != nil { + t.Fatalf("unexpected error during resume: %v", event.Err) + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil { + postResumeOutputs = append(postResumeOutputs, agenticTextContent(event.Output.MessageOutput.Message)) + } + } + + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeCalled), "resume function should have been called") + assert.Contains(t, postResumeOutputs, "after resume") +} + +func TestAgenticIntegration_CheckpointWithMCPListToolsResult(t *testing.T) { + ctx := context.Background() + + inputSchemaJSON := `{ + "type": "object", + "properties": { + "query": {"type": "string", "description": "search query"}, + "limit": {"type": "integer", "description": "max results"} + }, + "required": ["query"] + }` + var inputSchema jsonschema.Schema + require.NoError(t, json.Unmarshal([]byte(inputSchemaJSON), &inputSchema)) + + mcpMsg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeMCPListToolsResult, + MCPListToolsResult: &schema.MCPListToolsResult{ + ServerLabel: "test-server", + Tools: []*schema.MCPListToolsItem{ + { + Name: "search", + Description: "search the web", + InputSchema: &inputSchema, + }, + }, + }, + }, + schema.NewContentBlock(&schema.AssistantGenText{Text: "here are tools"}), + }, + } + + var resumeCalled int32 + agent := &myAgenticAgent{ + name: "mcp-agent", + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer gen.Close() + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "mcp-agent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: mcpMsg}, + }, + }) + gen.Send(TypedInterrupt[*schema.AgenticMessage](ctx, "approve tools")) + }() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + atomic.StoreInt32(&resumeCalled, 1) + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer gen.Close() + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "mcp-agent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("tools approved")}, + }, + }) + }() + return iter + }, + } + + store := newMyStore() + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "list tools", WithCheckPointID("mcp-1")) + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + ev, ok := iter.Next() + if !ok { + break + } + require.Nil(t, ev.Err) + if ev.Action != nil && ev.Action.Interrupted != nil { + interruptEvent = ev + } + } + require.NotNil(t, interruptEvent) + interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID + + resumeIter, err := runner.ResumeWithParams(ctx, "mcp-1", &ResumeParams{ + Targets: map[string]any{interruptID: nil}, + }) + require.NoError(t, err) + + var outputs []string + for { + ev, ok := resumeIter.Next() + if !ok { + break + } + require.Nil(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.Message != nil { + outputs = append(outputs, agenticTextContent(ev.Output.MessageOutput.Message)) + } + } + + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeCalled)) + assert.Contains(t, outputs, "tools approved") +} diff --git a/adk/agentic_react_test.go b/adk/agentic_react_test.go new file mode 100644 index 0000000..df505c6 --- /dev/null +++ b/adk/agentic_react_test.go @@ -0,0 +1,1288 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type agenticAgentEvent = TypedAgentEvent[*schema.AgenticMessage] + +func agenticToolCallMsg(toolName, callID, args string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{Name: toolName, CallID: callID, Arguments: args}, + }, + }, + } +} + +type sequentialAgenticModel struct { + responses []*schema.AgenticMessage + callCount int32 +} + +func (m *sequentialAgenticModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&m.callCount, 1) - 1 + if int(idx) >= len(m.responses) { + return nil, fmt.Errorf("sequentialAgenticModel: no more responses (call #%d)", idx) + } + return m.responses[idx], nil +} + +func (m *sequentialAgenticModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + result, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { defer w.Close(); w.Send(result, nil) }() + return r, nil +} + +type agenticEchoTool struct { + name string +} + +func (t *agenticEchoTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "echoes input"}, nil +} + +func (t *agenticEchoTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + return "echo:" + argumentsInJSON, nil +} + +type agenticInterruptTool struct { + name string +} + +func (t *agenticInterruptTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "interrupts on first call, returns on resume"}, nil +} + +func (t *agenticInterruptTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + wasInterrupted, _, _ := tool.GetInterruptState[any](ctx) + if !wasInterrupted { + return "", tool.Interrupt(ctx, "need_approval") + } + isResume, hasData, data := tool.GetResumeContext[string](ctx) + if isResume && hasData { + return "approved:" + data, nil + } + return "resumed_no_data", nil +} + +type agenticArgCaptureTool struct { + name string + onInvoke func(args string) string +} + +func (t *agenticArgCaptureTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "captures args"}, nil +} + +func (t *agenticArgCaptureTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + return t.onInvoke(argumentsInJSON), nil +} + +type agenticSignalTool struct { + name string + started chan struct{} + result string + done chan struct{} + once sync.Once +} + +func (t *agenticSignalTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: "blocks until finish() is called"}, nil +} + +func (t *agenticSignalTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + t.once.Do(func() { t.done = make(chan struct{}) }) + select { + case t.started <- struct{}{}: + default: + } + <-t.done + return t.result, nil +} + +func (t *agenticSignalTool) finish() { + t.once.Do(func() { t.done = make(chan struct{}) }) + close(t.done) +} + +type agenticReactTestStore struct { + m map[string][]byte +} + +func (s *agenticReactTestStore) Set(_ context.Context, key string, value []byte) error { + s.m[key] = value + return nil +} + +func (s *agenticReactTestStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.m[key] + return v, ok, nil +} + +func newAgenticAgent(t *testing.T, ctx context.Context, mdl model.BaseModel[*schema.AgenticMessage], tools []tool.BaseTool) TypedAgent[*schema.AgenticMessage] { + t.Helper() + config := &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: t.Name(), + Description: "test agentic agent", + Model: mdl, + } + if len(tools) > 0 { + config.ToolsConfig = ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: tools, + }, + } + } + agent, err := NewTypedChatModelAgent(ctx, config) + require.NoError(t, err) + return agent +} + +func newAgenticRunner(t *testing.T, ctx context.Context, mdl model.BaseModel[*schema.AgenticMessage], tools []tool.BaseTool) *TypedRunner[*schema.AgenticMessage] { + t.Helper() + agent := newAgenticAgent(t, ctx, mdl, tools) + return NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) +} + +func newAgenticRunnerWithStore(t *testing.T, ctx context.Context, mdl model.BaseModel[*schema.AgenticMessage], tools []tool.BaseTool, store CheckPointStore) *TypedRunner[*schema.AgenticMessage] { + t.Helper() + agent := newAgenticAgent(t, ctx, mdl, tools) + return NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) +} + +func drainAgenticEvents(iter *AsyncIterator[*agenticAgentEvent]) []*agenticAgentEvent { + var events []*agenticAgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + } + return events +} + +func lastAgenticEvent(events []*agenticAgentEvent) *agenticAgentEvent { + if len(events) == 0 { + return nil + } + return events[len(events)-1] +} + +func firstAgenticEventError(events []*agenticAgentEvent) error { + for _, ev := range events { + if ev.Err != nil { + return ev.Err + } + } + return nil +} + +func findInterruptEvent(events []*agenticAgentEvent) *agenticAgentEvent { + for _, ev := range events { + if ev.Action != nil && ev.Action.Interrupted != nil { + return ev + } + } + return nil +} + +func TestAgenticReact_BasicInvoke(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "call-1", `"hello"`), + agenticMsg("done: echo result received"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + events := drainAgenticEvents(runner.Query(ctx, "test input")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "done: echo result received", agenticTextContent(last.Output.MessageOutput.Message)) + assert.Equal(t, int32(2), atomic.LoadInt32(&mdl.callCount)) +} + +func TestAgenticReact_MultiTurnToolCalling(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "call-1", `"step1"`), + agenticToolCallMsg("echo", "call-2", `"step2"`), + agenticToolCallMsg("echo", "call-3", `"step3"`), + agenticMsg("all done"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + events := drainAgenticEvents(runner.Query(ctx, "do three steps")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "all done", agenticTextContent(last.Output.MessageOutput.Message)) + assert.Equal(t, int32(4), atomic.LoadInt32(&mdl.callCount)) +} + +func TestAgenticReact_Stream(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "call-1", `"hello"`), + agenticMsg("stream done"), + }, + } + + agent := newAgenticAgent(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + events := drainAgenticEvents(runner.Query(ctx, "stream test")) + + var finalText string + for _, ev := range events { + if ev.Output != nil && ev.Output.MessageOutput != nil { + msg, err := ev.Output.MessageOutput.GetMessage() + if err == nil && msg != nil { + txt := agenticTextContent(msg) + if txt != "" { + finalText = txt + } + } + } + } + + assert.Equal(t, "stream done", finalText) +} + +func TestAgenticReact_MaxIterations(t *testing.T) { + ctx := context.Background() + + t.Run("within_limit", func(t *testing.T) { + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("echo", "c1", `"1"`), + agenticToolCallMsg("echo", "c2", `"2"`), + agenticMsg("done within limit"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{&agenticEchoTool{name: "echo"}}) + events := drainAgenticEvents(runner.Query(ctx, "go")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "done within limit", agenticTextContent(last.Output.MessageOutput.Message)) + }) + + t.Run("exceeded", func(t *testing.T) { + responses := make([]*schema.AgenticMessage, 25) + for i := range responses { + responses[i] = agenticToolCallMsg("echo", fmt.Sprintf("c%d", i), `"x"`) + } + + mdl := &sequentialAgenticModel{responses: responses} + config := &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "exceed-agent", + Description: "test max iterations exceeded", + Model: mdl, + MaxIterations: 3, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&agenticEchoTool{name: "echo"}}, + }, + }, + } + agent, err := NewTypedChatModelAgent(ctx, config) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + events := drainAgenticEvents(runner.Query(ctx, "go")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.NotNil(t, last.Err) + assert.ErrorIs(t, last.Err, ErrExceedMaxIterations) + }) +} + +func TestAgenticReact_ReturnDirectly(t *testing.T) { + ctx := context.Background() + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + // Model calls the return-directly tool. + agenticToolCallMsg("direct", "call-1", `"final answer"`), + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: t.Name(), + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&agenticEchoTool{name: "direct"}}, + }, + ReturnDirectly: map[string]bool{"direct": true}, + }, + }) + require.NoError(t, err) + + t.Run("Invoke", func(t *testing.T) { + atomic.StoreInt32(&mdl.callCount, 0) + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, EnableStreaming: false, + }) + events := drainAgenticEvents(runner.Query(ctx, "test")) + + // Model should be called only once (for the tool call), not a second + // time, because the tool is return-directly. + assert.Equal(t, int32(1), atomic.LoadInt32(&mdl.callCount)) + + // Find the final output event — should be the return-directly tool result. + last := lastAgenticEvent(events) + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + + msg := last.Output.MessageOutput.Message + require.NotNil(t, msg) + require.GreaterOrEqual(t, len(msg.ContentBlocks), 1) + ftr := msg.ContentBlocks[0].FunctionToolResult + require.NotNil(t, ftr, "expected FunctionToolResult in final output, got type=%v", msg.ContentBlocks[0].Type) + assert.Equal(t, "call-1", ftr.CallID) + }) + + t.Run("Stream", func(t *testing.T) { + atomic.StoreInt32(&mdl.callCount, 0) + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, EnableStreaming: true, + }) + events := drainAgenticEvents(runner.Query(ctx, "test")) + + assert.Equal(t, int32(1), atomic.LoadInt32(&mdl.callCount)) + + last := lastAgenticEvent(events) + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + + mo := last.Output.MessageOutput + if mo.IsStreaming { + var finalMsg *schema.AgenticMessage + for { + chunk, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + break + } + finalMsg = chunk + } + require.NotNil(t, finalMsg) + require.GreaterOrEqual(t, len(finalMsg.ContentBlocks), 1) + ftr := finalMsg.ContentBlocks[0].FunctionToolResult + require.NotNil(t, ftr) + assert.Equal(t, "call-1", ftr.CallID) + } else { + msg := mo.Message + require.NotNil(t, msg) + require.GreaterOrEqual(t, len(msg.ContentBlocks), 1) + ftr := msg.ContentBlocks[0].FunctionToolResult + require.NotNil(t, ftr) + assert.Equal(t, "call-1", ftr.CallID) + } + }) +} + +func TestAgenticReact_CancelAfterChatModel(t *testing.T) { + ctx := context.Background() + + toolStarted := make(chan struct{}, 1) + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + switch count { + case 1: + return agenticToolCallMsg("slow", "c1", `"hi"`), nil + case 2: + return agenticToolCallMsg("slow", "c2", `"hi2"`), nil + default: + return agenticMsg("should not reach"), nil + } + }, + } + + slowTool := &agenticSignalTool{ + name: "slow", + started: toolStarted, + result: "slow result", + } + + agent := newAgenticAgent(t, ctx, mdl, []tool.BaseTool{slowTool}) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("trigger cancel")}, + }, cancelOpt) + + <-toolStarted + + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + _ = handle.Wait() + }() + + time.Sleep(10 * time.Millisecond) + slowTool.finish() + + var capturedErr error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + capturedErr = ev.Err + } + } + require.Error(t, capturedErr, "expected CancelError event") + var cancelErr *CancelError + require.ErrorAs(t, capturedErr, &cancelErr) +} + +func TestAgenticReact_CancelAfterToolCalls(t *testing.T) { + ctx := context.Background() + + toolStarted := make(chan struct{}, 1) + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + if count == 1 { + return agenticToolCallMsg("slow", "c1", `"hi"`), nil + } + return agenticMsg("should not reach on second call"), nil + }, + } + + slowTool := &agenticSignalTool{ + name: "slow", + started: toolStarted, + result: "slow result", + } + + agent := newAgenticAgent(t, ctx, mdl, []tool.BaseTool{slowTool}) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("trigger cancel")}, + }, cancelOpt) + + <-toolStarted + + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + _ = handle.Wait() + }() + + time.Sleep(10 * time.Millisecond) + slowTool.finish() + + var capturedErr error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + capturedErr = ev.Err + } + } + require.Error(t, capturedErr, "expected CancelError event") + var cancelErr *CancelError + require.ErrorAs(t, capturedErr, &cancelErr) + assert.Equal(t, int32(1), atomic.LoadInt32(&modelCallCount)) +} + +func TestAgenticReact_DoubleInterruptResume(t *testing.T) { + ctx := context.Background() + + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + switch count { + case 1: + return agenticToolCallMsg("approval_tool", "c1", `"first"`), nil + case 2: + return agenticToolCallMsg("approval_tool", "c2", `"second"`), nil + case 3: + return agenticMsg("all approved"), nil + default: + return nil, fmt.Errorf("unexpected call #%d", count) + } + }, + } + + store := &agenticReactTestStore{m: map[string][]byte{}} + runner := newAgenticRunnerWithStore(t, ctx, mdl, []tool.BaseTool{&agenticInterruptTool{name: "approval_tool"}}, store) + + events1 := drainAgenticEvents(runner.Query(ctx, "approve twice", WithCheckPointID("dbl-cp"))) + require.NoError(t, firstAgenticEventError(events1)) + int1Event := findInterruptEvent(events1) + require.NotNil(t, int1Event, "expected first interrupt") + int1ID := int1Event.Action.Interrupted.InterruptContexts[0].ID + + iter2, err := runner.ResumeWithParams(ctx, "dbl-cp", &ResumeParams{ + Targets: map[string]any{int1ID: "approved_1"}, + }) + require.NoError(t, err) + + events2 := drainAgenticEvents(iter2) + require.NoError(t, firstAgenticEventError(events2)) + int2Event := findInterruptEvent(events2) + require.NotNil(t, int2Event, "expected second interrupt") + int2ID := int2Event.Action.Interrupted.InterruptContexts[0].ID + + iter3, err := runner.ResumeWithParams(ctx, "dbl-cp", &ResumeParams{ + Targets: map[string]any{int2ID: "approved_2"}, + }) + require.NoError(t, err) + + events3 := drainAgenticEvents(iter3) + require.NoError(t, firstAgenticEventError(events3)) + last := lastAgenticEvent(events3) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Contains(t, agenticTextContent(last.Output.MessageOutput.Message), "all approved") +} + +func TestAgenticReact_ResumeThenCancelAfterChatModelCheckpoint(t *testing.T) { + ctx := context.Background() + + var modelCallCount int32 + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + count := atomic.AddInt32(&modelCallCount, 1) + switch count { + case 1: + return agenticToolCallMsg("approval_tool", "c1", `"first"`), nil + case 2: + return agenticToolCallMsg("approval_tool", "c2", `"second"`), nil + default: + return nil, fmt.Errorf("unexpected call #%d", count) + } + }, + } + + store := &agenticReactTestStore{m: map[string][]byte{}} + runner := newAgenticRunnerWithStore(t, ctx, mdl, []tool.BaseTool{&agenticInterruptTool{name: "approval_tool"}}, store) + + events1 := drainAgenticEvents(runner.Query(ctx, "approve then cancel", WithCheckPointID("resume-cancel-cp"))) + require.NoError(t, firstAgenticEventError(events1)) + int1Event := findInterruptEvent(events1) + require.NotNil(t, int1Event, "expected tool interrupt") + int1ID := int1Event.Action.Interrupted.InterruptContexts[0].ID + + cancelOpt, cancelFn := WithCancel() + handle, contributed := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + require.True(t, contributed) + + iter2, err := runner.ResumeWithParams(ctx, "resume-cancel-cp", &ResumeParams{ + Targets: map[string]any{int1ID: "approved"}, + }, cancelOpt) + require.NoError(t, err) + + events2 := drainAgenticEvents(iter2) + require.NoError(t, handle.Wait()) + for _, event := range events2 { + if event.Err == nil { + continue + } + var cancelErr *CancelError + require.ErrorAs(t, event.Err, &cancelErr) + } +} + +func TestAgenticReact_ChatModelAgent_NoTools(t *testing.T) { + ctx := context.Background() + + mdl := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("no tools response"), nil + }, + } + + runner := newAgenticRunner(t, ctx, mdl, nil) + events := drainAgenticEvents(runner.Query(ctx, "hello")) + last := lastAgenticEvent(events) + + require.NotNil(t, last) + require.Nil(t, last.Err) + require.NotNil(t, last.Output) + require.NotNil(t, last.Output.MessageOutput) + assert.Equal(t, "no tools response", agenticTextContent(last.Output.MessageOutput.Message)) +} + +func TestAgenticReact_ChatModelAgent_ToolsReceiveArgs(t *testing.T) { + ctx := context.Background() + + var receivedArgs string + captureTool := &agenticArgCaptureTool{ + name: "capture", + onInvoke: func(args string) string { + receivedArgs = args + return "captured" + }, + } + + mdl := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCallMsg("capture", "c1", `{"foo":"bar"}`), + agenticMsg("done"), + }, + } + + runner := newAgenticRunner(t, ctx, mdl, []tool.BaseTool{captureTool}) + drainAgenticEvents(runner.Query(ctx, "call capture")) + + assert.Equal(t, `{"foo":"bar"}`, receivedArgs) +} + +func TestCoverage_AgenticReact_Streaming(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + streamFn: func(_ context.Context, input []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { + defer w.Close() + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "streamed response"}), + }, + }, nil) + }() + return r, nil + }, + } + + echoTool := &agenticEchoTool{name: "echo"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "stream-react", + Description: "streaming agentic react", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{echoTool}, + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + iter := runner.Query(ctx, "stream me") + + var events []*TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + stream := event.Output.MessageOutput.MessageStream + for { + _, sErr := stream.Recv() + if sErr != nil { + break + } + } + } + events = append(events, event) + } + + require.NotEmpty(t, events) + assertAgenticEventRoleFields(t, events) +} + +func TestCoverage_ConcatMessageStream_Agentic(t *testing.T) { + t.Run("Success", func(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}), + }, + }, nil) + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + }, nil) + }() + + result, err := concatMessageStream(r) + assert.NoError(t, err) + assert.NotNil(t, result) + }) + + t.Run("ErrorDuringRecv", func(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + w.Send(nil, fmt.Errorf("recv error")) + w.Close() + }() + + _, err := concatMessageStream(r) + assert.Error(t, err) + }) +} + +func TestCoverage_AgenticReact_InterruptResume(t *testing.T) { + ctx := context.Background() + + interruptTool := &agenticInterruptTool{name: "approval"} + + var callIdx int32 + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&callIdx, 1) + if idx == 1 { + return agenticToolCallMsg("approval", "call1", `{}`), nil + } + return agenticMsg("approved and done"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "interrupt-agent", + Description: "tests interrupt and resume", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{interruptTool}, + }, + }, + }) + require.NoError(t, err) + + store := newDTTestStore() + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("need approval"), + }, WithCheckPointID("cp-int")) + + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + } + + require.NotNil(t, interruptEvent, "should have interrupt event") + + var rootCauseID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + rootCauseID = intCtx.ID + break + } + } + require.NotEmpty(t, rootCauseID) + + resumeIter, err := runner.ResumeWithParams(ctx, "cp-int", &ResumeParams{ + Targets: map[string]any{rootCauseID: "approved"}, + }) + require.NoError(t, err) + + var events []*TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := resumeIter.Next() + if !ok { + break + } + events = append(events, event) + } + require.NotEmpty(t, events) +} + +func TestCoverage_AgenticMessageHasToolCalls(t *testing.T) { + t.Run("NilMessage", func(t *testing.T) { + assert.False(t, agenticMessageHasToolCalls(nil)) + }) + + t.Run("NoToolCalls", func(t *testing.T) { + msg := agenticMsg("just text") + assert.False(t, agenticMessageHasToolCalls(msg)) + }) + + t.Run("HasToolCalls", func(t *testing.T) { + msg := agenticToolCallMsg("tool1", "id1", `{}`) + assert.True(t, agenticMessageHasToolCalls(msg)) + }) + + t.Run("NilBlock", func(t *testing.T) { + msg := &schema.AgenticMessage{ + ContentBlocks: []*schema.ContentBlock{nil}, + } + assert.False(t, agenticMessageHasToolCalls(msg)) + }) + + t.Run("ToolCallBlockNilFunctionToolCall", func(t *testing.T) { + msg := &schema.AgenticMessage{ + ContentBlocks: []*schema.ContentBlock{ + {Type: schema.ContentBlockTypeFunctionToolCall, FunctionToolCall: nil}, + }, + } + assert.False(t, agenticMessageHasToolCalls(msg)) + }) +} + +func TestCoverage_ChatModelAgent_StreamError(t *testing.T) { + ctx := context.Background() + + testErr := errors.New("stream failed") + m := &mockAgenticModel{ + streamFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + return nil, testErr + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "stream-error-agent", + Description: "tests stream error", + Model: m, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + + iter := runner.Query(ctx, "trigger stream error") + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should propagate stream error") +} + +func TestCoverage_AgenticReact_GobStateRoundTrip(t *testing.T) { + ctx := context.Background() + + var callIdx int32 + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&callIdx, 1) + if idx == 1 { + return agenticToolCallMsg("interrupt_tool", "call1", `{}`), nil + } + return agenticMsg("completed"), nil + }, + } + + interruptTool := &agenticInterruptTool{name: "interrupt_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "gob-test", + Description: "tests gob state round trip", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{interruptTool}, + }, + }, + }) + require.NoError(t, err) + + store := newDTTestStore() + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("test gob"), + }, WithCheckPointID("gob-cp")) + + var interrupted bool + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interrupted = true + interruptEvent = event + } + } + + if !interrupted || interruptEvent == nil { + t.Skip("no interrupt occurred, skipping gob round-trip test") + } + + _, exists, err := store.Get(ctx, "gob-cp") + assert.NoError(t, err) + assert.True(t, exists, "checkpoint should be saved") + + var rootCauseID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + rootCauseID = intCtx.ID + break + } + } + require.NotEmpty(t, rootCauseID) + + resumeIter, err := runner.ResumeWithParams(ctx, "gob-cp", &ResumeParams{ + Targets: map[string]any{rootCauseID: "approved"}, + }) + require.NoError(t, err) + + var resumed bool + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + resumed = true + } + } + assert.True(t, resumed, "should successfully resume from gob checkpoint") +} + +func TestCoverage_GetMessageFromTypedWrappedEvent_Agentic(t *testing.T) { + t.Run("NilOutput", func(t *testing.T) { + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + event: &TypedAgentEvent[*schema.AgenticMessage]{}, + } + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Nil(t, msg) + }) + + t.Run("NonStreaming", func(t *testing.T) { + expected := agenticMsg("hello") + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: expected, + }, + }, + }, + } + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("StreamingAlreadyConcatenated", func(t *testing.T) { + expected := agenticMsg("already concatenated") + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + concatenatedMessage: expected, + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + }, + }, + }, + } + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("StreamingWithPriorError", func(t *testing.T) { + testErr := errors.New("prior stream error") + wrapper := &typedAgentEventWrapper[*schema.AgenticMessage]{ + event: &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + }, + }, + }, + } + wrapper.StreamErr = testErr + msg, err := getMessageFromTypedWrappedEvent(wrapper) + assert.Equal(t, testErr, err) + assert.Nil(t, msg) + }) +} + +func TestCoverage_GetMessageFromWrappedEvent_ErrorPaths(t *testing.T) { + t.Run("NilOutput", func(t *testing.T) { + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{}, + } + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Nil(t, msg) + }) + + t.Run("NonStreaming", func(t *testing.T) { + expected := schema.AssistantMessage("hello", nil) + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: expected, + }, + }, + }, + } + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("AlreadyConcatenated", func(t *testing.T) { + expected := schema.AssistantMessage("concatenated", nil) + wrapper := &agentEventWrapper{ + concatenatedMessage: expected, + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + }, + }, + }, + } + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, expected, msg) + }) + + t.Run("PriorStreamError", func(t *testing.T) { + testErr := errors.New("prior error") + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + }, + }, + }, + } + wrapper.StreamErr = testErr + msg, err := getMessageFromWrappedEvent(wrapper) + assert.Equal(t, testErr, err) + assert.Nil(t, msg) + }) +} + +func TestCoverage_ConsumeStream_ErrorDuringRecv(t *testing.T) { + testErr := errors.New("stream recv error") + r, w := schema.Pipe[*schema.Message](2) + go func() { + w.Send(schema.AssistantMessage("partial", nil), nil) + w.Send(nil, testErr) + w.Close() + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + assert.NotNil(t, wrapper.StreamErr) + assert.Nil(t, wrapper.concatenatedMessage) +} + +func TestCoverage_ConsumeStream_EmptyStream(t *testing.T) { + r, w := schema.Pipe[*schema.Message](1) + go func() { w.Close() }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + require.NotNil(t, wrapper.StreamErr) + assert.Contains(t, wrapper.StreamErr.Error(), "no messages") +} + +func TestCoverage_ConsumeStream_MultipleMessages(t *testing.T) { + r, w := schema.Pipe[*schema.Message](3) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("chunk1", nil), nil) + w.Send(schema.AssistantMessage("chunk2", nil), nil) + w.Send(schema.AssistantMessage("chunk3", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + assert.Nil(t, wrapper.StreamErr) + assert.NotNil(t, wrapper.concatenatedMessage) +} + +func TestCoverage_ConsumeStream_SingleMessage(t *testing.T) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("single", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + + assert.Nil(t, wrapper.StreamErr) + require.NotNil(t, wrapper.concatenatedMessage) + assert.Equal(t, "single", wrapper.concatenatedMessage.Content) +} + +func TestCoverage_ConsumeStream_Idempotent(t *testing.T) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("once", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + }, + } + + wrapper.consumeStream() + msg1 := wrapper.concatenatedMessage + + wrapper.consumeStream() + msg2 := wrapper.concatenatedMessage + + assert.Equal(t, msg1, msg2, "second call should be no-op") +} diff --git a/adk/agentic_test.go b/adk/agentic_test.go new file mode 100644 index 0000000..5a4ae88 --- /dev/null +++ b/adk/agentic_test.go @@ -0,0 +1,1745 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "io" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type mockAgenticModel struct { + generateFn func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) + streamFn func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) +} + +func (m *mockAgenticModel) Generate(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return m.generateFn(ctx, input, opts...) +} + +func (m *mockAgenticModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + if m.streamFn != nil { + return m.streamFn(ctx, input, opts...) + } + result, err := m.generateFn(ctx, input, opts...) + if err != nil { + return nil, err + } + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { defer w.Close(); w.Send(result, nil) }() + return r, nil +} + +type testAgenticMiddleware struct { + *TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + beforeFn func(context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) + afterFn func(context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) +} + +func (m *testAgenticMiddleware) BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + if m.beforeFn != nil { + return m.beforeFn(ctx, state, mc) + } + return ctx, state, nil +} + +func (m *testAgenticMiddleware) AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + if m.afterFn != nil { + return m.afterFn(ctx, state, mc) + } + return ctx, state, nil +} + +func TestAgenticChatModelAgentRun_NoTools(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello from agentic model"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticTestAgent", + Description: "Agentic test agent", + Instruction: "You are helpful.", + Model: m, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }, + } + iter := agent.Run(ctx, input) + require.NotNil(t, iter) + + event, ok := iter.Next() + assert.True(t, ok) + require.NotNil(t, event) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + require.NotNil(t, msg) + assert.Equal(t, schema.AgenticRoleTypeAssistant, msg.Role) + assert.Len(t, msg.ContentBlocks, 1) + assert.Equal(t, "Hello from agentic model", msg.ContentBlocks[0].AssistantGenText.Text) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestAgenticChatModelAgentRun_WithTools(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Used tool and got result"}), + }, + } + + var receivedToolInfos []*schema.ToolInfo + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + o := model.GetCommonOptions(&model.Options{}, opts...) + receivedToolInfos = o.Tools + return agenticResponse, nil + }, + } + + dummyTool := newSlowTool("dummy_tool", 0, "ok") + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticToolAgent", + Description: "Agentic agent with tools", + Instruction: "You are helpful.", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{dummyTool}, + }, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Call a tool"), + }, + } + iter := agent.Run(ctx, input) + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + + _, ok = iter.Next() + assert.False(t, ok) + + require.Len(t, receivedToolInfos, 1) + assert.Equal(t, "dummy_tool", receivedToolInfos[0].Name) +} + +func TestAgenticChatModelAgentRun_Streaming(t *testing.T) { + ctx := context.Background() + + chunk1 := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}), + }, + } + chunk2 := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + } + + m := &mockAgenticModel{ + streamFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(chunk1, nil) + w.Send(chunk2, nil) + }() + return r, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticStreamAgent", + Description: "Agentic streaming agent", + Instruction: "You are helpful.", + Model: m, + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }, + EnableStreaming: true, + } + iter := agent.Run(ctx, input) + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.NotNil(t, event.Output.MessageOutput.MessageStream) + event.Output.MessageOutput.MessageStream.Close() + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestDefaultAgenticGenModelInput(t *testing.T) { + ctx := context.Background() + + t.Run("WithInstruction", func(t *testing.T) { + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + msgs, err := newDefaultGenModelInput[*schema.AgenticMessage]()(ctx, "Be helpful", input) + assert.NoError(t, err) + assert.Len(t, msgs, 2) + assert.Equal(t, schema.AgenticRoleTypeSystem, msgs[0].Role) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[1].Role) + }) + + t.Run("WithoutInstruction", func(t *testing.T) { + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + msgs, err := newDefaultGenModelInput[*schema.AgenticMessage]()(ctx, "", input) + assert.NoError(t, err) + assert.Len(t, msgs, 1) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[0].Role) + }) +} + +func TestAgenticRunnerQuery(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "query response"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "QueryAgent", + Description: "Query test agent", + Instruction: "Be helpful.", + Model: m, + }) + assert.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + + iter := runner.Query(ctx, "What's up?") + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func agenticAssistantMessage(text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: text}), + }, + } +} + +type mockAgenticRunnerAgent struct { + name string + description string + responses []*TypedAgentEvent[*schema.AgenticMessage] + callCount int + lastInput *TypedAgentInput[*schema.AgenticMessage] + enableStreaming bool +} + +func (a *mockAgenticRunnerAgent) Name(_ context.Context) string { return a.name } +func (a *mockAgenticRunnerAgent) Description(_ context.Context) string { return a.description } +func (a *mockAgenticRunnerAgent) Run(_ context.Context, input *TypedAgentInput[*schema.AgenticMessage], _ ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + a.callCount++ + a.lastInput = input + a.enableStreaming = input.EnableStreaming + + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + for _, event := range a.responses { + generator.Send(event) + if event.Action != nil && event.Action.Exit { + break + } + } + }() + return iterator +} + +type mockAgenticAgent struct { + name string + description string + responses []*TypedAgentEvent[*schema.AgenticMessage] +} + +func (a *mockAgenticAgent) Name(_ context.Context) string { return a.name } +func (a *mockAgenticAgent) Description(_ context.Context) string { return a.description } +func (a *mockAgenticAgent) Run(_ context.Context, _ *TypedAgentInput[*schema.AgenticMessage], _ ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer generator.Close() + for _, event := range a.responses { + generator.Send(event) + if event.Action != nil && event.Action.Exit { + break + } + } + }() + return iterator +} + +type myAgenticAgent struct { + name string + runFn func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] + resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] +} + +func (m *myAgenticAgent) Name(_ context.Context) string { + if len(m.name) > 0 { + return m.name + } + return "myAgenticAgent" +} +func (m *myAgenticAgent) Description(_ context.Context) string { return "my agentic agent description" } +func (m *myAgenticAgent) Run(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + return m.runFn(ctx, input, options...) +} +func (m *myAgenticAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + return m.resumeFn(ctx, info, opts...) +} + +func TestAgenticChatModelAgentRun_WithMiddleware(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello from agentic agent"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + afterModelExecuted := false + + mw := &testAgenticMiddleware{ + beforeFn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + state.Messages = append(state.Messages, schema.UserAgenticMessage("extra")) + return ctx, state, nil + }, + afterFn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + assert.Len(t, state.Messages, 4) + afterModelExecuted = true + return ctx, state, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticMiddlewareAgent", + Description: "Agentic agent with middleware", + Instruction: "You are helpful.", + Model: m, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{mw}, + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hi"), + }, + } + iter := agent.Run(ctx, input) + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.NotNil(t, event.Output.MessageOutput.Message) + assert.Equal(t, schema.AgenticRoleTypeAssistant, event.Output.MessageOutput.Message.Role) + _, ok = iter.Next() + assert.False(t, ok) + assert.True(t, afterModelExecuted) +} + +func TestAgenticAfterModel_NoTools_ModifyDoesNotAffectEvent(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticAssistantMessage("original content"), nil + }, + } + + var capturedMessages []*schema.AgenticMessage + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticAfterModelAgent", + Description: "Test AfterModelRewriteState", + Instruction: "You are helpful.", + Model: m, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{ + &testAgenticMiddleware{ + afterFn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage], mc *TypedModelContext[*schema.AgenticMessage]) (context.Context, *TypedChatModelAgentState[*schema.AgenticMessage], error) { + capturedMessages = make([]*schema.AgenticMessage, len(state.Messages)) + copy(capturedMessages, state.Messages) + state.Messages = append(state.Messages, agenticAssistantMessage("appended content")) + return ctx, state, nil + }, + }, + }, + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + require.NotNil(t, msg) + assert.Equal(t, "original content", msg.ContentBlocks[0].AssistantGenText.Text) + + _, ok = iterator.Next() + assert.False(t, ok) + + assert.Len(t, capturedMessages, 3) +} + +func TestAgenticGetComposeOptions_WithChatModelOptions(t *testing.T) { + ctx := context.Background() + + var capturedTemperature float32 + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + options := model.GetCommonOptions(&model.Options{}, opts...) + if options.Temperature != nil { + capturedTemperature = *options.Temperature + } + return agenticAssistantMessage("response"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticOptionsAgent", + Description: "Test agent", + Model: m, + }) + assert.NoError(t, err) + + temp := float32(0.7) + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}}, + WithChatModelOptions([]model.Option{model.WithTemperature(temp)})) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, temp, capturedTemperature) +} + +func TestAgenticChatModelAgent_PrepareExecContextError(t *testing.T) { + ctx := context.Background() + + expectedErr := errors.New("tool info error") + errTool := &errorTool{infoErr: expectedErr} + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticAssistantMessage("response"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticErrToolAgent", + Description: "Test agent", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{errTool}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}}) + + event, ok := iter.Next() + assert.True(t, ok) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "tool info error") + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestAgenticChatModelAgentOutputKey(t *testing.T) { + t.Run("OutputKeyStoresInSession", func(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticAssistantMessage("Hello from agentic assistant."), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticOutputKeyAgent", + Description: "Test agent for output key", + Instruction: "You are helpful.", + Model: m, + OutputKey: "agent_output", + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + } + ctx, runCtx := initTypedRunCtx(ctx, "AgenticOutputKeyAgent", input) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + msg := event.Output.MessageOutput.Message + assert.Equal(t, "Hello from agentic assistant.", msg.ContentBlocks[0].AssistantGenText.Text) + + _, ok = iterator.Next() + assert.False(t, ok) + + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "agent_output") + assert.Equal(t, "Hello from agentic assistant.", sessionValues["agent_output"]) + }) + + t.Run("OutputKeyWithStreamingStoresInSession", func(t *testing.T) { + ctx := context.Background() + + chunk1 := agenticAssistantMessage("Hello") + chunk2 := agenticAssistantMessage(", world.") + + m := &mockAgenticModel{ + streamFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(chunk1, nil) + w.Send(chunk2, nil) + }() + return r, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticStreamOutputKeyAgent", + Description: "Test agent for streaming output key", + Instruction: "You are helpful.", + Model: m, + OutputKey: "agent_output", + }) + assert.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }, + EnableStreaming: true, + } + ctx, runCtx := initTypedRunCtx(ctx, "AgenticStreamOutputKeyAgent", input) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) + }) + + t.Run("SetOutputToSessionAgenticMessage", func(t *testing.T) { + ctx := context.Background() + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}, + } + ctx, runCtx := initTypedRunCtx(ctx, "TestAgent", input) + require.NotNil(t, runCtx) + require.NotNil(t, runCtx.Session) + + msg := agenticAssistantMessage("Test response") + err := setOutputToSession(ctx, msg, nil, "test_output") + assert.NoError(t, err) + + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "test_output") + assert.Equal(t, "Test response", sessionValues["test_output"]) + }) +} + +func TestAgenticRunner_Run_WithStreaming(t *testing.T) { + ctx := context.Background() + + mockAgent_ := &mockAgenticRunnerAgent{ + name: "AgenticStreamRunnerAgent", + description: "Test agent for agentic runner streaming", + responses: []*TypedAgentEvent[*schema.AgenticMessage]{ + { + AgentName: "AgenticStreamRunnerAgent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]*schema.AgenticMessage{ + agenticAssistantMessage("Streaming response"), + }), + }, + }, + }, + }, + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{EnableStreaming: true, Agent: mockAgent_}) + + msgs := []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello, agent!"), + } + + iterator := runner.Run(ctx, msgs) + + assert.Equal(t, 1, mockAgent_.callCount) + assert.Equal(t, msgs, mockAgent_.lastInput.Messages) + assert.True(t, mockAgent_.enableStreaming) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "AgenticStreamRunnerAgent", event.AgentName) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestAgenticRunner_Query_WithStreaming(t *testing.T) { + ctx := context.Background() + + mockAgent_ := &mockAgenticRunnerAgent{ + name: "AgenticStreamQueryAgent", + description: "Test agent for agentic runner query streaming", + responses: []*TypedAgentEvent[*schema.AgenticMessage]{ + { + AgentName: "AgenticStreamQueryAgent", + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]*schema.AgenticMessage{ + agenticAssistantMessage("Streaming query response"), + }), + }, + }, + }, + }, + } + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{EnableStreaming: true, Agent: mockAgent_}) + + iterator := runner.Query(ctx, "Test query") + + assert.Equal(t, 1, mockAgent_.callCount) + assert.Len(t, mockAgent_.lastInput.Messages, 1) + assert.True(t, mockAgent_.enableStreaming) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "AgenticStreamQueryAgent", event.AgentName) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestAgenticSimpleInterrupt(t *testing.T) { + data := "hello world" + agent := &myAgenticAgent{ + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: schema.StreamReaderFromArray([]*schema.AgenticMessage{ + schema.UserAgenticMessage("hello "), + schema.UserAgenticMessage("world"), + }), + }, + }, + }) + intEvent := TypedInterrupt[*schema.AgenticMessage](ctx, data) + intEvent.Action.Interrupted.Data = data + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + assert.True(t, info.WasInterrupted) + assert.Nil(t, info.InterruptState) + assert.True(t, info.EnableStreaming) + assert.Equal(t, data, info.Data) + + assert.True(t, info.IsResumeTarget) + iter, generator := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + generator.Close() + return iter + }, + } + store := newMyStore() + ctx := context.Background() + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + CheckPointStore: store, + }) + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + + var interruptEvent *TypedAgentEvent[*schema.AgenticMessage] + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + } + + require.NotNil(t, interruptEvent) + assert.Equal(t, data, interruptEvent.Action.Interrupted.Data) + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID) + assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause) + assert.Equal(t, data, interruptEvent.Action.Interrupted.InterruptContexts[0].Info) + assert.Equal(t, Address{{Type: AddressSegmentAgent, ID: "myAgenticAgent"}}, + interruptEvent.Action.Interrupted.InterruptContexts[0].Address) +} + +func TestCascadingFrom_NewChatModelAgentFrom(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "from response"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "FromAgent", + Description: "Test cascading constructor", + Instruction: "Be helpful.", + Model: m, + }) + assert.NoError(t, err) + assert.Equal(t, "FromAgent", agent.Name(ctx)) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("Hello"), + }) + + event, ok := iter.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestCascadingTyped_TypedStatefulInterrupt(t *testing.T) { + ctx := context.Background() + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, "test-agent") + + type myState struct { + Count int + } + + event := TypedStatefulInterrupt[*schema.AgenticMessage](ctx, "please confirm", &myState{Count: 42}) + require.NotNil(t, event) + require.NotNil(t, event.Action) + require.NotNil(t, event.Action.Interrupted) +} + +func TestCascadingTyped_EventFromAgenticMessage(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "hello"}), + }, + } + + event := EventFromAgenticMessage(msg, nil, schema.AgenticRoleTypeAssistant) + require.NotNil(t, event) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + assert.Equal(t, msg, event.Output.MessageOutput.Message) + assert.False(t, event.Output.MessageOutput.IsStreaming) + assert.Equal(t, schema.RoleType(""), event.Output.MessageOutput.Role) + assert.Equal(t, schema.AgenticRoleTypeAssistant, event.Output.MessageOutput.AgenticRole) + assert.Empty(t, event.Output.MessageOutput.ToolName) +} + +// assertAgenticEventRoleFields asserts that all AgenticMessage events in the +// list have zero-valued Role and ToolName fields (which are *schema.Message-only), +// and that AgenticRole is populated with a non-zero value. +func assertAgenticEventRoleFields(t *testing.T, events []*TypedAgentEvent[*schema.AgenticMessage]) { + t.Helper() + for i, event := range events { + if event.Output == nil || event.Output.MessageOutput == nil { + continue + } + mo := event.Output.MessageOutput + assert.Equal(t, schema.RoleType(""), mo.Role, "event[%d]: AgenticMessage must have zero Role", i) + assert.Empty(t, mo.ToolName, "event[%d]: AgenticMessage must have empty ToolName", i) + assert.NotEmpty(t, mo.AgenticRole, "event[%d]: AgenticMessage must have non-zero AgenticRole", i) + } +} + +func TestCoverage_FlowAgent_ResumeNotResumable(t *testing.T) { + ctx := context.Background() + + agent := &mockAgenticAgent{ + name: "non-resumable", + description: "cannot resume", + responses: []*TypedAgentEvent[*schema.AgenticMessage]{ + {Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("done"), + }, + }}, + }, + } + + fa := toTypedFlowAgent[*schema.AgenticMessage](agent) + + info := &ResumeInfo{WasInterrupted: true} + iter := fa.Resume(ctx, info) + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should get error for non-resumable agent") +} + +func TestCoverage_GenAgenticErrorIter(t *testing.T) { + testErr := errors.New("test agentic error") + iter := genAgenticErrorIter(testErr) + + event, ok := iter.Next() + require.True(t, ok) + assert.Equal(t, testErr, event.Err) + + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestCoverage_ChatModelAgent_OnSetSubAgents_FrozenError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "freeze-test", + Description: "frozen test agent", + Model: m, + }) + require.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + } + iter := agent.Run(ctx, input) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + err = agent.OnSetSubAgents(ctx, []TypedAgent[*schema.AgenticMessage]{ + &mockAgenticAgent{name: "late-child"}, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "frozen") +} + +func TestCoverage_ChatModelAgent_OnSetAsSubAgent_FrozenError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "freeze-child", + Description: "frozen child agent", + Model: m, + }) + require.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + } + iter := agent.Run(ctx, input) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + err = agent.OnSetAsSubAgent(ctx, &mockAgenticAgent{name: "parent"}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "frozen") +} + +func TestCoverage_ChatModelAgent_OnSetAsSubAgent_DuplicateError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "dup-child", + Description: "duplicate child agent", + Model: m, + }) + require.NoError(t, err) + + err = agent.OnSetAsSubAgent(ctx, &mockAgenticAgent{name: "parent1"}) + assert.NoError(t, err) + + err = agent.OnSetAsSubAgent(ctx, &mockAgenticAgent{name: "parent2"}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "already been set as a sub-agent") +} + +func TestCoverage_ChatModelAgent_OnDisallowTransferToParent_FrozenError(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("done"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "disallow-test", + Description: "disallow transfer test", + Model: m, + }) + require.NoError(t, err) + + input := &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("Hi")}, + } + iter := agent.Run(ctx, input) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + err = agent.OnDisallowTransferToParent(ctx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "frozen") +} + +func TestCoverage_TypedGetMessage_AgenticNonStreaming(t *testing.T) { + msg := agenticMsg("hello") + event := &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: msg, + }, + }, + } + + result, retEvent, err := TypedGetMessage(event) + assert.NoError(t, err) + assert.Equal(t, msg, result) + assert.Equal(t, event, retEvent) +} + +func TestCoverage_TypedGetMessage_AgenticStreaming(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hello "}), + }, + }, nil) + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "world"}), + }, + }, nil) + }() + + event := &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: r, + }, + }, + } + + result, retEvent, err := TypedGetMessage(event) + assert.NoError(t, err) + assert.NotNil(t, result) + require.NotNil(t, retEvent) + assert.NotNil(t, retEvent.Output.MessageOutput.MessageStream) +} + +func TestCoverage_TypedGetMessage_NilOutput(t *testing.T) { + event := &TypedAgentEvent[*schema.AgenticMessage]{} + + result, retEvent, err := TypedGetMessage(event) + assert.NoError(t, err) + assert.Nil(t, result) + assert.Equal(t, event, retEvent) +} + +func TestCoverage_GetMessage_NonStreaming(t *testing.T) { + msg := schema.AssistantMessage("hello", nil) + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: msg, + }, + }, + } + + result, retEvent, err := GetMessage(event) + assert.NoError(t, err) + assert.Equal(t, msg, result) + assert.Equal(t, event, retEvent) +} + +func TestCoverage_GetMessage_Streaming(t *testing.T) { + r, w := schema.Pipe[*schema.Message](2) + go func() { + defer w.Close() + w.Send(schema.AssistantMessage("Hello ", nil), nil) + w.Send(schema.AssistantMessage("world", nil), nil) + }() + + event := &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + }, + }, + } + + result, retEvent, err := GetMessage(event) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.NotNil(t, retEvent) +} + +func TestCoverage_NewTypedAgentTool_Agentic(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("tool response"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "tool-agent", + Description: "agent wrapped as tool", + Model: m, + }) + require.NoError(t, err) + + agentTool := NewTypedAgentTool[*schema.AgenticMessage](ctx, agent) + + info, err := agentTool.Info(ctx) + require.NoError(t, err) + assert.Equal(t, "tool-agent", info.Name) + + result, err := agentTool.(tool.InvokableTool).InvokableRun(ctx, `{"request":"test"}`) + require.NoError(t, err) + assert.Contains(t, result, "tool response") +} +func TestCoverage_CopyAgenticEvent(t *testing.T) { + original := &TypedAgentEvent[*schema.AgenticMessage]{ + AgentName: "agent1", + RunPath: []RunStep{{agentName: "root"}, {agentName: "agent1"}}, + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("hello"), + AgenticRole: schema.AgenticRoleTypeAssistant, + }, + }, + Action: &AgentAction{ + TransferToAgent: &TransferToAgentAction{DestAgentName: "agent2"}, + }, + } + + copied := copyTypedAgentEvent(original) + assert.Equal(t, original.AgentName, copied.AgentName) + assert.Equal(t, len(original.RunPath), len(copied.RunPath)) + assert.Equal(t, original.Action, copied.Action) + assert.Equal(t, schema.AgenticRoleTypeAssistant, copied.Output.MessageOutput.AgenticRole) + + copied.RunPath[0].agentName = "mutated" + assert.NotEqual(t, original.RunPath[0].agentName, copied.RunPath[0].agentName) +} + +func TestCoverage_ChatModelAgent_ModelGenerateError(t *testing.T) { + ctx := context.Background() + + testErr := errors.New("model generate failed") + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return nil, testErr + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "error-model-agent", + Description: "tests model generate error", + Model: m, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + }) + + iter := runner.Query(ctx, "trigger error") + + var capturedErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + capturedErr = event.Err + } + } + require.Error(t, capturedErr, "should propagate model error") +} + +func TestCoverage_NewTypedUserMessages(t *testing.T) { + t.Run("Message", func(t *testing.T) { + msgs := newTypedUserMessages[*schema.Message]("hello") + require.Len(t, msgs, 1) + assert.Equal(t, schema.User, msgs[0].Role) + assert.Equal(t, "hello", msgs[0].Content) + }) + + t.Run("AgenticMessage", func(t *testing.T) { + msgs := newTypedUserMessages[*schema.AgenticMessage]("hello") + require.Len(t, msgs, 1) + assert.Equal(t, schema.AgenticRoleTypeUser, msgs[0].Role) + }) +} + +func TestCoverage_TypedEndpointModel_NilEndpoints(t *testing.T) { + ctx := context.Background() + + m := &typedEndpointModel[*schema.AgenticMessage]{} + + _, err := m.Generate(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "generate endpoint not set") + + _, err = m.Stream(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "stream endpoint not set") +} + +func TestCoverage_TypedEndpointModel_WithEndpoints(t *testing.T) { + ctx := context.Background() + + expected := agenticMsg("generated") + m := &typedEndpointModel[*schema.AgenticMessage]{ + generate: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return expected, nil + }, + stream: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { + defer w.Close() + w.Send(expected, nil) + }() + return r, nil + }, + } + + result, err := m.Generate(ctx, nil) + assert.NoError(t, err) + assert.Equal(t, expected, result) + + stream, err := m.Stream(ctx, nil) + assert.NoError(t, err) + require.NotNil(t, stream) + msg, err := stream.Recv() + assert.NoError(t, err) + assert.Equal(t, expected, msg) + _, err = stream.Recv() + assert.Equal(t, io.EOF, err) +} + +func TestCoverage_SetAutomaticClose(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { + defer w.Close() + w.Send(agenticMsg("data"), nil) + }() + + event := &TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + IsStreaming: true, + MessageStream: r, + }, + }, + } + + typedSetAutomaticClose(event) +} + +func TestConcatMessageStream_AgenticClosesStream(t *testing.T) { + r, w := schema.Pipe[*schema.AgenticMessage](2) + go func() { + defer w.Close() + w.Send(agenticMsg("a"), nil) + w.Send(agenticMsg("b"), nil) + }() + + result, err := concatMessageStream(r) + require.NoError(t, err) + require.NotNil(t, result) + + _, recvErr := r.Recv() + assert.Error(t, recvErr, + "stream should be closed after concatMessageStream returns") +} + +// --- Agentic retry/failover stream test helpers --- + +func agenticStreamWithMidError(chunks []*schema.AgenticMessage, err error) *schema.StreamReader[*schema.AgenticMessage] { + sr, sw := schema.Pipe[*schema.AgenticMessage](len(chunks) + 1) + go func() { + defer sw.Close() + for _, c := range chunks { + sw.Send(c, nil) + } + sw.Send(nil, err) + }() + return sr +} + +func agenticStreamOK(chunks []*schema.AgenticMessage) *schema.StreamReader[*schema.AgenticMessage] { + sr, sw := schema.Pipe[*schema.AgenticMessage](len(chunks)) + go func() { + defer sw.Close() + for _, c := range chunks { + sw.Send(c, nil) + } + }() + return sr +} + +func drainTypedAgenticEvents(t *testing.T, iter *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) *schema.AgenticMessage { + t.Helper() + var lastMsg *schema.AgenticMessage + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + var willRetry *WillRetryError + if errors.As(ev.Err, &willRetry) { + continue + } + t.Fatalf("unexpected error event: %v", ev.Err) + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + if ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + sr := ev.Output.MessageOutput.MessageStream + for { + chunk, err := sr.Recv() + if err != nil { + break + } + lastMsg = chunk + } + } else if ev.Output.MessageOutput.Message != nil { + lastMsg = ev.Output.MessageOutput.Message + } + } + } + return lastMsg +} + +func TestAgenticRetryWithShouldRetry_Generate(t *testing.T) { + ctx := context.Background() + + var callCount int32 + var shouldRetryCalls int32 + genErr := errors.New("transient generate error") + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + n := atomic.AddInt32(&callCount, 1) + if n == 1 { + return nil, genErr + } + return agenticMsg("retry ok"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "retry-gen-agent", + Description: "test retry generate", + Model: m, + ModelRetryConfig: &TypedModelRetryConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *TypedRetryContext[*schema.AgenticMessage]) *TypedRetryDecision[*schema.AgenticMessage] { + n := atomic.AddInt32(&shouldRetryCalls, 1) + if n == 1 { + assert.Nil(t, retryCtx.OutputMessage, "OutputMessage should be nil when Generate returns error") + assert.ErrorIs(t, retryCtx.Err, genErr, "Err should be the generate error") + assert.Equal(t, 1, retryCtx.RetryAttempt) + return &TypedRetryDecision[*schema.AgenticMessage]{Retry: true} + } + return &TypedRetryDecision[*schema.AgenticMessage]{Retry: false} + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return time.Millisecond }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("hello")}) + + msg := drainTypedAgenticEvents(t, iter) + require.NotNil(t, msg, "should have received a final message") + assert.Equal(t, "retry ok", agenticTextContent(msg)) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount), "model should be called twice") + assert.Equal(t, int32(2), atomic.LoadInt32(&shouldRetryCalls), "ShouldRetry should be called for both attempts") +} + +func TestAgenticRetryWithShouldRetry_Stream(t *testing.T) { + ctx := context.Background() + + var streamCallCount int32 + var shouldRetryCalls int32 + streamErr := errors.New("mid-stream error") + + m := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return nil, errors.New("generate should not be called") + }, + streamFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + n := atomic.AddInt32(&streamCallCount, 1) + if n == 1 { + return agenticStreamWithMidError( + []*schema.AgenticMessage{agenticMsg("partial")}, + streamErr, + ), nil + } + return agenticStreamOK([]*schema.AgenticMessage{agenticMsg("stream ok")}), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "retry-stream-agent", + Description: "test retry stream", + Model: m, + ModelRetryConfig: &TypedModelRetryConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *TypedRetryContext[*schema.AgenticMessage]) *TypedRetryDecision[*schema.AgenticMessage] { + n := atomic.AddInt32(&shouldRetryCalls, 1) + if n == 1 { + assert.NotNil(t, retryCtx.OutputMessage, "OutputMessage should be non-nil from partial stream") + assert.Error(t, retryCtx.Err, "Err should be the stream error") + return &TypedRetryDecision[*schema.AgenticMessage]{Retry: true} + } + return nil + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return time.Millisecond }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("hello")}) + + var lastMsg *schema.AgenticMessage + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + var willRetry *WillRetryError + if errors.As(ev.Err, &willRetry) { + continue + } + t.Fatalf("unexpected error: %v", ev.Err) + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + if ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + sr := ev.Output.MessageOutput.MessageStream + for { + chunk, err := sr.Recv() + if err != nil { + break + } + lastMsg = chunk + } + } else if ev.Output.MessageOutput.Message != nil { + lastMsg = ev.Output.MessageOutput.Message + } + } + } + require.NotNil(t, lastMsg, "should have received final stream message") + assert.Contains(t, agenticTextContent(lastMsg), "stream ok") + assert.Equal(t, int32(2), atomic.LoadInt32(&shouldRetryCalls), "ShouldRetry should be called for both attempts") +} + +func TestAgenticFailoverGenerate(t *testing.T) { + ctx := context.Background() + + m1Err := errors.New("m1 generate failed") + var m1Calls, m2Calls int32 + + m1 := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, m1Err + }, + } + m2 := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + atomic.AddInt32(&m2Calls, 1) + return agenticMsg("failover ok"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "failover-gen-agent", + Description: "test failover generate", + Model: m1, + ModelFailoverConfig: &ModelFailoverConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.AgenticMessage, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { + assert.Equal(t, uint(1), failoverCtx.FailoverAttempt) + assert.Nil(t, failoverCtx.LastOutputMessage, "LastOutputMessage should be nil when Generate returns error") + assert.ErrorIs(t, failoverCtx.LastErr, m1Err) + return m2, nil, nil + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("hello")}) + + msg := drainTypedAgenticEvents(t, iter) + require.NotNil(t, msg) + assert.Equal(t, "failover ok", agenticTextContent(msg)) + assert.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + assert.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) +} + +func TestAgenticFailoverStream_MidStreamError(t *testing.T) { + ctx := context.Background() + + streamErr := errors.New("m1 mid-stream error") + var m1Calls, m2Calls int32 + var capturedLastOutput *schema.AgenticMessage + + m1 := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return nil, errors.New("unused") + }, + streamFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + atomic.AddInt32(&m1Calls, 1) + return agenticStreamWithMidError( + []*schema.AgenticMessage{agenticMsg("partial chunk")}, + streamErr, + ), nil + }, + } + m2 := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + return nil, errors.New("unused") + }, + streamFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + atomic.AddInt32(&m2Calls, 1) + return agenticStreamOK([]*schema.AgenticMessage{agenticMsg("failover stream ok")}), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "failover-stream-agent", + Description: "test failover stream", + Model: m1, + ModelFailoverConfig: &ModelFailoverConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.AgenticMessage, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { + capturedLastOutput = failoverCtx.LastOutputMessage + return m2, nil, nil + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + EnableStreaming: true, + }) + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("hello")}) + + var lastMsg *schema.AgenticMessage + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + var willRetry *WillRetryError + if errors.As(ev.Err, &willRetry) { + continue + } + t.Fatalf("unexpected error: %v", ev.Err) + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + if ev.Output.MessageOutput.IsStreaming && ev.Output.MessageOutput.MessageStream != nil { + sr := ev.Output.MessageOutput.MessageStream + for { + chunk, err := sr.Recv() + if err != nil { + break + } + lastMsg = chunk + } + } else if ev.Output.MessageOutput.Message != nil { + lastMsg = ev.Output.MessageOutput.Message + } + } + } + + require.NotNil(t, lastMsg, "should have received final stream from m2") + assert.Contains(t, agenticTextContent(lastMsg), "failover stream ok") + assert.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + assert.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + assert.NotNil(t, capturedLastOutput, "failoverCtx.LastOutputMessage should contain partial stream from m1") +} + +// TestAgenticFailoverGenerate_WithTools verifies that ModelFailoverConfig is +// honored on the ReAct (with-tools) path for *schema.AgenticMessage. This +// guards against the regression where buildAgenticReActRunFunc dropped the +// failover config, leaving ModelFailoverConfig as a no-op for typed agents +// that have any tools configured. +func TestAgenticFailoverGenerate_WithTools(t *testing.T) { + ctx := context.Background() + + m1Err := errors.New("m1 generate failed") + var m1Calls, m2Calls, getFailoverCalls int32 + + m1 := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, m1Err + }, + } + m2 := &mockAgenticModel{ + generateFn: func(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + atomic.AddInt32(&m2Calls, 1) + return agenticMsg("failover ok with tools"), nil + }, + } + + dummyTool := newSlowTool("dummy_tool", 0, "ok") + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "failover-react-agent", + Description: "test failover on ReAct path", + Model: m1, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{dummyTool}, + }, + }, + ModelFailoverConfig: &ModelFailoverConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.AgenticMessage, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { + atomic.AddInt32(&getFailoverCalls, 1) + assert.Equal(t, uint(1), failoverCtx.FailoverAttempt) + assert.ErrorIs(t, failoverCtx.LastErr, m1Err) + return m2, nil, nil + }, + }, + }) + require.NoError(t, err) + + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent}) + iter := runner.Run(ctx, []*schema.AgenticMessage{schema.UserAgenticMessage("hello")}) + + msg := drainTypedAgenticEvents(t, iter) + require.NotNil(t, msg) + assert.Equal(t, "failover ok with tools", agenticTextContent(msg)) + assert.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + assert.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + assert.Equal(t, int32(1), atomic.LoadInt32(&getFailoverCalls), + "GetFailoverModel must be invoked on the ReAct path; if zero, the failover config was dropped by buildAgenticReActRunFunc") +} diff --git a/adk/call_option.go b/adk/call_option.go new file mode 100644 index 0000000..7a1cc1b --- /dev/null +++ b/adk/call_option.go @@ -0,0 +1,211 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import "github.com/cloudwego/eino/callbacks" + +type options struct { + sharedParentSession bool + sessionValues map[string]any + checkPointID *string + skipTransferMessages bool + handlers []callbacks.Handler + cancelCtx *cancelContext +} + +// AgentRunOption is the call option for adk Agent. +type AgentRunOption struct { + implSpecificOptFn any + + // specify which Agent can see this AgentRunOption, if empty, all Agents can see this AgentRunOption + agentNames []string +} + +func (o AgentRunOption) DesignateAgent(name ...string) AgentRunOption { + o.agentNames = append(o.agentNames, name...) + return o +} + +func getCommonOptions(base *options, opts ...AgentRunOption) *options { + if base == nil { + base = &options{} + } + + return GetImplSpecificOptions(base, opts...) +} + +// WithSessionValues sets session-scoped values for the agent run. +func WithSessionValues(v map[string]any) AgentRunOption { + return WrapImplSpecificOptFn(func(o *options) { + o.sessionValues = v + }) +} + +// WithSkipTransferMessages disables forwarding transfer messages during execution. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func WithSkipTransferMessages() AgentRunOption { + return WrapImplSpecificOptFn(func(t *options) { + t.skipTransferMessages = true + }) +} + +func withSharedParentSession() AgentRunOption { + return WrapImplSpecificOptFn(func(o *options) { + o.sharedParentSession = true + }) +} + +// WithCallbacks adds callback handlers to receive agent lifecycle events. +// Handlers receive OnStart with AgentCallbackInput and OnEnd with AgentCallbackOutput. +// Multiple handlers can be added; each receives an independent copy of the event stream. +func WithCallbacks(handlers ...callbacks.Handler) AgentRunOption { + return WrapImplSpecificOptFn(func(o *options) { + o.handlers = append(o.handlers, handlers...) + }) +} + +// WrapImplSpecificOptFn is the option to wrap the implementation specific option function. +func WrapImplSpecificOptFn[T any](optFn func(*T)) AgentRunOption { + return AgentRunOption{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions extract the implementation specific options from AgentRunOption list, optionally providing a base options with default values. +// e.g. +// +// myOption := &MyOption{ +// Field1: "default_value", +// } +// +// myOption := model.GetImplSpecificOptions(myOption, opts...) +func GetImplSpecificOptions[T any](base *T, opts ...AgentRunOption) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} + +// filterCallbackHandlersForNestedAgents removes callback handlers that have already been applied +// to the current agent before passing opts to nested inner agents. +// +// This is necessary for workflow agents (LoopAgent, SequentialAgent, ParallelAgent) because: +// 1. Callback handlers designated for the current agent are applied via initAgentCallbacks(), +// which stores them in the context. +// 2. Nested inner agents inherit this context, so they automatically receive these callbacks. +// 3. If we also pass these handlers in opts to inner agents, they would be applied twice, +// causing duplicate callback invocations. +// +// Note: This only applies to workflow agents where inner agents inherit context from the parent. +// For flowAgent's sub-agents (which are peer agents that transfer to each other), the full opts +// are passed since they don't inherit the parent's callback context. +func filterCallbackHandlersForNestedAgents(currentAgentName string, opts []AgentRunOption) []AgentRunOption { + if len(opts) == 0 { + return nil + } + var filteredOpts []AgentRunOption + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn == nil { + filteredOpts = append(filteredOpts, opt) + continue + } + if _, isCallbackOpt := opt.implSpecificOptFn.(func(*options)); isCallbackOpt { + testOpt := &options{} + opt.implSpecificOptFn.(func(*options))(testOpt) + if len(testOpt.handlers) > 0 { + if len(opt.agentNames) == 0 { + continue + } + matched := false + for _, name := range opt.agentNames { + if name == currentAgentName { + matched = true + break + } + } + if matched { + continue + } + } + } + filteredOpts = append(filteredOpts, opt) + } + return filteredOpts +} + +// filterCancelOption removes any AgentRunOption that sets a cancelCtx on *options. +// This prevents inner (nested) agents from receiving the cancel option when the +// outer flowAgent owns the cancel lifecycle. Inner agents access the cancelContext +// via the Go context (getCancelContext) instead. +func filterCancelOption(opts []AgentRunOption) []AgentRunOption { + if len(opts) == 0 { + return nil + } + var filteredOpts []AgentRunOption + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn == nil { + filteredOpts = append(filteredOpts, opt) + continue + } + if _, isCommonOpt := opt.implSpecificOptFn.(func(*options)); isCommonOpt { + testOpt := &options{} + opt.implSpecificOptFn.(func(*options))(testOpt) + if testOpt.cancelCtx != nil { + continue + } + } + filteredOpts = append(filteredOpts, opt) + } + return filteredOpts +} + +func filterOptions(agentName string, opts []AgentRunOption) []AgentRunOption { + if len(opts) == 0 { + return nil + } + var filteredOpts []AgentRunOption + for i := range opts { + opt := opts[i] + if len(opt.agentNames) == 0 { + filteredOpts = append(filteredOpts, opt) + continue + } + for j := range opt.agentNames { + if opt.agentNames[j] == agentName { + filteredOpts = append(filteredOpts, opt) + break + } + } + } + return filteredOpts +} diff --git a/adk/callback.go b/adk/callback.go new file mode 100644 index 0000000..3818500 --- /dev/null +++ b/adk/callback.go @@ -0,0 +1,202 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + icb "github.com/cloudwego/eino/internal/callbacks" +) + +// AgentCallbackInput represents the input passed to agent callbacks during OnStart. +// Use ConvAgentCallbackInput to safely convert from callbacks.CallbackInput. +type AgentCallbackInput struct { + // Input contains the agent input for a new run. Nil when resuming. + Input *AgentInput + // ResumeInfo contains resume information when resuming from an interrupt. Nil for new runs. + ResumeInfo *ResumeInfo +} + +// AgentCallbackOutput represents the output passed to agent callbacks during OnEnd. +// Use ConvAgentCallbackOutput to safely convert from callbacks.CallbackOutput. +// +// Important: The Events iterator should be consumed asynchronously to avoid blocking +// the agent execution. Each callback handler receives an independent copy of the iterator. +type AgentCallbackOutput struct { + // Events provides a stream of agent events. Each handler receives its own copy. + Events *AsyncIterator[*AgentEvent] +} + +func copyTypedEventIterator[M MessageType](iter *AsyncIterator[*TypedAgentEvent[M]], n int) []*AsyncIterator[*TypedAgentEvent[M]] { + if n <= 0 { + return nil + } + if n == 1 { + return []*AsyncIterator[*TypedAgentEvent[M]]{iter} + } + + iterators := make([]*AsyncIterator[*TypedAgentEvent[M]], n) + generators := make([]*AsyncGenerator[*TypedAgentEvent[M]], n) + for i := 0; i < n; i++ { + iterators[i], generators[i] = NewAsyncIteratorPair[*TypedAgentEvent[M]]() + } + + go func() { + defer func() { + for _, g := range generators { + g.Close() + } + }() + + for { + event, ok := iter.Next() + if !ok { + break + } + for i := 0; i < n-1; i++ { + generators[i].Send(copyTypedAgentEvent(event)) + } + generators[n-1].Send(event) + } + }() + + return iterators +} + +func copyAgentCallbackOutput(out *AgentCallbackOutput, n int) []*AgentCallbackOutput { + if out == nil || out.Events == nil { + result := make([]*AgentCallbackOutput, n) + for i := 0; i < n; i++ { + result[i] = out + } + return result + } + iters := copyTypedEventIterator(out.Events, n) + result := make([]*AgentCallbackOutput, n) + for i, iter := range iters { + result[i] = &AgentCallbackOutput{Events: iter} + } + return result +} + +// ConvAgentCallbackInput converts a generic CallbackInput to AgentCallbackInput. +// Returns nil if the input is not an AgentCallbackInput. +func ConvAgentCallbackInput(input callbacks.CallbackInput) *AgentCallbackInput { + if v, ok := input.(*AgentCallbackInput); ok { + return v + } + return nil +} + +// ConvAgentCallbackOutput converts a generic CallbackOutput to AgentCallbackOutput. +// Returns nil if the output is not an AgentCallbackOutput. +func ConvAgentCallbackOutput(output callbacks.CallbackOutput) *AgentCallbackOutput { + if v, ok := output.(*AgentCallbackOutput); ok { + return v + } + return nil +} + +func initAgentCallbacks(ctx context.Context, agentName, agentType string, opts ...AgentRunOption) context.Context { + ri := &callbacks.RunInfo{ + Name: agentName, + Type: agentType, + Component: ComponentOfAgent, + } + + o := getCommonOptions(nil, opts...) + if len(o.handlers) == 0 { + return icb.ReuseHandlers(ctx, ri) + } + return icb.AppendHandlers(ctx, ri, o.handlers...) +} + +func getAgentType(agent Agent) string { + if typer, ok := agent.(components.Typer); ok { + return typer.GetType() + } + return "" +} + +// TypedAgentCallbackInput represents the input passed to typed agent callbacks during OnStart. +// Use ConvTypedCallbackInput to safely convert from callbacks.CallbackInput. +type TypedAgentCallbackInput[M MessageType] struct { + // Input contains the agent input for a new run. Nil when resuming. + Input *TypedAgentInput[M] + // ResumeInfo contains resume information when resuming from an interrupt. Nil for new runs. + ResumeInfo *ResumeInfo +} + +// TypedAgentCallbackOutput represents the output passed to typed agent callbacks during OnEnd. +// Use ConvTypedCallbackOutput to safely convert from callbacks.CallbackOutput. +// +// Important: The Events iterator should be consumed asynchronously to avoid blocking +// the agent execution. Each callback handler receives an independent copy of the iterator. +type TypedAgentCallbackOutput[M MessageType] struct { + // Events provides a stream of agent events. Each handler receives its own copy. + Events *AsyncIterator[*TypedAgentEvent[M]] +} + +// ConvTypedCallbackInput converts a callbacks.CallbackInput to *TypedAgentCallbackInput[M]. +// Returns nil if the input is not of the expected type. +func ConvTypedCallbackInput[M MessageType](input callbacks.CallbackInput) *TypedAgentCallbackInput[M] { + if v, ok := input.(*TypedAgentCallbackInput[M]); ok { + return v + } + return nil +} + +// ConvTypedCallbackOutput converts a callbacks.CallbackOutput to *TypedAgentCallbackOutput[M]. +// Returns nil if the output is not of the expected type. +func ConvTypedCallbackOutput[M MessageType](output callbacks.CallbackOutput) *TypedAgentCallbackOutput[M] { + if v, ok := output.(*TypedAgentCallbackOutput[M]); ok { + return v + } + return nil +} + +func copyTypedCallbackOutput[M MessageType](out *TypedAgentCallbackOutput[M], n int) []*TypedAgentCallbackOutput[M] { + if out == nil || out.Events == nil { + result := make([]*TypedAgentCallbackOutput[M], n) + for i := 0; i < n; i++ { + result[i] = out + } + return result + } + iters := copyTypedEventIterator(out.Events, n) + result := make([]*TypedAgentCallbackOutput[M], n) + for i, iter := range iters { + result[i] = &TypedAgentCallbackOutput[M]{Events: iter} + } + return result +} + +func initAgenticCallbacks(ctx context.Context, agentName, agentType string, opts ...AgentRunOption) context.Context { + ri := &callbacks.RunInfo{ + Name: agentName, + Type: agentType, + Component: ComponentOfAgenticAgent, + } + + o := getCommonOptions(nil, opts...) + if len(o.handlers) == 0 { + return icb.ReuseHandlers(ctx, ri) + } + return icb.AppendHandlers(ctx, ri, o.handlers...) +} diff --git a/adk/callback_integration_test.go b/adk/callback_integration_test.go new file mode 100644 index 0000000..fd708a0 --- /dev/null +++ b/adk/callback_integration_test.go @@ -0,0 +1,1062 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/callbacks" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type callbackRecorder struct { + mu sync.Mutex + onStartCalled bool + onEndCalled bool + runInfo *callbacks.RunInfo + inputReceived *AgentCallbackInput + eventsReceived []*AgentEvent + eventsDone chan struct{} + closeOnce sync.Once +} + +func (r *callbackRecorder) getOnStartCalled() bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.onStartCalled +} + +func (r *callbackRecorder) getOnEndCalled() bool { + r.mu.Lock() + defer r.mu.Unlock() + return r.onEndCalled +} + +func (r *callbackRecorder) getEventsReceived() []*AgentEvent { + r.mu.Lock() + defer r.mu.Unlock() + result := make([]*AgentEvent, len(r.eventsReceived)) + copy(result, r.eventsReceived) + return result +} + +func newRecordingHandler(recorder *callbackRecorder) callbacks.Handler { + recorder.eventsDone = make(chan struct{}) + return callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + recorder.mu.Lock() + defer recorder.mu.Unlock() + recorder.onStartCalled = true + recorder.runInfo = info + if agentInput := ConvAgentCallbackInput(input); agentInput != nil { + recorder.inputReceived = agentInput + } + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + recorder.mu.Lock() + recorder.onEndCalled = true + recorder.runInfo = info + recorder.mu.Unlock() + + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil { + if agentOutput.Events != nil { + go func() { + defer recorder.closeOnce.Do(func() { close(recorder.eventsDone) }) + for { + event, ok := agentOutput.Events.Next() + if !ok { + break + } + recorder.mu.Lock() + recorder.eventsReceived = append(recorder.eventsReceived, event) + recorder.mu.Unlock() + } + }() + return ctx + } + } + recorder.closeOnce.Do(func() { close(recorder.eventsDone) }) + return ctx + }). + Build() +} + +func TestCallbackOnStartInvocation(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("test response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for callback", + Instruction: "You are a test agent", + Model: cm, + }) + assert.NoError(t, err) + + recorder := &callbackRecorder{} + handler := newRecordingHandler(recorder) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + <-recorder.eventsDone + + assert.True(t, recorder.onStartCalled, "OnStart should be called") + assert.NotNil(t, recorder.inputReceived, "Input should be received") + assert.NotNil(t, recorder.inputReceived.Input, "AgentInput should be set") + assert.Len(t, recorder.inputReceived.Input.Messages, 1) +} + +func TestCallbackOnEndInvocation(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("test response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for callback", + Instruction: "You are a test agent", + Model: cm, + }) + assert.NoError(t, err) + + recorder := &callbackRecorder{} + handler := newRecordingHandler(recorder) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + <-recorder.eventsDone + + assert.True(t, recorder.onEndCalled, "OnEnd should be called") + assert.NotEmpty(t, recorder.eventsReceived, "Events should be received") +} + +func TestCallbackRunInfoForChatModelAgent(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("test response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestChatAgent", + Description: "Test chat agent", + Instruction: "You are a test agent", + Model: cm, + }) + assert.NoError(t, err) + + recorder := &callbackRecorder{} + handler := newRecordingHandler(recorder) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + <-recorder.eventsDone + + assert.NotNil(t, recorder.runInfo) + assert.Equal(t, "TestChatAgent", recorder.runInfo.Name) + assert.Equal(t, "ChatModel", recorder.runInfo.Type) + assert.Equal(t, ComponentOfAgent, recorder.runInfo.Component) +} + +func TestMultipleCallbackHandlers(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("test response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test agent", + Model: cm, + }) + assert.NoError(t, err) + + recorder1 := &callbackRecorder{} + recorder2 := &callbackRecorder{} + handler1 := newRecordingHandler(recorder1) + handler2 := newRecordingHandler(recorder2) + + runner := NewRunner(ctx, RunnerConfig{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler1, handler2)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + <-recorder1.eventsDone + <-recorder2.eventsDone + + assert.True(t, recorder1.onStartCalled, "Handler1 OnStart should be called") + assert.True(t, recorder2.onStartCalled, "Handler2 OnStart should be called") + assert.True(t, recorder1.onEndCalled, "Handler1 OnEnd should be called") + assert.True(t, recorder2.onEndCalled, "Handler2 OnEnd should be called") + + assert.NotEmpty(t, recorder1.eventsReceived, "Handler1 should receive events") + assert.NotEmpty(t, recorder2.eventsReceived, "Handler2 should receive events") +} + +func TestCallbackWithWorkflowAgent(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response 1", nil), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response 2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "SequentialAgent", + Description: "Sequential workflow", + SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + var callbackInfos []*callbacks.RunInfo + handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component == ComponentOfAgent { + callbackInfos = append(callbackInfos, info) + } + return ctx + }). + Build() + + runner := NewRunner(ctx, RunnerConfig{Agent: seqAgent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.NotEmpty(t, callbackInfos, "OnStart should be called for agents") + foundAgent1 := false + foundAgent2 := false + for _, info := range callbackInfos { + if info.Name == "Agent1" && info.Type == "ChatModel" { + foundAgent1 = true + } + if info.Name == "Agent2" && info.Type == "ChatModel" { + foundAgent2 = true + } + } + assert.True(t, foundAgent1, "Agent1 callback should be invoked") + assert.True(t, foundAgent2, "Agent2 callback should be invoked") +} + +func TestCallbackEventsMatchAgentOutput(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + expectedContent := "This is the test response content" + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage(expectedContent, nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test agent", + Model: cm, + }) + assert.NoError(t, err) + + recorder := &callbackRecorder{} + handler := newRecordingHandler(recorder) + + var agentEvents []*AgentEvent + runner := NewRunner(ctx, RunnerConfig{Agent: agent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + event, ok := iter.Next() + if !ok { + break + } + agentEvents = append(agentEvents, event) + } + + <-recorder.eventsDone + + assert.NotEmpty(t, agentEvents, "Agent should emit events") + assert.NotEmpty(t, recorder.eventsReceived, "Callback should receive events") + + foundExpectedContent := false + for _, event := range recorder.eventsReceived { + if event.Output != nil && event.Output.MessageOutput != nil { + msg := event.Output.MessageOutput.Message + if msg != nil && msg.Content == expectedContent { + foundExpectedContent = true + break + } + } + } + assert.True(t, foundExpectedContent, "Callback events should contain the expected content") +} + +func TestCallbackOnEndForWorkflowAgent(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response 1", nil), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response 2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "SequentialAgent", + Description: "Sequential workflow", + SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + recorder := &callbackRecorder{} + handler := newRecordingHandler(recorder) + + runner := NewRunner(ctx, RunnerConfig{Agent: seqAgent}) + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + <-recorder.eventsDone + + assert.True(t, recorder.getOnStartCalled(), "OnStart should be called for workflow agent") + assert.True(t, recorder.getOnEndCalled(), "OnEnd should be called for workflow agent") + assert.NotEmpty(t, recorder.getEventsReceived(), "Events should be received for workflow agent") +} + +type ctxKeyForTest string + +const testOnStartMarkerKey ctxKeyForTest = "onStartMarker" + +func TestSubAgentContextIsolation(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("transferring to Agent2", + []schema.ToolCall{ + { + ID: "transfer_1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "Agent2"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response from Agent2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent that transfers to Agent2", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents}) + + var mu sync.Mutex + onStartContextMarkers := make(map[string][]string) + + handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + marker, _ := ctx.Value(testOnStartMarkerKey).(string) + onStartContextMarkers[info.Name] = append(onStartContextMarkers[info.Name], marker) + mu.Unlock() + + return context.WithValue(ctx, testOnStartMarkerKey, info.Name+"_marker") + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.NotEmpty(t, onStartContextMarkers["Agent1"], "Agent1's OnStart should be called") + assert.NotEmpty(t, onStartContextMarkers["Agent2"], "Agent2's OnStart should be called") + + if len(onStartContextMarkers["Agent1"]) > 0 { + assert.Equal(t, "", onStartContextMarkers["Agent1"][0], + "Agent1's OnStart should receive context without marker (initial context)") + } + if len(onStartContextMarkers["Agent2"]) > 0 { + assert.Equal(t, "", onStartContextMarkers["Agent2"][0], + "Agent2's first OnStart should NOT inherit Agent1's marker - context should be isolated") + } +} + +func TestCallbackDesignatedToSpecificAgent(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("transferring to Agent2", + []schema.ToolCall{ + { + ID: "transfer_1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "Agent2"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response from Agent2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent that transfers to Agent2", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents}) + + var mu sync.Mutex + onStartCalls := make(map[string]int) + + agent2OnlyHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + onStartCalls[info.Name]++ + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + iter := runner.Query(ctx, "hello", WithCallbacks(agent2OnlyHandler).DesignateAgent("Agent2")) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.Equal(t, 0, onStartCalls["Agent1"], "Agent1's OnStart should NOT be called when handler is designated to Agent2") + assert.Equal(t, 1, onStartCalls["Agent2"], "Agent2's OnStart should be called exactly once") +} + +func TestCallbackDesignatedToMultipleAgents(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("transferring to Agent2", + []schema.ToolCall{ + { + ID: "transfer_1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "Agent2"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response from Agent2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents}) + + var mu sync.Mutex + onStartCalls := make(map[string]int) + + agent1And2Handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + onStartCalls[info.Name]++ + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + iter := runner.Query(ctx, "hello", WithCallbacks(agent1And2Handler).DesignateAgent("Agent1", "Agent2")) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.Equal(t, 1, onStartCalls["Agent1"], "Agent1's OnStart should be called exactly once") + assert.Equal(t, 1, onStartCalls["Agent2"], "Agent2's OnStart should be called exactly once") +} + +func TestCallbackDesignatedExcludesNonMatchingAgents(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("transferring to Agent2", + []schema.ToolCall{ + { + ID: "transfer_1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "Agent2"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response from Agent2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents}) + + var mu sync.Mutex + onStartCalls := make(map[string]int) + + agent1OnlyHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + onStartCalls[info.Name]++ + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + iter := runner.Query(ctx, "hello", WithCallbacks(agent1OnlyHandler).DesignateAgent("Agent1")) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.Equal(t, 1, onStartCalls["Agent1"], "Agent1's OnStart should be called exactly once") + assert.Equal(t, 0, onStartCalls["Agent2"], "Agent2's OnStart should NOT be called when handler is designated only to Agent1") +} + +func TestMixedDesignatedAndGlobalCallbacks(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("transferring to Agent2", + []schema.ToolCall{ + { + ID: "transfer_1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "Agent2"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response from Agent2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent that transfers to Agent2", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents}) + + var mu sync.Mutex + globalHandlerCalls := make(map[string]int) + agent2OnlyHandlerCalls := make(map[string]int) + + globalHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + globalHandlerCalls[info.Name]++ + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + agent2OnlyHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + agent2OnlyHandlerCalls[info.Name]++ + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + iter := runner.Query(ctx, "hello", + WithCallbacks(globalHandler), + WithCallbacks(agent2OnlyHandler).DesignateAgent("Agent2"), + ) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.Equal(t, 1, globalHandlerCalls["Agent1"], "Global handler should fire for Agent1") + assert.Equal(t, 1, globalHandlerCalls["Agent2"], "Global handler should fire for Agent2") + + assert.Equal(t, 0, agent2OnlyHandlerCalls["Agent1"], "Agent2-only handler should NOT fire for Agent1") + assert.Equal(t, 1, agent2OnlyHandlerCalls["Agent2"], "Agent2-only handler should fire for Agent2") +} + +func TestOnStartCalledOncePerAgentWithDesignation(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("transferring to Agent2", + []schema.ToolCall{ + { + ID: "transfer_1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "Agent2"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response from Agent2", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent1", + Description: "First agent that transfers to Agent2", + Instruction: "You are agent 1", + Model: cm1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "Agent2", + Description: "Second agent", + Instruction: "You are agent 2", + Model: cm2, + }) + assert.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, agent1, []Agent{agent2}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: agentWithSubAgents}) + + var mu sync.Mutex + onStartCalls := make(map[string]int) + + handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + mu.Lock() + onStartCalls[info.Name]++ + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != ComponentOfAgent { + return ctx + } + if agentOutput := ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + iter := runner.Query(ctx, "hello", WithCallbacks(handler)) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.Equal(t, 1, onStartCalls["Agent1"], "Agent1's OnStart should be called exactly once") + assert.Equal(t, 1, onStartCalls["Agent2"], "Agent2's OnStart should be called exactly once") +} diff --git a/adk/callback_test.go b/adk/callback_test.go new file mode 100644 index 0000000..efd66f5 --- /dev/null +++ b/adk/callback_test.go @@ -0,0 +1,388 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +func TestCopyTypedEventIterator(t *testing.T) { + t.Run("n=0 returns nil", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + gen.Send(&AgentEvent{AgentName: "test"}) + gen.Close() + }() + + result := copyTypedEventIterator(iter, 0) + assert.Nil(t, result) + }) + + t.Run("n=1 returns original iterator", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + gen.Send(&AgentEvent{AgentName: "test"}) + gen.Close() + }() + + result := copyTypedEventIterator(iter, 1) + assert.Len(t, result, 1) + assert.Equal(t, iter, result[0]) + }) + + t.Run("n>1 creates n independent copies", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + events := []*AgentEvent{ + {AgentName: "agent1", Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage("msg1", nil)}}}, + {AgentName: "agent2", Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.AssistantMessage("msg2", nil)}}}, + } + + go func() { + for _, e := range events { + gen.Send(e) + } + gen.Close() + }() + + n := 3 + copies := copyTypedEventIterator(iter, n) + assert.Len(t, copies, n) + + var wg sync.WaitGroup + receivedEvents := make([][]*AgentEvent, n) + + for i := 0; i < n; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + for { + event, ok := copies[idx].Next() + if !ok { + break + } + receivedEvents[idx] = append(receivedEvents[idx], event) + } + }(i) + } + + wg.Wait() + + for i := 0; i < n; i++ { + assert.Len(t, receivedEvents[i], len(events), "iterator %d should receive all events", i) + for j, e := range receivedEvents[i] { + assert.Equal(t, events[j].AgentName, e.AgentName) + } + } + }) +} + +func TestCopyAgentCallbackOutput(t *testing.T) { + t.Run("nil output", func(t *testing.T) { + result := copyAgentCallbackOutput(nil, 3) + assert.Len(t, result, 3) + for _, r := range result { + assert.Nil(t, r) + } + }) + + t.Run("output with nil Events", func(t *testing.T) { + out := &AgentCallbackOutput{Events: nil} + result := copyAgentCallbackOutput(out, 3) + assert.Len(t, result, 3) + for _, r := range result { + assert.Equal(t, out, r) + } + }) + + t.Run("valid output with events", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + gen.Send(&AgentEvent{AgentName: "test"}) + gen.Close() + }() + + out := &AgentCallbackOutput{Events: iter} + result := copyAgentCallbackOutput(out, 2) + assert.Len(t, result, 2) + + for i, r := range result { + require.NotNil(t, r, "result[%d] should not be nil", i) + assert.NotNil(t, r.Events, "result[%d].Events should not be nil", i) + } + }) +} + +func TestConvAgentCallbackInput(t *testing.T) { + t.Run("valid AgentCallbackInput", func(t *testing.T) { + input := &AgentCallbackInput{ + Input: &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + } + result := ConvAgentCallbackInput(input) + assert.Equal(t, input, result) + }) + + t.Run("invalid type returns nil", func(t *testing.T) { + result := ConvAgentCallbackInput("invalid") + assert.Nil(t, result) + }) + + t.Run("nil returns nil", func(t *testing.T) { + result := ConvAgentCallbackInput(nil) + assert.Nil(t, result) + }) +} + +func TestConvAgentCallbackOutput(t *testing.T) { + t.Run("valid AgentCallbackOutput", func(t *testing.T) { + iter, _ := NewAsyncIteratorPair[*AgentEvent]() + output := &AgentCallbackOutput{Events: iter} + result := ConvAgentCallbackOutput(output) + assert.Equal(t, output, result) + }) + + t.Run("invalid type returns nil", func(t *testing.T) { + result := ConvAgentCallbackOutput("invalid") + assert.Nil(t, result) + }) + + t.Run("nil returns nil", func(t *testing.T) { + result := ConvAgentCallbackOutput(nil) + assert.Nil(t, result) + }) +} + +type mockTyperAgent struct { + name string + agentType string +} + +func (a *mockTyperAgent) Name(_ context.Context) string { return a.name } +func (a *mockTyperAgent) Description(_ context.Context) string { return "mock agent" } +func (a *mockTyperAgent) GetType() string { return a.agentType } +func (a *mockTyperAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Close() + return iter +} + +type mockNonTyperAgent struct { + name string +} + +func (a *mockNonTyperAgent) Name(_ context.Context) string { return a.name } +func (a *mockNonTyperAgent) Description(_ context.Context) string { return "mock agent" } +func (a *mockNonTyperAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Close() + return iter +} + +func TestGetAgentType(t *testing.T) { + t.Run("agent implementing Typer", func(t *testing.T) { + agent := &mockTyperAgent{name: "test", agentType: "CustomType"} + result := getAgentType(agent) + assert.Equal(t, "CustomType", result) + }) + + t.Run("agent not implementing Typer", func(t *testing.T) { + agent := &mockNonTyperAgent{name: "test"} + result := getAgentType(agent) + assert.Equal(t, "", result) + }) +} + +func TestWithCallbacksOption(t *testing.T) { + handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + return ctx + }). + Build() + + opt := WithCallbacks(handler) + opts := getCommonOptions(nil, opt) + + assert.Len(t, opts.handlers, 1) +} + +func TestWithMultipleCallbacksOption(t *testing.T) { + handler1 := callbacks.NewHandlerBuilder().Build() + handler2 := callbacks.NewHandlerBuilder().Build() + opt := WithCallbacks(handler1, handler2) + + opts := getCommonOptions(nil, opt) + + assert.Len(t, opts.handlers, 2) +} + +func TestCopyTypedEventIteratorAgentic(t *testing.T) { + t.Run("n=0 returns nil", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"}) + gen.Close() + }() + + result := copyTypedEventIterator(iter, 0) + assert.Nil(t, result) + }) + + t.Run("n=1 returns original iterator", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"}) + gen.Close() + }() + + result := copyTypedEventIterator(iter, 1) + assert.Len(t, result, 1) + assert.Equal(t, iter, result[0]) + }) + + t.Run("n>1 creates n independent copies", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + events := []*TypedAgentEvent[*schema.AgenticMessage]{ + {AgentName: "agent1", Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("msg1")}, + }}, + {AgentName: "agent2", Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{Message: agenticMsg("msg2")}, + }}, + } + + go func() { + for _, e := range events { + gen.Send(e) + } + gen.Close() + }() + + n := 3 + copies := copyTypedEventIterator(iter, n) + assert.Len(t, copies, n) + + var wg sync.WaitGroup + receivedEvents := make([][]*TypedAgentEvent[*schema.AgenticMessage], n) + + for i := 0; i < n; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + for { + event, ok := copies[idx].Next() + if !ok { + break + } + receivedEvents[idx] = append(receivedEvents[idx], event) + } + }(i) + } + + wg.Wait() + + for i := 0; i < n; i++ { + assert.Len(t, receivedEvents[i], len(events), "iterator %d should receive all events", i) + for j, e := range receivedEvents[i] { + assert.Equal(t, events[j].AgentName, e.AgentName) + } + } + }) +} + +func TestCopyTypedCallbackOutput(t *testing.T) { + t.Run("nil output", func(t *testing.T) { + result := copyTypedCallbackOutput[*schema.AgenticMessage](nil, 3) + assert.Len(t, result, 3) + for _, r := range result { + assert.Nil(t, r) + } + }) + + t.Run("output with nil Events", func(t *testing.T) { + out := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: nil} + result := copyTypedCallbackOutput(out, 3) + assert.Len(t, result, 3) + for _, r := range result { + assert.Equal(t, out, r) + } + }) + + t.Run("valid output with events", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{AgentName: "test"}) + gen.Close() + }() + + out := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: iter} + result := copyTypedCallbackOutput(out, 2) + assert.Len(t, result, 2) + + for i, r := range result { + require.NotNil(t, r, "result[%d] should not be nil", i) + assert.NotNil(t, r.Events, "result[%d].Events should not be nil", i) + } + }) +} + +func TestConvTypedCallbackInput(t *testing.T) { + t.Run("valid TypedAgentCallbackInput", func(t *testing.T) { + input := &TypedAgentCallbackInput[*schema.AgenticMessage]{ + Input: &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}, + }, + } + result := ConvTypedCallbackInput[*schema.AgenticMessage](input) + assert.Equal(t, input, result) + }) + + t.Run("invalid type returns nil", func(t *testing.T) { + result := ConvTypedCallbackInput[*schema.AgenticMessage]("invalid") + assert.Nil(t, result) + }) + + t.Run("nil returns nil", func(t *testing.T) { + result := ConvTypedCallbackInput[*schema.AgenticMessage](nil) + assert.Nil(t, result) + }) +} + +func TestConvTypedCallbackOutput(t *testing.T) { + t.Run("valid TypedAgentCallbackOutput", func(t *testing.T) { + iter, _ := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + output := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: iter} + result := ConvTypedCallbackOutput[*schema.AgenticMessage](output) + assert.Equal(t, output, result) + }) + + t.Run("invalid type returns nil", func(t *testing.T) { + result := ConvTypedCallbackOutput[*schema.AgenticMessage]("invalid") + assert.Nil(t, result) + }) + + t.Run("nil returns nil", func(t *testing.T) { + result := ConvTypedCallbackOutput[*schema.AgenticMessage](nil) + assert.Nil(t, result) + }) +} diff --git a/adk/cancel.go b/adk/cancel.go new file mode 100644 index 0000000..6639144 --- /dev/null +++ b/adk/cancel.go @@ -0,0 +1,975 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "io" + "sync" + "sync/atomic" + "time" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*CancelError]("_eino_adk_cancel_error") + schema.RegisterName[*AgentCancelInfo]("_eino_adk_agent_cancel_info") + schema.RegisterName[*StreamCanceledError]("_eino_adk_stream_cancelled_error") +} + +// CancelMode specifies when an agent should be canceled. +// Modes can be combined with bitwise OR to cancel at multiple safe-points. +// For example, CancelAfterChatModel | CancelAfterToolCalls cancels the agent +// after whichever safe-point is reached first. +type CancelMode int + +const ( + // CancelImmediate cancels the agent as soon as the signal is received, + // without waiting for a ChatModel or ToolCalls safe-point. + // By default, only the root agent is interrupted; descendant agents inside + // AgentTools are torn down via context cancellation as a side effect. + // Use WithRecursive to propagate explicit immediate-cancel signals to + // descendants for clean teardown with grace period. + CancelImmediate CancelMode = 0 + // CancelAfterChatModel cancels after the root agent's next chat model call + // completes. By default, only the root agent checks this safe-point; + // nested sub-agents inside AgentTools are unaware of the cancel. + // Use WithRecursive to propagate the cancel to all descendants — whichever + // ChatModel finishes first triggers the cancel. + CancelAfterChatModel CancelMode = 1 << iota + // CancelAfterToolCalls cancels after the root agent's next set of tool calls + // completes. By default, only the root agent checks this safe-point. + // Use WithRecursive to propagate to all descendants. + CancelAfterToolCalls +) + +// CancelHandle represents a cancel operation that can be waited on. +type CancelHandle struct { + wait func() error +} + +// Wait blocks until the cancel request reaches a terminal outcome. +// +// It reports the result of the cancel operation itself, not the agent's final +// business error: +// - nil: cancellation succeeded, including the case where a business interrupt +// was absorbed into CancelError while cancellation was active +// - ErrCancelTimeout: the requested safe-point cancellation timed out and was +// escalated to immediate cancellation +// - ErrExecutionEnded: the execution ended before cancellation took effect, +// meaning the stream drained to completion without any interrupt +func (h *CancelHandle) Wait() error { + return h.wait() +} + +// AgentCancelFunc is called to request cancellation of a running agent. +// It returns after the cancel request is committed; use the returned handle's +// Wait to block for completion and outcome. +// +// The returned bool reports whether this call contributed to the CancelError +// for the current execution. "Contributed" means this call's cancel options +// were included before cancellation was finalized. It is false when cancellation +// was already finalized (handled or execution completed). +type AgentCancelFunc func(...AgentCancelOption) (*CancelHandle, bool) + +type agentCancelConfig struct { + Mode CancelMode + Recursive bool + Timeout *time.Duration +} + +// AgentCancelOption configures cancel behavior. +type AgentCancelOption func(*agentCancelConfig) + +// WithAgentCancelMode sets the cancel mode for the agent cancel operation. +func WithAgentCancelMode(mode CancelMode) AgentCancelOption { + return func(config *agentCancelConfig) { + config.Mode = mode + } +} + +// WithAgentCancelTimeout sets a timeout for the cancel operation. +// This only applies to safe-point modes (CancelAfterChatModel, CancelAfterToolCalls): +// if the safe-point hasn't fired within this duration, the cancel escalates to +// CancelImmediate. The escalated cancel still saves a checkpoint, so the execution +// can be resumed via Runner.Resume or Runner.ResumeWithParams. +// For CancelImmediate this timeout is ignored — the cancel fires immediately. +func WithAgentCancelTimeout(timeout time.Duration) AgentCancelOption { + return func(config *agentCancelConfig) { + config.Timeout = &timeout + } +} + +// WithRecursive opts into recursive cancel propagation. By default, cancel +// modes only affect the root agent; descendant agents inside AgentTools are +// not notified. WithRecursive makes the cancel propagate to all descendants: +// - CancelAfterChatModel / CancelAfterToolCalls: descendants check their own safe-points. +// - CancelImmediate: descendants receive explicit immediate-cancel signals for +// clean teardown; the root uses a grace period to collect child interrupts. +// +// With recursive cancellation, each descendant agent also triggers cancellation +// and cascades its interrupt information upward. The root agent ultimately +// produces a complete checkpoint that includes descendant checkpoints, enabling +// resumption from the exact point where each descendant was interrupted. +// +// Once any cancel call includes WithRecursive, the flag stays set for the +// entire cancel lifecycle (monotonic escalation). +func WithRecursive() AgentCancelOption { + return func(config *agentCancelConfig) { + config.Recursive = true + } +} + +// AgentCancelInfo contains information about a cancel operation. +type AgentCancelInfo struct { + Mode CancelMode + Escalated bool + Timeout bool +} + +// CancelError is sent via AgentEvent.Err when an agent is canceled. +// Use errors.As to match and extract *CancelError from event errors. +// +// Interrupt absorption: when a cancel is active (shouldCancel() == true), ANY +// interrupt — whether from a cancel safe-point node or from business logic +// (e.g. tool.Interrupt in a tool) — is converted to a CancelError. The +// cancel "absorbs" the business interrupt. This is intentional: +// +// - In concurrent execution (parallel workflows, concurrent tool calls), +// cancel-induced and business interrupts can arrive as a single composite +// signal that cannot be split apart. +// - Even in sequential execution, treating business interrupts as CancelError +// during active cancel gives consistent semantics. +// - The business interrupt is NOT lost — the checkpoint preserves the full +// interrupt hierarchy. On resume (Runner.Resume or Runner.ResumeWithParams), +// the agent re-executes the interrupting code path and the business +// interrupt re-fires naturally. +type CancelError struct { + Info *AgentCancelInfo + + // InterruptContexts provides the interrupt contexts needed for targeted + // resumption via Runner.ResumeWithParams. Each context represents a step + // in the agent hierarchy that was interrupted. This is a slice because + // composite agents (e.g. parallel workflows) may interrupt at multiple + // points simultaneously, matching the shape of AgentAction.Interrupted.InterruptContexts. + // Use each InterruptCtx.ID as a key in ResumeParams.Targets. + InterruptContexts []*InterruptCtx + + interruptSignal *InterruptSignal // unexported — only Runner needs it for checkpoint +} + +func (e *CancelError) Error() string { + return fmt.Sprintf("agent canceled: mode=%v, escalated=%v", e.Info.Mode, e.Info.Escalated) +} + +// Sentinel errors for cancel outcomes. +var ( + // ErrCancelTimeout is returned by CancelHandle.Wait when the cancel operation timed out. + ErrCancelTimeout = errors.New("cancel timed out") + + // ErrExecutionEnded is returned by CancelHandle.Wait when the agent ended + // before the cancel took effect. "Ended" means the event stream was fully + // drained without any interrupt — normal completion or a fatal error. + // + // Note: business interrupts that occur while cancel is active are absorbed + // into CancelError (see CancelError doc), so they result in nil (cancel + // succeeded), NOT ErrExecutionEnded. Only execution that completes with + // no interrupt at all produces this error. + ErrExecutionEnded = errors.New("execution already ended") + + // ErrStreamCanceled is the error sent through the stream when CancelImmediate aborts it. + // It is a *StreamCanceledError so it can be gob-serialized during checkpoint save + // (when stored as agentEventWrapper.StreamErr). + ErrStreamCanceled error = &StreamCanceledError{} +) + +// StreamCanceledError is the concrete error type for ErrStreamCanceled. +// It is exported so that gob can serialize it during checkpoint save when the error +// is stored in agentEventWrapper.StreamErr. +type StreamCanceledError struct{} + +func (e *StreamCanceledError) Error() string { + return "stream canceled" +} + +// WithCancel creates an AgentRunOption that enables cancellation for an agent run. +// It returns the option to pass to Run/Resume and a cancel function. +// Cancel options (mode, timeout) are passed to the returned AgentCancelFunc at call time. +func WithCancel() (AgentRunOption, AgentCancelFunc) { + cc := newCancelContext() + opt := WrapImplSpecificOptFn(func(o *options) { + o.cancelCtx = cc + }) + cancelFn := cc.buildCancelFunc() + return opt, cancelFn +} + +// cancelContext state constants (for int32 CAS). +// +// State transition rules: +// +// stateRunning -> stateCancelling (cancel requested by AgentCancelFunc) +// stateRunning -> stateDone (execution finished without interrupt) +// stateCancelling -> stateCancelHandled (ANY interrupt absorbed as CancelError) +// stateCancelling -> stateDone (execution finished without interrupt while cancel pending) +// +// Terminal states: stateDone, stateCancelHandled. +// +// Note: We intentionally do NOT distinguish between "completed" and "errored" +// terminal states. End-users get the actual outcome from AgentEvent. +// This simplification keeps the state machine minimal — only the cancel/non-cancel +// distinction matters for the AgentCancelFunc return value. +// +// Business interrupt handling: when cancel is active (stateCancelling) and any +// interrupt arrives — cancel-induced OR business — wrapIterWithCancelCtx absorbs +// it as a CancelError and transitions to stateCancelHandled. The business interrupt +// data is preserved in the checkpoint for re-emission on resume. +const ( + // stateRunning is the initial state: agent is executing normally. + stateRunning int32 = 0 + // stateCancelling means AgentCancelFunc has been called and cancelChan is + // closed, but the cancel has not yet been handled by the runFunc. + stateCancelling int32 = 1 + // stateDone means execution has finished through any non-cancel path: + // normal completion, business interrupt, or error. The specific outcome + // is conveyed through AgentEvent, not through the cancel state machine. + stateDone int32 = 2 + // stateCancelHandled means the cancel was processed by the runFunc and a + // CancelError was emitted through the event stream. This is the success + // terminal state for cancellation. + stateCancelHandled int32 = 5 +) + +// interruptSent constants (for int32 CAS). +// +// Transition rules: +// +// interruptNotSent -> interruptImmediate (CancelImmediate or escalation) +const ( + // interruptNotSent means no compose graph interrupt has been sent. + interruptNotSent int32 = 0 + // interruptImmediate means an immediate graph interrupt was sent with + // timeout=0, forcing the graph to stop as soon as possible. + interruptImmediate int32 = 1 +) + +// defaultCancelImmediateGracePeriod is the bounded time a recursive +// AgentTool cancel waits before forcing the current level's graph interrupt. +// This gives deeper AgentTool/internal-agent interrupts a chance to bubble up +// as CompositeInterrupts. If this proves insufficient for deeply nested +// structures or too slow for latency-sensitive use cases, consider making it +// configurable via an AgentCancelOption. +const defaultCancelImmediateGracePeriod = 1 * time.Second + +type cancelContextKey struct{} + +// withCancelContext stores a cancelContext in the Go context. +func withCancelContext(ctx context.Context, cc *cancelContext) context.Context { + if cc == nil { + return ctx + } + return context.WithValue(ctx, cancelContextKey{}, cc) +} + +// getCancelContext retrieves the cancelContext from the Go context, or nil. +func getCancelContext(ctx context.Context) *cancelContext { + if v := ctx.Value(cancelContextKey{}); v != nil { + return v.(*cancelContext) + } + return nil +} + +type cancelContext struct { + mode int32 // atomic, CancelMode + + cancelChan chan struct{} // closed when cancel is requested (all modes, not just safe-point) + immediateChan chan struct{} // closed when an immediate graph interrupt fires + doneChan chan struct{} // closed when execution completes (by any mark* method) + doneOnce sync.Once // ensures doneChan is closed exactly once + + state int32 // stateRunning, stateCancelling, stateDone, stateCancelHandled + interruptSent int32 // interruptNotSent, interruptImmediate + escalated int32 // 1 if escalated from safe-point to immediate + timeoutEscalated int32 // 1 if escalation was triggered by timeout + startedMode int32 // atomic, mode when state transitioned to cancelling + deadlineUnixNano int64 // atomic, 0 means no deadline + + recursive int32 // atomic; 1 if cancel should propagate into AgentTool internal agents + recursiveChan chan struct{} // closed when recursive transitions from 0 to 1 + + root bool // true for the original cancelContext created by WithCancel(); false for AgentTool internal agents + parent *cancelContext // non-nil for AgentTool internal agents; used to propagate AgentTool boundary markers upward + + agentToolDescendant int32 // atomic; 1 once an AgentTool runs under this cancel context + + cancelMu sync.Mutex + timeoutOnce sync.Once + timeoutNotify chan struct{} + + mu sync.Mutex + graphInterruptFuncs []func(...compose.GraphInterruptOption) +} + +func newCancelContext() *cancelContext { + return &cancelContext{ + cancelChan: make(chan struct{}), + immediateChan: make(chan struct{}), + doneChan: make(chan struct{}), + timeoutNotify: make(chan struct{}, 1), + recursiveChan: make(chan struct{}), + root: true, + } +} + +func (cc *cancelContext) isRoot() bool { + return cc != nil && cc.root +} + +func (cc *cancelContext) isRecursive() bool { + return cc != nil && atomic.LoadInt32(&cc.recursive) == 1 +} + +// setRecursive(false) is a no-op; recursive is monotonically escalating: +// once set to true, it cannot be reverted. +func (cc *cancelContext) setRecursive(v bool) { + if v && atomic.CompareAndSwapInt32(&cc.recursive, 0, 1) { + close(cc.recursiveChan) + } +} + +// deriveAgentToolCancelContext creates the cancelContext used by an AgentTool's +// internal agent. It receives recursive cancel propagation from the parent +// AgentTool call. The caller MUST ensure the child's markDone() is eventually +// called (e.g., via wrapIterWithCancelCtx's defer) or that ctx is canceled; +// otherwise the two propagation goroutines will leak. +func (cc *cancelContext) deriveAgentToolCancelContext(ctx context.Context) *cancelContext { + if cc == nil { + return nil + } + child := newCancelContext() + child.root = false + child.parent = cc + + // Each goroutine below propagates one signal class (cancel / immediate) to + // the child. The pattern is a two-phase select: + // Phase 1: wait for the parent signal (or child/ctx completion). + // Phase 2: if the signal fired but recursive mode is not active yet, + // enter a second select waiting for either recursive escalation + // (recursiveChan) or child/ctx completion. This ensures + // non-recursive cancels leave children unaware, while a late + // escalation to recursive still propagates. + go func() { + select { + case <-cc.cancelChan: + if cc.isRecursive() { + child.setRecursive(true) + child.triggerCancel(cc.getMode()) + return + } + select { + case <-cc.recursiveChan: + child.setRecursive(true) + child.triggerCancel(cc.getMode()) + case <-child.doneChan: + case <-ctx.Done(): + } + case <-child.doneChan: + case <-ctx.Done(): + } + }() + + go func() { + select { + case <-cc.immediateChan: + if cc.isRecursive() { + child.setRecursive(true) + child.triggerImmediateCancel() + return + } + select { + case <-cc.recursiveChan: + child.setRecursive(true) + child.triggerImmediateCancel() + case <-child.doneChan: + case <-ctx.Done(): + } + case <-child.doneChan: + case <-ctx.Done(): + } + }() + + return child +} + +func (cc *cancelContext) triggerCancel(mode CancelMode) { + cc.setMode(mode) + if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) { + close(cc.cancelChan) + } +} + +func (cc *cancelContext) triggerImmediateCancel() { + atomic.StoreInt32(&cc.escalated, 1) + cc.setMode(CancelImmediate) + if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) { + close(cc.cancelChan) + } + cc.sendImmediateInterrupt() +} + +func (cc *cancelContext) getMode() CancelMode { + if cc == nil { + return CancelImmediate + } + return CancelMode(atomic.LoadInt32(&cc.mode)) +} + +func (cc *cancelContext) setMode(mode CancelMode) { + atomic.StoreInt32(&cc.mode, int32(mode)) +} + +func (cc *cancelContext) getDeadlineUnixNano() int64 { + return atomic.LoadInt64(&cc.deadlineUnixNano) +} + +func (cc *cancelContext) setDeadlineUnixNano(v int64) { + atomic.StoreInt64(&cc.deadlineUnixNano, v) +} + +func (cc *cancelContext) wakeTimeoutController() { + select { + case cc.timeoutNotify <- struct{}{}: + default: + } +} + +// shouldCancel returns true if a cancel has been requested (cancelChan is closed). +func (cc *cancelContext) shouldCancel() bool { + if cc == nil { + return false + } + select { + case <-cc.cancelChan: + return true + default: + return false + } +} + +// isImmediateCancelled returns true if an immediate graph interrupt has been +// fired (CancelImmediate or timeout escalation). This is stronger than +// shouldCancel: it means the compose graph is being torn down right now and +// orphaned goroutines should not attempt to send events. +func (cc *cancelContext) isImmediateCancelled() bool { + if cc == nil { + return false + } + select { + case <-cc.immediateChan: + return true + default: + return false + } +} + +// sendImmediateInterrupt sends the compose graph interrupt signal via graphInterruptFuncs. +// Also closes immediateChan (used by cancelMonitoredModel to abort an in-progress stream). +// Returns false if an interrupt was already sent or if no graphInterruptFuncs have been +// registered yet (the deferred fire in setGraphInterruptFunc will handle that case). +func (cc *cancelContext) sendImmediateInterrupt() bool { + cc.mu.Lock() + + if !atomic.CompareAndSwapInt32(&cc.interruptSent, interruptNotSent, interruptImmediate) { + cc.mu.Unlock() + return false + } + + close(cc.immediateChan) + + fns := make([]func(...compose.GraphInterruptOption), len(cc.graphInterruptFuncs)) + copy(fns, cc.graphInterruptFuncs) + + if cc.isRecursive() && cc.hasAgentToolDescendant() { + select { + case <-cc.doneChan: + cc.mu.Unlock() + return true + case <-time.After(defaultCancelImmediateGracePeriod): + } + } + + if len(fns) == 0 { + cc.mu.Unlock() + return false + } + + for _, fn := range fns { + fn(compose.WithGraphInterruptTimeout(0)) + } + cc.mu.Unlock() + return true +} + +// setGraphInterruptFunc appends a graph interrupt function to the list. +// If an immediate cancel was already requested, fires it retroactively. +// Multiple functions can be registered (e.g. one per parallel sub-agent). +// +// Both this method and sendImmediateInterrupt hold cc.mu across the entire +// check-and-fire sequence, ensuring each interrupt function is called exactly +// once (compose.WithGraphInterrupt returns a non-idempotent closure that panics +// on double-call). +func (cc *cancelContext) setGraphInterruptFunc(interrupt func(...compose.GraphInterruptOption)) { + cc.mu.Lock() + cc.graphInterruptFuncs = append(cc.graphInterruptFuncs, interrupt) + + shouldFire := atomic.LoadInt32(&cc.interruptSent) == interruptImmediate + if shouldFire { + interrupt(compose.WithGraphInterruptTimeout(0)) + } + cc.mu.Unlock() +} + +// markDone marks the execution as finished through any non-cancel path +// (normal completion, business interrupt, or error). +// This is safe to call even if a cancel is in progress — it allows the +// cancel func to detect that execution finished before cancel took effect. +func (cc *cancelContext) markDone() { + if cc == nil { + return + } + if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateDone) || + atomic.CompareAndSwapInt32(&cc.state, stateCancelling, stateDone) { + cc.doneOnce.Do(func() { close(cc.doneChan) }) + } +} + +func (cc *cancelContext) hasAgentToolDescendant() bool { + return cc != nil && atomic.LoadInt32(&cc.agentToolDescendant) == 1 +} + +func (cc *cancelContext) markAgentToolDescendant() { + for cur := cc; cur != nil; cur = cur.parent { + atomic.StoreInt32(&cur.agentToolDescendant, 1) + } +} + +// markCancelHandled signals that the cancel path in the runFunc has created +// and sent a CancelError. Transitions state to stateCancelHandled so that: +// 1. The deferred markDone() becomes a no-op (CAS from cancelling fails). +// 2. buildCancelFunc sees stateCancelHandled and returns nil (cancel succeeded). +// Returns true if the transition succeeded, false if cancel was already handled +// (e.g., by a sub-agent). This prevents duplicate CancelError emission. +func (cc *cancelContext) markCancelHandled() bool { + if cc == nil { + return false + } + if atomic.CompareAndSwapInt32(&cc.state, stateCancelling, stateCancelHandled) { + cc.doneOnce.Do(func() { close(cc.doneChan) }) + return true + } + return false +} + +// createCancelError creates a CancelError based on the current cancel state. +func (cc *cancelContext) createCancelError() *CancelError { + info := &AgentCancelInfo{} + info.Mode = cc.getMode() + if atomic.LoadInt32(&cc.escalated) == 1 { + info.Escalated = true + info.Timeout = atomic.LoadInt32(&cc.timeoutEscalated) == 1 + } + return &CancelError{ + Info: info, + } +} + +func (cc *cancelContext) createAndMarkCancelHandled() (*CancelError, bool) { + cc.cancelMu.Lock() + defer cc.cancelMu.Unlock() + cancelErr := cc.createCancelError() + ok := cc.markCancelHandled() + return cancelErr, ok +} + +// buildCancelFunc builds the AgentCancelFunc for external use. +func (cc *cancelContext) buildCancelFunc() AgentCancelFunc { + joinMode := func(a, b CancelMode) CancelMode { + if a == CancelImmediate || b == CancelImmediate { + return CancelImmediate + } + return a | b + } + + parseReq := func(callOpts ...AgentCancelOption) *agentCancelConfig { + cfg := &agentCancelConfig{Mode: CancelImmediate} + for _, opt := range callOpts { + opt(cfg) + } + return cfg + } + + startTimeoutController := func() { + cc.timeoutOnce.Do(func() { + go func() { + for { + select { + case <-cc.doneChan: + return + default: + } + + mode := cc.getMode() + if mode == CancelImmediate { + return + } + + deadline := cc.getDeadlineUnixNano() + if deadline == 0 { + select { + case <-cc.timeoutNotify: + continue + case <-cc.doneChan: + return + } + } + + now := time.Now().UnixNano() + wait := time.Duration(deadline - now) + if wait <= 0 { + atomic.StoreInt32(&cc.escalated, 1) + atomic.StoreInt32(&cc.timeoutEscalated, 1) + cc.sendImmediateInterrupt() + return + } + + timer := time.NewTimer(wait) + select { + case <-timer.C: + timer.Stop() + atomic.StoreInt32(&cc.escalated, 1) + atomic.StoreInt32(&cc.timeoutEscalated, 1) + cc.sendImmediateInterrupt() + return + case <-cc.timeoutNotify: + timer.Stop() + continue + case <-cc.doneChan: + timer.Stop() + return + } + } + }() + }) + } + + newHandle := func(wait func() error) *CancelHandle { + return &CancelHandle{wait: wait} + } + + waitForCompletion := func() error { + <-cc.doneChan + + st := atomic.LoadInt32(&cc.state) + switch st { + case stateDone: + return ErrExecutionEnded + default: + if atomic.LoadInt32(&cc.timeoutEscalated) == 1 { + return ErrCancelTimeout + } + return nil + } + } + + return func(callOpts ...AgentCancelOption) (*CancelHandle, bool) { + req := parseReq(callOpts...) + + st := atomic.LoadInt32(&cc.state) + switch st { + case stateCancelHandled: + return newHandle(func() error { return nil }), false + case stateDone: + return newHandle(func() error { return ErrExecutionEnded }), false + } + + var needImmediate, needTimeoutCtl bool + + cc.cancelMu.Lock() + + st = atomic.LoadInt32(&cc.state) + switch st { + case stateCancelHandled: + cc.cancelMu.Unlock() + return newHandle(func() error { return nil }), false + case stateDone: + cc.cancelMu.Unlock() + return newHandle(func() error { return ErrExecutionEnded }), false + } + + curMode := cc.getMode() + if st == stateRunning { + if !atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) { + st = atomic.LoadInt32(&cc.state) + cc.cancelMu.Unlock() + if st == stateDone { + return newHandle(func() error { return ErrExecutionEnded }), false + } + return newHandle(waitForCompletion), true + } + + curMode = req.Mode + cc.setMode(curMode) + atomic.StoreInt32(&cc.startedMode, int32(curMode)) + cc.setRecursive(req.Recursive) + close(cc.cancelChan) + } else { + // Recursive is monotonic: once set, cannot be unset. The first + // cancel call uses the bool directly; subsequent calls only + // escalate (false → true) — setRecursive(false) is a no-op. + curMode = joinMode(curMode, req.Mode) + cc.setMode(curMode) + if req.Recursive { + cc.setRecursive(true) + } + } + + if curMode == CancelImmediate { + cc.setDeadlineUnixNano(0) + needImmediate = true + } else if req.Timeout != nil && *req.Timeout > 0 { + proposed := time.Now().Add(*req.Timeout).UnixNano() + existing := cc.getDeadlineUnixNano() + if existing == 0 || proposed < existing { + cc.setDeadlineUnixNano(proposed) + cc.wakeTimeoutController() + } + needTimeoutCtl = cc.getDeadlineUnixNano() != 0 + } + + cc.cancelMu.Unlock() + + if needImmediate { + if atomic.LoadInt32(&cc.startedMode) != int32(CancelImmediate) { + atomic.StoreInt32(&cc.escalated, 1) + } + cc.sendImmediateInterrupt() + } + if needTimeoutCtl { + startTimeoutController() + } + + return newHandle(waitForCompletion), true + } +} + +// wrapIterWithCancelCtx wraps an iterator with cancel lifecycle management. +// It calls markDone when the inner iterator is fully drained, ensuring the +// cancelContext's doneChan is closed and propagation goroutines can exit. +// +// For root cancelContexts (created by WithCancel, not deriveAgentToolCancelContext), it also +// converts interrupt ACTION events to CancelError when cancel is active. +// This is the single point of interrupt-to-CancelError conversion in the +// system — Runner.handleIter only enriches the resulting CancelError with +// checkpoint metadata. +// +// Interrupt absorption: ALL interrupts are converted when cancel is active, +// including business interrupts (compose.Interrupt from user code). Cancel and +// business interrupts cannot be reliably distinguished in concurrent execution +// (parallel workflows, concurrent tool calls) where they merge into a single +// composite signal. The business interrupt data is preserved in the checkpoint +// and re-fires naturally on resume. +// +// This conversion MUST happen in this wrapper (not deferred to Runner.handleIter) +// because markDone runs as a defer in this goroutine — if the interrupt event +// were passed through unconverted, markDone would transition stateCancelling→stateDone +// before the Runner goroutine could call createAndMarkCancelHandled, causing it +// to fail the CAS. +func wrapIterWithCancelCtx[M MessageType](iter *AsyncIterator[*TypedAgentEvent[M]], cancelCtx *cancelContext) *AsyncIterator[*TypedAgentEvent[M]] { + if cancelCtx == nil { + return iter + } + it, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go func() { + defer cancelCtx.markDone() + defer gen.Close() + for { + event, ok := iter.Next() + if !ok { + break + } + + if cancelCtx.isRoot() && event.Action != nil && event.Action.internalInterrupted != nil { + if cancelCtx.shouldCancel() { + cancelErr, ok := cancelCtx.createAndMarkCancelHandled() + if ok { + cancelErr.interruptSignal = event.Action.internalInterrupted + gen.Send(&TypedAgentEvent[M]{Err: cancelErr}) + } + return + } + } + + gen.Send(event) + } + }() + return it +} + +// typedCancelMonitoredModel wraps a model with cancel monitoring. +// Generate: pure delegate to the inner model (CancelAfterChatModel is handled +// by a dedicated node after the ChatModel in the compose graph). +// Stream: pipes chunks through a goroutine that selects on immediateChan for +// CancelImmediate abort. +type typedCancelMonitoredModel[M MessageType] struct { + inner model.BaseModel[M] + cancelContext *cancelContext +} + +type recvResult[T any] struct { + data T + err error +} + +func (m *typedCancelMonitoredModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + return m.inner.Generate(ctx, input, opts...) +} + +func (m *typedCancelMonitoredModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + stream, err := m.inner.Stream(ctx, input, opts...) + if err != nil { + return nil, err + } + wrapped := wrapStreamWithCancelMonitoring(stream, m.cancelContext) + return wrapped, nil +} + +// wrapStreamWithCancelMonitoring wraps a stream with cancel monitoring. +// When immediateChan fires (CancelImmediate or timeout escalation), the output +// stream is terminated with ErrStreamCanceled. +func wrapStreamWithCancelMonitoring[T any](stream *schema.StreamReader[T], cc *cancelContext) *schema.StreamReader[T] { + if cc == nil { + return stream + } + + // Already canceled — terminate immediately + select { + case <-cc.immediateChan: + stream.Close() + r, w := schema.Pipe[T](1) + var zero T + w.Send(zero, ErrStreamCanceled) + w.Close() + return r + default: + } + + reader, writer := schema.Pipe[T](1) + + go func() { + done := make(chan struct{}) + defer close(done) + defer writer.Close() + defer stream.Close() + + ch := make(chan recvResult[T]) + go func() { + defer close(ch) + for { + chunk, recvErr := stream.Recv() + select { + case ch <- recvResult[T]{chunk, recvErr}: + case <-done: + return + } + if recvErr != nil { + return + } + } + }() + + for { + select { + case <-cc.immediateChan: + var zero T + writer.Send(zero, ErrStreamCanceled) + return + + case r, ok := <-ch: + if !ok { + return + } + if r.err != nil { + if r.err == io.EOF { + return + } + var zero T + writer.Send(zero, r.err) + return + } + if closed := writer.Send(r.data, nil); closed { + return + } + } + } + }() + + return reader +} + +// cancelMonitoredToolHandler wraps streamable tool calls with cancel monitoring. +// When CancelImmediate fires, the tool output stream is terminated with ErrStreamCanceled. +// This handler reads the cancelContext from the Go context via getCancelContext. +type cancelMonitoredToolHandler struct{} + +func (h *cancelMonitoredToolHandler) WrapStreamableToolCall(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + output, err := next(ctx, input) + if err != nil { + return nil, err + } + + cc := getCancelContext(ctx) + if cc == nil { + return output, nil + } + + wrapped := wrapStreamWithCancelMonitoring(output.Result, cc) + return &compose.StreamToolOutput{Result: wrapped}, nil + } +} + +func (h *cancelMonitoredToolHandler) WrapEnhancedStreamableToolCall(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + output, err := next(ctx, input) + if err != nil { + return nil, err + } + + cc := getCancelContext(ctx) + if cc == nil { + return output, nil + } + + wrapped := wrapStreamWithCancelMonitoring(output.Result, cc) + return &compose.EnhancedStreamableToolOutput{Result: wrapped}, nil + } +} diff --git a/adk/cancel_edge_test.go b/adk/cancel_edge_test.go new file mode 100644 index 0000000..0c2a80d --- /dev/null +++ b/adk/cancel_edge_test.go @@ -0,0 +1,2085 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// --- helpers shared across edge-case tests --- + +// blockingChatModel blocks until unblockCh is closed, then returns a fixed response. +type blockingChatModel struct { + unblockCh chan struct{} + response *schema.Message + started chan struct{} + callCount int32 +} + +func newBlockingChatModel(response *schema.Message) *blockingChatModel { + return &blockingChatModel{ + unblockCh: make(chan struct{}), + response: response, + started: make(chan struct{}, 1), + } +} + +func (m *blockingChatModel) Generate(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m.callCount, 1) + select { + case m.started <- struct{}{}: + default: + } + <-m.unblockCh + return m.response, nil +} + +func (m *blockingChatModel) Stream(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m.callCount, 1) + select { + case m.started <- struct{}{}: + default: + } + <-m.unblockCh + return schema.StreamReaderFromArray([]*schema.Message{m.response}), nil +} + +func (m *blockingChatModel) BindTools(_ []*schema.ToolInfo) error { return nil } + +// errorChatModel returns an error from Generate/Stream. +type errorChatModel struct { + err error + started chan struct{} +} + +func (m *errorChatModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if m.started != nil { + select { + case m.started <- struct{}{}: + default: + } + } + return nil, m.err +} + +func (m *errorChatModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, m.err +} + +func (m *errorChatModel) BindTools(_ []*schema.ToolInfo) error { return nil } + +// plainResponseModel returns immediately with a fixed text response (no tool calls). +type plainResponseModel struct { + text string +} + +func (m *plainResponseModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage(m.text, nil), nil +} + +func (m *plainResponseModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage(m.text, nil)}), nil +} + +func (m *plainResponseModel) BindTools(_ []*schema.ToolInfo) error { return nil } + +// blockingTool blocks until unblockCh is closed. +type blockingTool struct { + name string + unblockCh chan struct{} + started chan struct{} + callCount int32 +} + +func newBlockingTool(name string) *blockingTool { + return &blockingTool{ + name: name, + unblockCh: make(chan struct{}), + started: make(chan struct{}, 4), + } +} + +func (t *blockingTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "blocking tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string"}, + }), + }, nil +} + +func (t *blockingTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + atomic.AddInt32(&t.callCount, 1) + select { + case t.started <- struct{}{}: + default: + } + <-t.unblockCh + return "result", nil +} + +func toolCallMsg(calls ...schema.ToolCall) *schema.Message { + return &schema.Message{Role: schema.Assistant, ToolCalls: calls} +} + +func toolCall(id, name, args string) schema.ToolCall { + return schema.ToolCall{ID: id, Type: "function", Function: schema.FunctionCall{Name: name, Arguments: args}} +} + +func drainEvents(iter *AsyncIterator[*AgentEvent]) ([]*AgentEvent, bool) { + var events []*AgentEvent + hasCancelError := false + for { + e, ok := iter.Next() + if !ok { + break + } + events = append(events, e) + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + hasCancelError = true + } + } + return events, hasCancelError +} + +// --- tests --- + +// TestWithCancel_BeforeExecutionStarts verifies that a cancel issued before +// the graph begins executing still produces a CancelError without invoking +// the model or tools. +func TestWithCancel_BeforeExecutionStarts(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(toolCallMsg(toolCall("c1", "bt", `{"input":"x"}`))) + bt := newBlockingTool("bt") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{bt}}, + }, + }) + assert.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + + // Extract the cancelContext so we can wait for cancelChan to close, + // ensuring the cancel is fully registered before Run starts. + cc := getCommonOptions(nil, cancelOpt).cancelCtx + + // Call cancel BEFORE calling agent.Run. + // The cancelFunc must succeed (not hang) even though execution hasn't started. + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn() + cancelDone <- handle.Wait() + }() + + // Wait for cancelChan to close so the pre-execution check in runFunc + // deterministically sees shouldCancel()=true (eliminates goroutine scheduling race). + <-cc.cancelChan + + // Now start the run — it should see shouldCancel()=true and emit CancelError immediately. + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, cancelOpt) + + _, hasCancelError := drainEvents(iter) + assert.True(t, hasCancelError, "expected CancelError when cancel precedes execution") + + // cancelFn must have already returned (or return quickly now that doneChan is closed). + select { + case cancelErr := <-cancelDone: + // Either nil (cancel handled) or ErrExecutionEnded is acceptable + // depending on exact timing; what matters is it didn't hang. + _ = cancelErr + case <-time.After(3 * time.Second): + t.Fatal("cancelFn blocked indefinitely after pre-start cancel") + } + + // Model and tool must not have been invoked. + assert.Equal(t, int32(0), atomic.LoadInt32(&bt.callCount), "tool must not be called") +} + +// TestWithCancel_AfterCompletion verifies cancelFn returns ErrExecutionEnded +// when called after a normal run finishes. +func TestWithCancel_AfterCompletion(t *testing.T) { + ctx := context.Background() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: &plainResponseModel{text: "done"}, + }) + require.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, cancelOpt) + + // Drain all events so the run completes. + for { + _, ok := iter.Next() + if !ok { + break + } + } + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.ErrorIs(t, cancelErr, ErrExecutionEnded) +} + +// TestWithCancel_DerivedAgentToolCancelContextMarkedDoneAfterRun verifies that +// an explicitly derived AgentTool child cancel context is owned by the child run, +// even when the Go context also carries the parent cancel context. +func TestWithCancel_DerivedAgentToolCancelContextMarkedDoneAfterRun(t *testing.T) { + ctx := context.Background() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ChildAgent", + Description: "test child agent", + Model: &plainResponseModel{text: "done"}, + }) + require.NoError(t, err) + + parent := newCancelContext() + parentCtx := withCancelContext(ctx, parent) + child := parent.deriveAgentToolCancelContext(parentCtx) + + childOpt := WrapImplSpecificOptFn(func(o *options) { + o.cancelCtx = child + }) + iter := agent.Run(parentCtx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, childOpt) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + select { + case <-child.doneChan: + case <-time.After(time.Second): + t.Fatal("derived AgentTool cancel context was not marked done after child run completion") + } +} + +// TestWithCancel_AfterBusinessInterrupt verifies cancelFn returns ErrExecutionEnded +// when called after the agent has been interrupted by business logic. +func TestWithCancel_AfterBusinessInterrupt(t *testing.T) { + ctx := context.Background() + + // Use a model that triggers a compose.Interrupt so the agent stops with an interrupt. + interruptModel := &interruptingChatModel{} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: interruptModel, + }) + require.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hi")}, cancelOpt, WithCheckPointID("biz-interrupt-1")) + + // Drain — expect an interrupt action event, no cancel error. + var gotInterrupt bool + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Action != nil && e.Action.Interrupted != nil { + gotInterrupt = true + } + } + assert.True(t, gotInterrupt, "expected business interrupt event") + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.ErrorIs(t, cancelErr, ErrExecutionEnded) +} + +// TestWithCancel_AfterError verifies cancelFn returns ErrExecutionEnded +// when called after the agent errors out. +func TestWithCancel_AfterError(t *testing.T) { + ctx := context.Background() + + modelErr := errors.New("model exploded") + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: &errorChatModel{err: modelErr}, + }) + require.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, cancelOpt) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.ErrorIs(t, cancelErr, ErrExecutionEnded) +} + +// TestWithCancel_TimeoutEscalation tests that WithAgentCancelTimeout causes the +// cancel to escalate to immediate when the safe-point hasn't fired yet, and +// that the resulting CancelError has Escalated=true. +// +// Strategy: use CancelAfterChatModel mode. The model blocks (never completes), +// so the safe-point can't fire naturally. After the timeout, escalateToImmediate +// closes immediateChan which aborts the model stream via cancelMonitoredModel +// and causes a CancelError — no compose graph-interrupt races involved. +func TestWithCancel_TimeoutEscalation(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(schema.AssistantMessage("hello", nil)) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, // use streaming so cancelMonitoredModel.Stream is exercised + }) + + timeout := 300 * time.Millisecond + // CancelAfterChatModel + timeout: safe-point can't fire (model never finishes), + // so after 300ms the timeout goroutine escalates to immediate. + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("go")}, cancelOpt) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start") + } + + // Fire cancelFn; it will wait for escalation to complete. + start := time.Now() + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel), WithAgentCancelTimeout(timeout)) + cancelErr := handle.Wait() + elapsed := time.Since(start) + + assert.ErrorIs(t, cancelErr, ErrCancelTimeout, "cancel should return ErrCancelTimeout after timeout escalation") + assert.True(t, elapsed >= timeout, "should wait at least the timeout duration, elapsed=%v", elapsed) + assert.True(t, elapsed < 3*time.Second, "should complete shortly after timeout, elapsed=%v", elapsed) + + var cancelError *CancelError + for { + e, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + cancelError = ce + } + } + if assert.NotNil(t, cancelError, "expected CancelError after timeout escalation") { + assert.True(t, cancelError.Info.Escalated, "CancelError should report Escalated=true") + assert.True(t, cancelError.Info.Timeout, "CancelError should report Timeout=true") + } +} + +// TestWithCancel_AfterChatModel_WithTools verifies CancelAfterChatModel fires +// when the model returns tool calls (the safe-point is on the tool-calls path). +func TestWithCancel_AfterChatModel_WithTools(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(toolCallMsg(toolCall("c1", "bt", `{"input":"x"}`))) + bt := newBlockingTool("bt") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{bt}}, + }, + }) + require.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, cancelOpt) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start") + } + + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + cancelDone <- handle.Wait() + }() + + time.Sleep(20 * time.Millisecond) + + close(blk.unblockCh) + + cancelErr := <-cancelDone + assert.NoError(t, cancelErr) + + _, hasCancelError := drainEvents(iter) + assert.True(t, hasCancelError, "CancelError expected after model returns tool calls") +} + +// TestWithCancel_CancelImmediate_StreamAborted verifies that CancelImmediate +// during model execution surfaces CancelError and completes quickly. +// Uses blockingChatModel which blocks in Stream(), keeping the agent's run +// function alive so the cancel context stays in stateRunning. +func TestWithCancel_CancelImmediate_StreamAborted(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(schema.AssistantMessage("hello", nil)) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hi")}, cancelOpt) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start") + } + time.Sleep(50 * time.Millisecond) + + start := time.Now() + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + elapsed := time.Since(start) + assert.True(t, elapsed < 2*time.Second, "cancel should complete quickly, elapsed=%v", elapsed) + + var foundCancelError bool + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Action != nil && e.Action.Interrupted != nil { + foundCancelError = true + } + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + foundCancelError = true + } + } + assert.True(t, foundCancelError, "expected CancelError in event stream") +} + +// TestWithCancel_MultipleToolsConcurrent verifies that CancelAfterToolCalls +// waits for ALL concurrent tool calls to complete before cancelling. +func TestWithCancel_MultipleToolsConcurrent(t *testing.T) { + ctx := context.Background() + + bt1 := newBlockingTool("tool1") + bt2 := newBlockingTool("tool2") + + // Model calls both tools in one response. + modelResp := toolCallMsg( + toolCall("c1", "tool1", `{"input":"a"}`), + toolCall("c2", "tool2", `{"input":"b"}`), + ) + modelWithTools := &simpleChatModel{response: modelResp} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: modelWithTools, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{bt1, bt2}}, + }, + }) + assert.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}}, cancelOpt) + + // Wait for both tools to start. + for i := 0; i < 2; i++ { + select { + case <-bt1.started: + case <-bt2.started: + case <-time.After(5 * time.Second): + t.Fatal("tools did not start") + } + } + + // Request cancel after tool calls while both are still blocking. + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + cancelDone <- handle.Wait() + }() + + // Unblock both tools — cancel should fire only after both complete. + time.Sleep(50 * time.Millisecond) + close(bt1.unblockCh) + time.Sleep(50 * time.Millisecond) + close(bt2.unblockCh) + + cancelErr := <-cancelDone + assert.NoError(t, cancelErr) + + assert.Equal(t, int32(1), atomic.LoadInt32(&bt1.callCount), "tool1 should complete") + assert.Equal(t, int32(1), atomic.LoadInt32(&bt2.callCount), "tool2 should complete") + + _, hasCancelError := drainEvents(iter) + assert.True(t, hasCancelError, "expected CancelError after concurrent tools completed") +} + +// TestWithCancel_GraphInterruptRaceBeforeSet verifies that a CancelImmediate +// issued before setGraphInterruptFunc is called still results in cancellation. +// This exercises the retroactive-fire path in setGraphInterruptFunc. +func TestWithCancel_GraphInterruptRaceBeforeSet(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(schema.AssistantMessage("hi", nil)) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + }) + require.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + + // Cancel immediately before run starts. + go func() { + handle, _ := cancelFn() + _ = handle.Wait() + }() + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, cancelOpt) + + done := make(chan struct{}) + go func() { + defer close(done) + drainEvents(iter) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("iteration did not complete after pre-start CancelImmediate") + } +} + +// TestWithCancel_NoCheckpointStore verifies cancel completes and does not panic +// when no checkpoint store is configured. +func TestWithCancel_NoCheckpointStore(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(schema.AssistantMessage("hi", nil)) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + // No CheckPointStore set. + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hi")}, cancelOpt) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start") + } + time.Sleep(30 * time.Millisecond) + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + var ce *CancelError + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Err != nil && errors.As(e.Err, &ce) { + break + } + } + assert.NotNil(t, ce, "expected CancelError even without checkpoint store") +} + +// TestWithCancel_ModelError verifies that a model error marks the cancelCtx as +// done so that a subsequent cancelFn call returns ErrExecutionEnded. +func TestWithCancel_ModelError(t *testing.T) { + ctx := context.Background() + + modelErr := errors.New("model failed") + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: &errorChatModel{err: modelErr}, + }) + require.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hi")}}, cancelOpt) + + var gotModelErr bool + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Err != nil && !errors.As(e.Err, new(*CancelError)) { + gotModelErr = true + } + } + assert.True(t, gotModelErr, "expected non-cancel error event from model failure") + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.ErrorIs(t, cancelErr, ErrExecutionEnded, "cancelFn should return ErrExecutionEnded after model error") +} + +// TestWithCancel_Resume_SafePoint covers CancelAfterChatModel and +// CancelAfterToolCalls on a Resume path. +func TestWithCancel_Resume_SafePoint(t *testing.T) { + ctx := context.Background() + + // --- phase 1: run to get a checkpoint via CancelImmediate --- + blk := newBlockingChatModel(toolCallMsg(toolCall("c1", "bt", `{"input":"x"}`))) + bt := newSlowTool("bt", 50*time.Millisecond, "result") + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{bt}}, + }, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner1 := NewRunner(ctx, RunnerConfig{ + Agent: agent1, + CheckPointStore: store, + }) + + cancelOpt1, cancelFn1 := WithCancel() + iter1 := runner1.Run(ctx, []Message{schema.UserMessage("hi")}, cancelOpt1, WithCheckPointID("resume-sp-1")) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start in phase 1") + } + _, _ = cancelFn1() + drainEvents(iter1) + + // --- phase 2: resume, cancel after chat model --- + resumeModel := newBlockingChatModel(toolCallMsg(toolCall("c1", "bt", `{"input":"x"}`))) + + bt2 := newSlowTool("bt", 50*time.Millisecond, "result") + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: resumeModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{bt2}}, + }, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: agent2, + CheckPointStore: store, + }) + + cancelOpt2, cancelFn2 := WithCancel() + resumeIter, err := runner2.Resume(ctx, "resume-sp-1", cancelOpt2) + require.NoError(t, err) + + select { + case <-resumeModel.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start in phase 2") + } + + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn2(WithAgentCancelMode(CancelAfterChatModel)) + cancelDone <- handle.Wait() + }() + + time.Sleep(50 * time.Millisecond) + + close(resumeModel.unblockCh) + + cancelErr := <-cancelDone + assert.NoError(t, cancelErr) + + _, hasCancelError := drainEvents(resumeIter) + assert.True(t, hasCancelError, "CancelError expected after resumed model returns tool calls") +} + +// callbackTool is a tool that calls onCall when invoked. +type callbackTool struct { + name string + onCall func() +} + +func (t *callbackTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "callback tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string"}, + }), + }, nil +} + +func (t *callbackTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + if t.onCall != nil { + t.onCall() + } + return "ok", nil +} + +// interruptingChatModel returns a compose.Interrupt error to simulate a +// business interrupt during execution. +type interruptingChatModel struct{} + +func (m *interruptingChatModel) Generate(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, compose.Interrupt(ctx, "test interrupt") +} + +func (m *interruptingChatModel) Stream(ctx context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, compose.Interrupt(ctx, "test interrupt") +} + +func (m *interruptingChatModel) BindTools(_ []*schema.ToolInfo) error { return nil } + +// TestWithCancel_TargetedResume_CancelImmediate cancels an agent via CancelImmediate, +// extracts InterruptContexts from the resulting CancelError, and uses them +// for targeted resumption via Runner.ResumeWithParams. +func TestWithCancel_TargetedResume_CancelImmediate(t *testing.T) { + ctx := context.Background() + + blk := newBlockingChatModel(toolCallMsg(toolCall("c1", "st", `{"input":"x"}`))) + st := newSlowTool("st", 50*time.Millisecond, "result") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("go")}, cancelOpt, WithCheckPointID("targeted-imm-1")) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start") + } + + handle, _ := cancelFn() // CancelImmediate (default) + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + var cancelError *CancelError + for { + e, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + cancelError = ce + } + } + + require.NotNil(t, cancelError, "expected CancelError") + require.NotEmpty(t, cancelError.InterruptContexts, "CancelError should have InterruptContexts for targeted resume") + + // --- resume with targeted params --- + targets := make(map[string]any) + for _, ic := range cancelError.InterruptContexts { + targets[ic.ID] = nil + } + + resumeModel := &plainResponseModel{text: "resumed"} + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: resumeModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: agent2, + CheckPointStore: store, + }) + + resumeIter, err := runner2.ResumeWithParams(ctx, "targeted-imm-1", &ResumeParams{Targets: targets}) + require.NoError(t, err) + + var gotOutput bool + for { + e, ok := resumeIter.Next() + if !ok { + break + } + if e.Err != nil { + t.Fatalf("unexpected error during targeted resume: %v", e.Err) + } + if e.Output != nil && e.Output.MessageOutput != nil { + gotOutput = true + } + } + assert.True(t, gotOutput, "targeted resume should produce output") +} + +// TestWithCancel_TargetedResume_SafePoint cancels an agent via CancelAfterChatModel +// (safe-point) and verifies that InterruptContexts are populated on the CancelError +// and that targeted resume via ResumeWithParams succeeds. +// Since safe-point cancels now use compose.Interrupt, compose saves checkpoint data, +// making the cancel fully resumable. +func TestWithCancel_TargetedResume_SafePoint(t *testing.T) { + ctx := context.Background() + + // The model returns a tool call so the react graph routes to toolPreHandle, + // which detects CancelAfterChatModel and fires compose.Interrupt. + blk := newBlockingChatModel(toolCallMsg(toolCall("c1", "st", `{"input":"x"}`))) + st := newSlowTool("st", 0, "result") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("go")}, cancelOpt, WithCheckPointID("targeted-sp-1")) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start") + } + + // Start cancelFn in background so the CAS happens before the model unblocks. + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + cancelDone <- handle.Wait() + }() + time.Sleep(50 * time.Millisecond) + close(blk.unblockCh) + + cancelErr := <-cancelDone + assert.NoError(t, cancelErr) + + var cancelError *CancelError + for { + e, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + cancelError = ce + } + } + + require.NotNil(t, cancelError, "expected CancelError") + require.NotEmpty(t, cancelError.InterruptContexts, "CancelError should have InterruptContexts for targeted resume") + + // --- resume with targeted params --- + targets := make(map[string]any) + for _, ic := range cancelError.InterruptContexts { + targets[ic.ID] = nil + } + + resumeModel := &plainResponseModel{text: "resumed"} + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: resumeModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: agent2, + CheckPointStore: store, + }) + + resumeIter, err := runner2.ResumeWithParams(ctx, "targeted-sp-1", &ResumeParams{Targets: targets}) + require.NoError(t, err) + + var gotOutput bool + for { + e, ok := resumeIter.Next() + if !ok { + break + } + if e.Err != nil { + t.Fatalf("unexpected error during targeted resume: %v", e.Err) + } + if e.Output != nil && e.Output.MessageOutput != nil { + gotOutput = true + } + } + assert.True(t, gotOutput, "targeted resume should produce output") +} + +// TestWithCancel_Resume_CancelAfterChatModel_MessagePreserved tests both the +// ReAct (with-tools) and noTools paths to ensure that when a +// CancelAfterChatModel safe-point fires and the run is later resumed, the +// original Message returned by the chat model is preserved through the +// StatefulInterrupt checkpoint. +// +// For the ReAct path: the model returns a tool-call message. On resume the +// cancelCheck node must return that same message so the branch routes to the +// ToolNode and the tool actually executes. +// +// For the noTools path: the model returns a plain text message. On resume the +// cancel-check lambda must return that same message as the chain output. +func TestWithCancel_Resume_CancelAfterChatModel_MessagePreserved(t *testing.T) { + t.Run("react_path_tool_call_preserved", func(t *testing.T) { + ctx := context.Background() + + // Phase-2 model returns no tool calls so the graph ends. + // We track whether the tool actually executes on resume. + toolExecuted := make(chan struct{}, 1) + st := &callbackTool{ + name: "my_tool", + onCall: func() { + select { + case toolExecuted <- struct{}{}: + default: + } + }, + } + + // Phase-1 model returns a tool call. + blk := newBlockingChatModel(toolCallMsg(toolCall("c1", "my_tool", `{"input":"x"}`))) + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: blk, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + store := newCancelTestStore() + runner1 := NewRunner(ctx, RunnerConfig{ + Agent: agent1, + CheckPointStore: store, + }) + + cancelOpt1, cancelFn1 := WithCancel() + iter1 := runner1.Run(ctx, []Message{schema.UserMessage("hi")}, + cancelOpt1, WithCheckPointID("react-msg-preserved-1")) + + select { + case <-blk.started: + case <-time.After(5 * time.Second): + t.Fatal("model did not start in phase 1") + } + + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn1(WithAgentCancelMode(CancelAfterChatModel)) + cancelDone <- handle.Wait() + }() + time.Sleep(50 * time.Millisecond) + close(blk.unblockCh) + + cancelErr := <-cancelDone + assert.NoError(t, cancelErr) + + _, hasCancelError := drainEvents(iter1) + assert.True(t, hasCancelError, "expected CancelError from phase 1") + + // Phase 2: resume. The model for phase-2 returns plain text (no tool + // calls) so the react graph ends after one iteration. But first the + // tool from the checkpoint must execute. + resumeModel := &plainResponseModel{text: "done"} + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: resumeModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: agent2, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, "react-msg-preserved-1") + require.NoError(t, err) + + for { + e, ok := resumeIter.Next() + if !ok { + break + } + if e.Err != nil { + t.Fatalf("unexpected error during resume: %v", e.Err) + } + } + + // The key assertion: the tool must have been called during resume, + // which can only happen if the tool-call message was preserved. + select { + case <-toolExecuted: + // success + default: + t.Fatal("tool was not executed on resume — the tool-call message was lost") + } + }) + +} + +// TestHandleRunFuncError_AlreadyHandled_NoDuplicate verifies that when +// markCancelHandled() was already claimed by a sub-agent's handleRunFuncError, +// the sequential workflow's checkCancel does not emit a second CancelError. +// +// Setup: sequential[cma1, cma2] with CancelAfterToolCalls. cma1 has tools, +// cancel fires while tool is running. After tool completes, the safe-point +// fires in cma1's handleRunFuncError (claiming markCancelHandled). The +// sequential workflow's checkCancel at the transition point should find +// markCancelHandled returns false and skip — producing exactly 1 CancelError. +func TestHandleRunFuncError_AlreadyHandled_NoDuplicate(t *testing.T) { + ctx := context.Background() + + bt := newBlockingTool("bt") + + // cma1: model returns a tool call immediately, tool blocks until unblocked + cma1Model := newBlockingChatModel(toolCallMsg(toolCall("c1", "bt", `{"input":"x"}`))) + close(cma1Model.unblockCh) // model returns immediately + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", + Model: cma1Model, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{bt}}, + }, + }) + require.NoError(t, err) + + agent2Model := &plainResponseModel{text: "agent2-response"} + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", + Model: agent2Model, + }) + require.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "sequential", SubAgents: []Agent{agent1, agent2}, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + // Wait for tool to start + select { + case <-bt.started: + case <-time.After(5 * time.Second): + t.Fatal("Tool did not start") + } + + // Cancel while tool is still running (in goroutine because cancelFn blocks + // until execution finishes), then unblock tool so safe-point fires + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + _ = handle.Wait() + }() + + // Give cancel time to register, then unblock tool + time.Sleep(50 * time.Millisecond) + close(bt.unblockCh) + + cancelCount := 0 + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + cancelCount++ + } + } + + assert.Equal(t, 1, cancelCount, "Should have exactly one CancelError, no duplicate from handleRunFuncError + checkCancel") +} + +func TestWithCancel_CancelAfterChatModel_NestedAgentTool(t *testing.T) { + ctx := context.Background() + + subAgentModel := newBlockingChatModel(toolCallMsg(toolCall("c1", "sub_tool", `{"input":"x"}`))) + subAgentModelStarted := subAgentModel.started + subTool := newBlockingTool("sub_tool") + + subAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "sub_agent", + Description: "test sub agent", + Instruction: "you are a sub agent", + Model: subAgentModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{subTool}}, + }, + }) + require.NoError(t, err) + + supervisorModel := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call_1", Type: "function", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "sub_agent"}`, + }, + }}, + }, + } + + supervisorAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "supervisor", + Description: "supervisor agent (equivalent to DeepAgent)", + Instruction: "you are a supervisor", + Model: supervisorModel, + }) + require.NoError(t, err) + + agentWithSubAgents, err := SetSubAgents(ctx, supervisorAgent, []Agent{subAgent}) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agentWithSubAgents, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + select { + case <-subAgentModelStarted: + case <-time.After(10 * time.Second): + t.Fatal("Sub-agent model did not start") + } + + time.Sleep(50 * time.Millisecond) + + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel), WithRecursive()) + cancelDone <- handle.Wait() + }() + + time.Sleep(100 * time.Millisecond) + close(subAgentModel.unblockCh) + + cancelErr := <-cancelDone + assert.NoError(t, cancelErr) + + hasCancelError := false + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "CancelError expected from nested agent tool with tools") +} + +// slowStreamingTool implements StreamableTool (but NOT InvokableTool), streaming +// chunks slowly so CancelImmediate can fire mid-stream. +type slowStreamingTool struct { + name string + chunkInterval time.Duration + chunks []string + started chan struct{} + gate chan struct{} // if non-nil, blocks after first chunk until closed +} + +func (t *slowStreamingTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "slow streaming tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string"}, + }), + }, nil +} + +func (t *slowStreamingTool) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) { + r, w := schema.Pipe[string](1) + go func() { + defer w.Close() + select { + case t.started <- struct{}{}: + default: + } + for i, chunk := range t.chunks { + time.Sleep(t.chunkInterval) + if closed := w.Send(chunk, nil); closed { + return + } + // After the second chunk, block on gate so the caller can + // issue a cancel while the tool is deterministically still streaming. + // We wait until chunk index 1 (second chunk) so that the framework + // has time to receive the first chunk and forward the streaming + // event to the iterator, ensuring ErrStreamCanceled is observable. + if i == 1 && t.gate != nil { + <-t.gate + } + } + }() + return r, nil +} + +// toolCallStreamModel returns a tool-call message on the first Stream call, +// then a plain text response on subsequent calls. +type toolCallStreamModel struct { + callCount int32 +} + +func (m *toolCallStreamModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if atomic.AddInt32(&m.callCount, 1) == 1 { + return toolCallMsg(toolCall("c1", "slow_tool", `{"input":"x"}`)), nil + } + return schema.AssistantMessage("done", nil), nil +} + +func (m *toolCallStreamModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *toolCallStreamModel) BindTools(_ []*schema.ToolInfo) error { return nil } + +// TestWithCancel_CancelImmediate_StreamableToolAborted verifies that CancelImmediate +// during StreamableTool streaming surfaces ErrStreamCanceled on the tool's +// MessageStream.Recv(), just like it does for ChatModel streaming. +func TestWithCancel_CancelImmediate_StreamableToolAborted(t *testing.T) { + ctx := context.Background() + + tcm := &toolCallStreamModel{} + gate := make(chan struct{}) + st := &slowStreamingTool{ + name: "slow_tool", + chunkInterval: 100 * time.Millisecond, + chunks: []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}, + started: make(chan struct{}, 1), + gate: gate, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: tcm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{st}}, + }, + }) + require.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hi")}, cancelOpt) + + // Wait for the tool to start streaming and send its first chunk. + // The tool then blocks on the gate, guaranteeing the execution is + // still in progress when we issue the cancel. + select { + case <-st.started: + case <-time.After(5 * time.Second): + t.Fatal("tool did not start streaming") + } + + // Drain events in a separate goroutine so we can issue the cancel + // from the main goroutine after confirming the tool stream event + // has been received. + type result struct { + foundStreamCanceled bool + foundCancelError bool + } + resultCh := make(chan result, 1) + toolStreamReady := make(chan struct{}) + go func() { + var r result + for { + e, ok := iter.Next() + if !ok { + break + } + + // ErrStreamCanceled appears on the tool's MessageStream.Recv() + if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.IsStreaming && + e.Output.MessageOutput.Role == schema.Tool { + // Signal that the tool stream event has been received. + close(toolStreamReady) + stream := e.Output.MessageOutput.MessageStream + for { + _, recvErr := stream.Recv() + if recvErr != nil { + if errors.Is(recvErr, ErrStreamCanceled) { + r.foundStreamCanceled = true + } + break + } + } + } + + if e.Action != nil && e.Action.Interrupted != nil { + r.foundCancelError = true + } + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + r.foundCancelError = true + } + } + resultCh <- r + }() + + // Wait for the iterator goroutine to receive the tool streaming event. + // At this point the tool goroutine is blocked on the gate, and the + // iterator goroutine is blocked on stream.Recv(), so the execution is + // guaranteed to still be in progress. + select { + case <-toolStreamReady: + case <-time.After(5 * time.Second): + t.Fatal("tool stream event was not received by the iterator") + } + + // Issue cancel while the tool goroutine is blocked on gate. + // wrapStreamWithCancelMonitoring detects immediateChan and sends + // ErrStreamCanceled to the consumer side. We do NOT close gate here — + // keeping the tool goroutine blocked ensures the graph interrupt (timeout=0) + // wins the race against normal completion. Close gate in defer for cleanup. + defer close(gate) + handle, _ := cancelFn() + cancelErr := handle.Wait() + + r := <-resultCh + + if errors.Is(cancelErr, ErrExecutionEnded) { + // On slower runtimes (e.g. Go 1.19 CI), the execution can complete + // before the cancel signal is delivered — this is a valid race outcome. + t.Log("cancel raced with completion (ErrExecutionEnded) — skipping cancel assertions") + return + } + assert.NoError(t, cancelErr) + assert.True(t, r.foundStreamCanceled, "expected ErrStreamCanceled on tool's MessageStream.Recv()") + assert.True(t, r.foundCancelError, "expected CancelError in event stream") +} + +// TestWithCancel_CancelImmediate_NestedAgentTool_ResumeFromToolsNode verifies that +// when a nested ChatModelAgent (wrapped as an AgentTool inside an outer ChatModelAgent) +// is canceled via CancelImmediate and then resumed with Runner.Resume (no params), +// the outer agent resumes from the ToolsNode rather than restarting from the beginning. +// +// Regression test: previously, the outer ChatModelAgent would restart from its Init/ChatModel +// node instead of resuming from the ToolsNode, causing the outer model to be called again +// with the original user message before the AgentTool and inner ChatModelAgent were resumed. +func TestWithCancel_CancelImmediate_NestedAgentTool_ResumeFromToolsNode(t *testing.T) { + for _, tc := range []struct { + name string + enableStreaming bool + innerHasTools bool + recursive bool + }{ + {"Invoke_InnerNoTools_NonRecursive", false, false, false}, + {"Stream_InnerNoTools_NonRecursive", true, false, false}, + {"Invoke_InnerWithTools_NonRecursive", false, true, false}, + {"Stream_InnerWithTools_NonRecursive", true, true, false}, + {"Invoke_InnerNoTools_Recursive", false, false, true}, + {"Stream_InnerNoTools_Recursive", true, false, true}, + {"Invoke_InnerWithTools_Recursive", false, true, true}, + {"Stream_InnerWithTools_Recursive", true, true, true}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + // --- inner agent: its model blocks so we can cancel mid-execution --- + var innerTools []tool.BaseTool + var innerModelResp *schema.Message + if tc.innerHasTools { + innerModelResp = toolCallMsg(toolCall("ic1", "inner_tool", `{"input":"x"}`)) + innerTools = []tool.BaseTool{newBlockingTool("inner_tool")} + } else { + innerModelResp = &schema.Message{Role: schema.Assistant, Content: "inner agent done"} + } + innerModel := newBlockingChatModel(innerModelResp) + t.Cleanup(func() { + close(innerModel.unblockCh) + }) + + innerCfg := &ChatModelAgentConfig{ + Name: "InnerAgent", + Description: "inner agent that blocks", + Instruction: "you are an inner agent", + Model: innerModel, + } + if len(innerTools) > 0 { + innerCfg.ToolsConfig = ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: innerTools}, + } + } + innerAgent, err := NewChatModelAgent(ctx, innerCfg) + require.NoError(t, err) + + // --- outer agent: counting model --- + // Call 1: returns a tool call that invokes InnerAgent. + // Call 2 (only needed on resume): returns a plain final answer. + outerModelCallCount := int32(0) + outerModel := &countingChatModel{ + callCount: &outerModelCallCount, + responses: []*schema.Message{ + toolCallMsg(toolCall("c1", "InnerAgent", `{"request":"do something"}`)), + schema.AssistantMessage("outer completed", nil), + }, + } + + outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "outer agent with nested agent tool", + Instruction: "you are an outer agent", + Model: outerModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, innerAgent)}, + }, + }, + }) + require.NoError(t, err) + + store := newCancelTestStore() + checkpointID := "cancel-nested-resume-" + tc.name + + runner1 := NewRunner(ctx, RunnerConfig{ + Agent: outerAgent, + EnableStreaming: tc.enableStreaming, + CheckPointStore: store, + }) + + // --- phase 1: run and cancel while inner agent model is blocked --- + cancelOpt, cancelFn := WithCancel() + iter := runner1.Run(ctx, []Message{schema.UserMessage("go")}, cancelOpt, WithCheckPointID(checkpointID)) + + // Wait for inner model to start (meaning outer model already returned tool call). + select { + case <-innerModel.started: + case <-time.After(5 * time.Second): + t.Fatal("inner model did not start") + } + + // At this point outerModel should have been called exactly once. + assert.Equal(t, int32(1), atomic.LoadInt32(&outerModelCallCount), + "outer model should have been called once before cancel") + + // Cancel immediately. Recursive cases additionally propagate the cancel + // request into the AgentTool's internal ChatModelAgent. + var handle *CancelHandle + if tc.recursive { + handle, _ = cancelFn(WithRecursive()) + } else { + handle, _ = cancelFn() + } + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + _, hasCancelError := drainEvents(iter) + assert.True(t, hasCancelError, "expected CancelError from canceled nested agent tool") + + // --- phase 2: resume with Runner.Resume (no ResumeWithParams, no interrupt ID) --- + // Build fresh agents for resume. Recursive cancel should resume the + // inner ChatModelAgent inside AgentTool before the top-level + // ChatModelAgent produces the final answer. + resumeFirstModelCall := make(chan string, 5) + resumeOuterModelCallCount := int32(0) + resumeOuterModel := &countingChatModel{ + callCount: &resumeOuterModelCallCount, + callCh: resumeFirstModelCall, + callLabel: "outer", + responses: []*schema.Message{ + schema.AssistantMessage("outer completed after resume", nil), + }, + } + + resumeInnerModelCallCount := int32(0) + resumeInnerResponses := []*schema.Message{schema.AssistantMessage("inner agent done after resume", nil)} + if len(innerTools) > 0 { + resumeInnerResponses = []*schema.Message{ + toolCallMsg(toolCall("ic1", "inner_tool", `{"input":"x"}`)), + schema.AssistantMessage("inner agent done after resume", nil), + } + } + resumeInnerModel := &countingChatModel{ + callCount: &resumeInnerModelCallCount, + callCh: resumeFirstModelCall, + callLabel: "inner", + responses: resumeInnerResponses, + } + resumeInnerCfg := &ChatModelAgentConfig{ + Name: "InnerAgent", + Description: "inner agent that returns immediately on resume", + Instruction: "you are an inner agent", + Model: resumeInnerModel, + } + if len(innerTools) > 0 { + resumeInnerCfg.ToolsConfig = ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{newSlowTool("inner_tool", 0, "inner tool result")}, + }, + } + } + resumeInnerAgent, err := NewChatModelAgent(ctx, resumeInnerCfg) + require.NoError(t, err) + + resumeOuterAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "outer agent with nested agent tool", + Instruction: "you are an outer agent", + Model: resumeOuterModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, resumeInnerAgent)}, + }, + }, + }) + require.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumeOuterAgent, + EnableStreaming: tc.enableStreaming, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, checkpointID) + require.NoError(t, err) + + select { + case firstModel := <-resumeFirstModelCall: + if tc.recursive { + assert.Equal(t, "inner", firstModel, + "recursive cancel should resume the AgentTool/internal ChatModelAgent first") + } else { + assert.Contains(t, []string{"outer", "inner"}, firstModel, + "non-recursive cancel does not define whether a root or already-persisted inner checkpoint resumes first") + } + case <-time.After(5 * time.Second): + t.Fatal("no model call observed during resume") + } + + var resumeEvents []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Err != nil { + t.Fatalf("unexpected error during resume: %v", event.Err) + } + resumeEvents = append(resumeEvents, event) + } + + // The outer model should have been called exactly once during resume + // (to produce the final answer after receiving tool results). + // If it was called with the original user message (restarting from scratch), + // the counting model would either exceed its response list or the call count + // would be wrong. + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeOuterModelCallCount), + "outer model should be called exactly once during resume (for final answer after tool results), "+ + "not restarted from the beginning") + + // Verify we got the completion output. + var gotOutput bool + for _, event := range resumeEvents { + content, err := messageOutputContent(event) + require.NoError(t, err) + if content == "outer completed after resume" { + gotOutput = true + } + } + assert.True(t, gotOutput, "should get final output from resumed outer agent") + }) + } +} + +func TestWithCancel_CancelImmediate_RecursiveAgentTool_ResumeDeepestAgentTool(t *testing.T) { + ctx := context.Background() + + leafModel := newBlockingChatModel(schema.AssistantMessage("leaf done", nil)) + t.Cleanup(func() { + close(leafModel.unblockCh) + }) + leafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "LeafAgent", + Description: "leaf agent that blocks", + Instruction: "you are a leaf agent", + Model: leafModel, + }) + require.NoError(t, err) + + middleModelCallCount := int32(0) + middleModel := &countingChatModel{ + callCount: &middleModelCallCount, + responses: []*schema.Message{ + toolCallMsg(toolCall("middle-leaf", "LeafAgent", `{"request":"leaf work"}`)), + schema.AssistantMessage("middle done", nil), + }, + } + middleAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "MiddleAgent", + Description: "middle agent with an agent tool", + Instruction: "you are a middle agent", + Model: middleModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, leafAgent)}, + }, + }, + }) + require.NoError(t, err) + + outerModelCallCount := int32(0) + outerModel := &countingChatModel{ + callCount: &outerModelCallCount, + responses: []*schema.Message{ + toolCallMsg(toolCall("outer-middle", "MiddleAgent", `{"request":"middle work"}`)), + schema.AssistantMessage("outer done", nil), + }, + } + outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "outer agent with recursive agent tool nesting", + Instruction: "you are an outer agent", + Model: outerModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, middleAgent)}, + }, + }, + }) + require.NoError(t, err) + + store := newCancelTestStore() + checkpointID := "cancel-recursive-agent-tool-resume" + runner1 := NewRunner(ctx, RunnerConfig{Agent: outerAgent, CheckPointStore: store}) + + cancelOpt, cancelFn := WithCancel() + iter := runner1.Run(ctx, []Message{schema.UserMessage("go")}, cancelOpt, WithCheckPointID(checkpointID)) + + select { + case <-leafModel.started: + case <-time.After(5 * time.Second): + t.Fatal("leaf model did not start") + } + + handle, _ := cancelFn(WithRecursive()) + require.NoError(t, handle.Wait()) + _, hasCancelError := drainEvents(iter) + assert.True(t, hasCancelError, "expected CancelError from recursive nested agent tool") + + firstModelCall := make(chan string, 8) + resumeLeafModelCallCount := int32(0) + resumeLeafModel := &countingChatModel{ + callCount: &resumeLeafModelCallCount, + callCh: firstModelCall, + callLabel: "leaf", + responses: []*schema.Message{schema.AssistantMessage("leaf done after resume", nil)}, + } + resumeLeafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "LeafAgent", + Description: "leaf agent that returns on resume", + Instruction: "you are a leaf agent", + Model: resumeLeafModel, + }) + require.NoError(t, err) + + resumeMiddleModelCallCount := int32(0) + resumeMiddleModel := &countingChatModel{ + callCount: &resumeMiddleModelCallCount, + callCh: firstModelCall, + callLabel: "middle", + responses: []*schema.Message{schema.AssistantMessage("middle done after resume", nil)}, + } + resumeMiddleAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "MiddleAgent", + Description: "middle agent with an agent tool", + Instruction: "you are a middle agent", + Model: resumeMiddleModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, resumeLeafAgent)}, + }, + }, + }) + require.NoError(t, err) + + resumeOuterModelCallCount := int32(0) + resumeOuterModel := &countingChatModel{ + callCount: &resumeOuterModelCallCount, + callCh: firstModelCall, + callLabel: "outer", + responses: []*schema.Message{schema.AssistantMessage("outer done after resume", nil)}, + } + resumeOuterAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "outer agent with recursive agent tool nesting", + Instruction: "you are an outer agent", + Model: resumeOuterModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, resumeMiddleAgent)}, + }, + }, + }) + require.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{Agent: resumeOuterAgent, CheckPointStore: store}) + resumeIter, err := runner2.Resume(ctx, checkpointID) + require.NoError(t, err) + + select { + case first := <-firstModelCall: + assert.Equal(t, "leaf", first, "recursive AgentTool nesting should resume the deepest internal agent first") + case <-time.After(5 * time.Second): + t.Fatal("no model call observed during resume") + } + + resumeEvents, hasResumeCancelError := drainEvents(resumeIter) + require.False(t, hasResumeCancelError, "resume should complete without another CancelError") + assert.NotEmpty(t, resumeEvents) + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeLeafModelCallCount)) + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeMiddleModelCallCount)) + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeOuterModelCallCount)) +} + +func TestWithCancel_CancelImmediate_ConcurrentAgentTools_ResumeWithoutRestart(t *testing.T) { + ctx := context.Background() + + innerAModel := newBlockingChatModel(schema.AssistantMessage("inner A done", nil)) + innerBModel := newBlockingChatModel(schema.AssistantMessage("inner B done", nil)) + t.Cleanup(func() { + close(innerAModel.unblockCh) + close(innerBModel.unblockCh) + }) + + innerAAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "InnerAgentA", + Description: "inner agent A", + Instruction: "you are inner agent A", + Model: innerAModel, + }) + require.NoError(t, err) + innerBAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "InnerAgentB", + Description: "inner agent B", + Instruction: "you are inner agent B", + Model: innerBModel, + }) + require.NoError(t, err) + + outerModelCallCount := int32(0) + outerModel := &countingChatModel{ + callCount: &outerModelCallCount, + responses: []*schema.Message{ + toolCallMsg( + toolCall("outer-a", "InnerAgentA", `{"request":"work A"}`), + toolCall("outer-b", "InnerAgentB", `{"request":"work B"}`), + ), + schema.AssistantMessage("outer done", nil), + }, + } + outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "outer agent with concurrent agent tools", + Instruction: "you are an outer agent", + Model: outerModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{ + NewAgentTool(ctx, innerAAgent), + NewAgentTool(ctx, innerBAgent), + }, + }, + }, + }) + require.NoError(t, err) + + store := newCancelTestStore() + checkpointID := "cancel-concurrent-agent-tools-resume" + runner1 := NewRunner(ctx, RunnerConfig{Agent: outerAgent, CheckPointStore: store}) + + cancelOpt, cancelFn := WithCancel() + iter := runner1.Run(ctx, []Message{schema.UserMessage("go")}, cancelOpt, WithCheckPointID(checkpointID)) + + for _, started := range []chan struct{}{innerAModel.started, innerBModel.started} { + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("both concurrent inner models should start before cancel") + } + } + + handle, _ := cancelFn(WithRecursive()) + require.NoError(t, handle.Wait()) + _, hasCancelError := drainEvents(iter) + assert.True(t, hasCancelError, "expected CancelError from concurrent agent tools") + + firstModelCall := make(chan string, 8) + resumeInnerAModelCallCount := int32(0) + resumeInnerAModel := &countingChatModel{ + callCount: &resumeInnerAModelCallCount, + callCh: firstModelCall, + callLabel: "innerA", + responses: []*schema.Message{schema.AssistantMessage("inner A done after resume", nil)}, + } + resumeInnerBModelCallCount := int32(0) + resumeInnerBModel := &countingChatModel{ + callCount: &resumeInnerBModelCallCount, + callCh: firstModelCall, + callLabel: "innerB", + responses: []*schema.Message{schema.AssistantMessage("inner B done after resume", nil)}, + } + + resumeInnerAAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "InnerAgentA", + Description: "inner agent A", + Instruction: "you are inner agent A", + Model: resumeInnerAModel, + }) + require.NoError(t, err) + resumeInnerBAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "InnerAgentB", + Description: "inner agent B", + Instruction: "you are inner agent B", + Model: resumeInnerBModel, + }) + require.NoError(t, err) + + resumeOuterModelCallCount := int32(0) + resumeOuterModel := &countingChatModel{ + callCount: &resumeOuterModelCallCount, + callCh: firstModelCall, + callLabel: "outer", + responses: []*schema.Message{schema.AssistantMessage("outer done after resume", nil)}, + } + resumeOuterAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "outer agent with concurrent agent tools", + Instruction: "you are an outer agent", + Model: resumeOuterModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{ + NewAgentTool(ctx, resumeInnerAAgent), + NewAgentTool(ctx, resumeInnerBAgent), + }, + }, + }, + }) + require.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{Agent: resumeOuterAgent, CheckPointStore: store}) + resumeIter, err := runner2.Resume(ctx, checkpointID) + require.NoError(t, err) + + select { + case first := <-firstModelCall: + assert.Contains(t, []string{"innerA", "innerB"}, first, + "concurrent AgentTools should resume an internal agent before the outer model") + case <-time.After(5 * time.Second): + t.Fatal("no model call observed during resume") + } + + resumeEvents, hasResumeCancelError := drainEvents(resumeIter) + require.False(t, hasResumeCancelError, "resume should complete without another CancelError") + assert.NotEmpty(t, resumeEvents) + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeInnerAModelCallCount)) + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeInnerBModelCallCount)) + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeOuterModelCallCount)) +} + +// countingChatModel is a chat model that counts calls and records inputs. +// It returns responses from a fixed slice, indexed by call count. +type countingChatModel struct { + callCount *int32 + inputsCh chan []*schema.Message // optional: receives a copy of each input + callCh chan string // optional: receives callLabel when Generate is called + callLabel string + responses []*schema.Message +} + +func (m *countingChatModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + idx := int(atomic.AddInt32(m.callCount, 1)) - 1 + if m.callCh != nil { + select { + case m.callCh <- m.callLabel: + default: + } + } + if m.inputsCh != nil { + cp := make([]*schema.Message, len(input)) + copy(cp, input) + select { + case m.inputsCh <- cp: + default: + } + } + if idx >= len(m.responses) { + return nil, fmt.Errorf("countingChatModel: call %d exceeds %d responses (outer model was called too many times - possible restart from beginning)", idx+1, len(m.responses)) + } + return m.responses[idx], nil +} + +func (m *countingChatModel) Stream(_ context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(context.Background(), input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *countingChatModel) BindTools(_ []*schema.ToolInfo) error { return nil } + +func messageOutputContent(event *AgentEvent) (string, error) { + if event.Output == nil || event.Output.MessageOutput == nil { + return "", nil + } + mo := event.Output.MessageOutput + if mo.IsStreaming { + msg, err := schema.ConcatMessageStream(mo.MessageStream) + if err != nil { + return "", err + } + if msg == nil { + return "", nil + } + return msg.Content, nil + } + if mo.Message == nil { + return "", nil + } + return mo.Message.Content, nil +} diff --git a/adk/cancel_multicall_test.go b/adk/cancel_multicall_test.go new file mode 100644 index 0000000..790d14f --- /dev/null +++ b/adk/cancel_multicall_test.go @@ -0,0 +1,125 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/compose" +) + +func TestAgentCancelFunc_MultiCall_EscalateToImmediate(t *testing.T) { + cc := newCancelContext() + var interruptCalls int32 + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&interruptCalls, 1) + }) + cancelFn := cc.buildCancelFunc() + + handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + handle2, _ := cancelFn(WithAgentCancelMode(CancelImmediate)) + assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls)) + + cancelErr := cc.createCancelError() + assert.Equal(t, CancelImmediate, cancelErr.Info.Mode) + assert.True(t, cancelErr.Info.Escalated) + assert.False(t, cancelErr.Info.Timeout) + + assert.True(t, cc.markCancelHandled()) + assert.NoError(t, handle1.Wait()) + assert.NoError(t, handle2.Wait()) +} + +func TestAgentCancelFunc_MultiCall_JoinSafePointModes(t *testing.T) { + cc := newCancelContext() + cancelFn := cc.buildCancelFunc() + + handle1, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + handle2, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + + want := CancelAfterChatModel | CancelAfterToolCalls + assert.Equal(t, want, cc.getMode()) + + assert.True(t, cc.markCancelHandled()) + assert.NoError(t, handle1.Wait()) + assert.NoError(t, handle2.Wait()) +} + +func TestAgentCancelFunc_MultiCall_TimeoutDeadlineJoinUsesAbsoluteTime(t *testing.T) { + cc := newCancelContext() + cancelFn := cc.buildCancelFunc() + + handle1, _ := cancelFn( + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(200*time.Millisecond), + ) + + firstDeadline := cc.getDeadlineUnixNano() + assert.NotZero(t, firstDeadline) + + time.Sleep(50 * time.Millisecond) + + handle2, _ := cancelFn( + WithAgentCancelMode(CancelAfterToolCalls), + WithAgentCancelTimeout(60*time.Millisecond), + ) + + secondDeadline := cc.getDeadlineUnixNano() + assert.NotZero(t, secondDeadline) + assert.Less(t, secondDeadline, firstDeadline) + + assert.True(t, cc.markCancelHandled()) + assert.NoError(t, handle1.Wait()) + assert.NoError(t, handle2.Wait()) +} + +func TestAgentCancelFunc_MultiCall_TimeoutEscalationReturnsErrCancelTimeout(t *testing.T) { + cc := newCancelContext() + var interruptCalls int32 + interruptCh := make(chan struct{}, 1) + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&interruptCalls, 1) + select { + case interruptCh <- struct{}{}: + default: + } + }) + cancelFn := cc.buildCancelFunc() + handle, _ := cancelFn( + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(30*time.Millisecond), + ) + + select { + case <-interruptCh: + case <-time.After(1 * time.Second): + t.Fatal("timeout escalation did not interrupt") + } + assert.Equal(t, int32(1), atomic.LoadInt32(&interruptCalls)) + + cancelErr := cc.createCancelError() + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + assert.True(t, cancelErr.Info.Escalated) + assert.True(t, cancelErr.Info.Timeout) + + assert.True(t, cc.markCancelHandled()) + assert.Equal(t, ErrCancelTimeout, handle.Wait()) +} diff --git a/adk/cancel_recursive_test.go b/adk/cancel_recursive_test.go new file mode 100644 index 0000000..cd7e427 --- /dev/null +++ b/adk/cancel_recursive_test.go @@ -0,0 +1,387 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "runtime" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func assertNotClosedWithin(t *testing.T, ch <-chan struct{}, d time.Duration) { + t.Helper() + select { + case <-ch: + t.Fatal("channel was closed but should not have been") + case <-time.After(d): + } +} + +func setupParentChild(t *testing.T) (parent, child *cancelContext, cleanup func()) { + parent = newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + child = parent.deriveAgentToolCancelContext(ctx) + cleanup = func() { + child.markDone() + cancel() + } + t.Cleanup(cleanup) + return parent, child, cleanup +} + +func TestDeriveAgentToolCancelContext(t *testing.T) { + t.Run("Shallow", func(t *testing.T) { + t.Run("DoesNotPropagateSafePoint", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerCancel(CancelAfterChatModel) + + assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) + }) + + t.Run("ImmediateDoesNotPropagate", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerImmediateCancel() + + assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond) + }) + + t.Run("GrandchildNoPropagation", func(t *testing.T) { + a := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + b := a.deriveAgentToolCancelContext(ctx) + c := b.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + c.markDone() + b.markDone() + cancel() + }) + + a.triggerCancel(CancelAfterChatModel) + + assertNotClosedWithin(t, b.cancelChan, 50*time.Millisecond) + assertNotClosedWithin(t, c.cancelChan, 50*time.Millisecond) + }) + + t.Run("NeverRecursive_GoroutineCleanup", func(t *testing.T) { + runtime.GC() + time.Sleep(50 * time.Millisecond) + before := runtime.NumGoroutine() + + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + child := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + time.Sleep(100 * time.Millisecond) + + child.markDone() + cancel() + + time.Sleep(200 * time.Millisecond) + runtime.GC() + time.Sleep(50 * time.Millisecond) + after := runtime.NumGoroutine() + + assert.InDelta(t, before, after, 5, "goroutine leak detected: before=%d after=%d", before, after) + }) + }) + + t.Run("Recursive", func(t *testing.T) { + t.Run("PropagatesSafePoint", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.setRecursive(true) + parent.triggerCancel(CancelAfterChatModel) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive cancel within 1s") + } + assert.True(t, child.shouldCancel()) + }) + + t.Run("ImmediatePropagates", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.setRecursive(true) + parent.triggerImmediateCancel() + + select { + case <-child.immediateChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive immediate cancel within 1s") + } + assert.True(t, child.isImmediateCancelled()) + }) + + t.Run("GrandchildPropagation", func(t *testing.T) { + a := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + b := a.deriveAgentToolCancelContext(ctx) + c := b.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + c.markDone() + b.markDone() + cancel() + }) + + a.setRecursive(true) + a.triggerCancel(CancelAfterChatModel) + + select { + case <-b.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("B did not receive cancel within 1s") + } + + select { + case <-c.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("C did not receive cancel within 1s") + } + + assert.True(t, b.shouldCancel()) + assert.True(t, c.shouldCancel()) + }) + + t.Run("SetBeforeCancel", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.setRecursive(true) + + parent.triggerCancel(CancelAfterChatModel) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive cancel within 1s") + } + assert.True(t, child.shouldCancel()) + }) + + t.Run("AfterRecursiveAndCancelAlreadySet", func(t *testing.T) { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + parent.setRecursive(true) + parent.triggerCancel(CancelAfterChatModel) + + child := parent.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + child.markDone() + cancel() + }) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not immediately receive cancel") + } + assert.True(t, child.shouldCancel()) + }) + }) + + t.Run("Escalation", func(t *testing.T) { + t.Run("EscalateFromNonRecursive", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerCancel(CancelAfterChatModel) + + assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) + + parent.setRecursive(true) + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive cancel after escalation within 1s") + } + assert.True(t, child.shouldCancel()) + }) + + t.Run("EscalateImmediate", func(t *testing.T) { + parent, child, _ := setupParentChild(t) + + parent.triggerImmediateCancel() + + assertNotClosedWithin(t, child.immediateChan, 50*time.Millisecond) + + parent.setRecursive(true) + + select { + case <-child.immediateChan: + case <-time.After(1 * time.Second): + t.Fatal("child did not receive immediate cancel after escalation within 1s") + } + assert.True(t, child.isImmediateCancelled()) + }) + }) +} + +func TestDeriveAgentToolCancelContext_Race(t *testing.T) { + t.Run("SetRecursiveConcurrentWithCancelChan", func(t *testing.T) { + for i := 0; i < 100; i++ { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + child := parent.deriveAgentToolCancelContext(ctx) + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + parent.setRecursive(true) + }() + + go func() { + defer wg.Done() + parent.triggerCancel(CancelAfterChatModel) + }() + + wg.Wait() + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatalf("iteration %d: child did not receive cancel within 1s", i) + } + + assert.True(t, child.shouldCancel()) + child.markDone() + cancel() + } + }) + + t.Run("ChildCompletesBeforeEscalation", func(t *testing.T) { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + child := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + time.Sleep(50 * time.Millisecond) + + child.markDone() + time.Sleep(50 * time.Millisecond) + + parent.setRecursive(true) + + assertNotClosedWithin(t, child.cancelChan, 50*time.Millisecond) + }) + + t.Run("MultipleChildren_PartialCompletion", func(t *testing.T) { + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + child1 := parent.deriveAgentToolCancelContext(ctx) + child2 := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + time.Sleep(50 * time.Millisecond) + + child1.markDone() + time.Sleep(50 * time.Millisecond) + + parent.setRecursive(true) + + select { + case <-child2.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("running child did not receive cancel within 1s") + } + + assert.True(t, child2.shouldCancel()) + assert.False(t, child1.shouldCancel()) + child2.markDone() + }) + + t.Run("ContextCancelConcurrentWithRecursive", func(t *testing.T) { + done := make(chan struct{}) + go func() { + defer close(done) + + parent := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + + child := parent.deriveAgentToolCancelContext(ctx) + + parent.triggerCancel(CancelAfterChatModel) + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + cancel() + }() + + go func() { + defer wg.Done() + parent.setRecursive(true) + }() + + wg.Wait() + child.markDone() + }() + + select { + case <-done: + case <-time.After(1 * time.Second): + t.Fatal("deadlock detected") + } + }) + + t.Run("ConcurrentSetRecursive", func(t *testing.T) { + parent := newCancelContext() + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + parent.setRecursive(true) + }() + } + + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(1 * time.Second): + t.Fatal("deadlock or panic in concurrent setRecursive") + } + + assert.True(t, parent.isRecursive()) + }) +} diff --git a/adk/cancel_test.go b/adk/cancel_test.go new file mode 100644 index 0000000..ea36afe --- /dev/null +++ b/adk/cancel_test.go @@ -0,0 +1,3794 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "io" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type cancelTestChatModel struct { + delayNs int64 + response *schema.Message + startedChan chan struct{} + doneChan chan struct{} +} + +func (m *cancelTestChatModel) getDelay() time.Duration { + return time.Duration(atomic.LoadInt64(&m.delayNs)) +} + +func (m *cancelTestChatModel) setDelay(d time.Duration) { + atomic.StoreInt64(&m.delayNs, int64(d)) +} + +func (m *cancelTestChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + select { + case m.startedChan <- struct{}{}: + default: + } + select { + case <-time.After(m.getDelay()): + case <-ctx.Done(): + return nil, ctx.Err() + } + select { + case m.doneChan <- struct{}{}: + default: + } + return m.response, nil +} + +func (m *cancelTestChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + m.startedChan <- struct{}{} + time.Sleep(m.getDelay()) + m.doneChan <- struct{}{} + return schema.StreamReaderFromArray([]*schema.Message{m.response}), nil +} + +func (m *cancelTestChatModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +type slowTool struct { + name string + delay time.Duration + result string + callCount int32 + startedChan chan struct{} +} + +func newSlowTool(name string, delay time.Duration, result string) *slowTool { + return &slowTool{ + name: name, + delay: delay, + result: result, + startedChan: make(chan struct{}, 10), + } +} + +func (t *slowTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "A slow tool for testing", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string", Desc: "Input parameter"}, + }), + }, nil +} + +func (t *slowTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + atomic.AddInt32(&t.callCount, 1) + select { + case t.startedChan <- struct{}{}: + default: + } + select { + case <-time.After(t.delay): + case <-ctx.Done(): + return "", ctx.Err() + } + return t.result, nil +} + +type cancelTestStore struct { + m map[string][]byte + mu sync.Mutex +} + +func newCancelTestStore() *cancelTestStore { + return &cancelTestStore{m: make(map[string][]byte)} +} + +func (s *cancelTestStore) Set(_ context.Context, key string, value []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.m[key] = value + return nil +} + +func (s *cancelTestStore) Get(_ context.Context, key string) ([]byte, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.m[key] + return v, ok, nil +} + +func assertHasCancelError(t *testing.T, events []*AgentEvent) { + t.Helper() + for _, e := range events { + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + return + } + } + t.Fatal("expected CancelError in events") +} + +func drainAndAssertCancelError(t *testing.T, iter *AsyncIterator[*AgentEvent]) { + t.Helper() + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + return + } + } + t.Fatal("expected CancelError in event stream") +} + +func drainEventsAndAssertCancelError(t *testing.T, iter *AsyncIterator[*AgentEvent]) []*AgentEvent { + t.Helper() + var events []*AgentEvent + hasCancelError := false + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + hasCancelError = true + } + events = append(events, event) + } + assert.True(t, hasCancelError, "expected CancelError in event stream") + return events +} + +func TestCancelContext(t *testing.T) { + t.Run("BasicCancelContext", func(t *testing.T) { + cc := newCancelContext() + assert.False(t, cc.shouldCancel(), "Should not be cancelled initially") + + cc.setMode(CancelImmediate) + close(cc.cancelChan) + + assert.True(t, cc.shouldCancel(), "Should be cancelled after close(cancelChan)") + assert.Equal(t, CancelImmediate, cc.getMode()) + }) +} + +func TestWithCancel_WithTools(t *testing.T) { + ctx := context.Background() + + t.Run("CancelImmediate_DuringModelCall", func(t *testing.T) { + modelStarted := make(chan struct{}, 1) + st := newSlowTool("slow_tool", 100*time.Millisecond, "tool result") + + slowModel := &cancelTestChatModel{ + delayNs: int64(2 * time.Second), + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt) + assert.NotNil(t, iter) + assert.NotNil(t, cancelFn) + + eventsCh := make(chan []*AgentEvent, 1) + go func() { + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + eventsCh <- events + }() + + select { + case <-modelStarted: + case <-time.After(5 * time.Second): + t.Fatal("Model did not start within 5 seconds") + } + + time.Sleep(100 * time.Millisecond) + + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err) + + var events []*AgentEvent + select { + case events = <-eventsCh: + case <-time.After(5 * time.Second): + t.Fatal("Timed out waiting for events") + } + + assert.NotEmpty(t, events) + + assertHasCancelError(t, events) + }) + + t.Run("CancelAfterChatModel_DuringToolCall", func(t *testing.T) { + toolStarted := make(chan struct{}, 1) + st := &slowToolWithSignal{ + name: "slow_tool", + delay: 2 * time.Second, + result: "tool result", + startedChan: toolStarted, + } + + modelWithToolCall := &simpleChatModel{ + delay: 1 * time.Second, + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: modelWithToolCall, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("Use the tool")}, + }, cancelOpt) + assert.NotNil(t, iter) + assert.NotNil(t, cancelFn) + + <-toolStarted + + time.Sleep(100 * time.Millisecond) + + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + err = handle.Wait() + assert.NoError(t, err) + + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + continue + } + events = append(events, event) + } + + assert.NotEmpty(t, events) + assert.True(t, atomic.LoadInt32(&st.callCount) >= 1, "Tool should have been called") + }) + + t.Run("CancelAfterToolCalls_CompletesToolExecution", func(t *testing.T) { + toolStarted := make(chan struct{}, 1) + st := &slowToolWithSignal{ + name: "slow_tool", + delay: 500 * time.Millisecond, + result: "tool result", + startedChan: toolStarted, + } + + modelWithToolCall := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: modelWithToolCall, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("Use the tool")}, + }, cancelOpt) + assert.NotNil(t, iter) + assert.NotNil(t, cancelFn) + + <-toolStarted + + time.Sleep(100 * time.Millisecond) + + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + err = handle.Wait() + assert.NoError(t, err) + + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + continue + } + events = append(events, event) + } + + assert.NotEmpty(t, events) + assert.True(t, atomic.LoadInt32(&st.callCount) >= 1, "Tool should have been called") + }) + + t.Run("NestedCancelPropagation", func(t *testing.T) { + cc := newCancelContext() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + child := cc.deriveAgentToolCancelContext(ctx) + assert.NotNil(t, child) + + cc.setRecursive(true) + cc.setMode(CancelImmediate) + + if atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) { + close(cc.cancelChan) + } + + select { + case <-child.cancelChan: + case <-time.After(1 * time.Second): + t.Fatal("Child did not receive cancel signal") + } + + assert.True(t, child.shouldCancel()) + assert.Equal(t, CancelImmediate, child.getMode()) + }) + + t.Run("DeepAgentIntegrationCancel", func(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + + leafModel := &cancelTestChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "Leaf result", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + leafModel.setDelay(500 * time.Millisecond) + leafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "LeafAgent", + Description: "desc", + Model: leafModel, + }) + assert.NoError(t, err) + + rootModel := &cancelTestChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "LeafAgent", + Arguments: `{}`, + }, + }, + }, + }, + startedChan: make(chan struct{}, 1), + doneChan: make(chan struct{}, 1), + } + rootAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RootAgent", + Description: "desc", + Model: rootModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, leafAgent)}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: rootAgent, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Run leaf")}, cancelOpt) + + <-modelStarted + + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterChatModel), WithRecursive()) + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + var ce *CancelError + if errors.As(event.Err, &ce) { + hasCancelError = true + assert.NotNil(t, ce.interruptSignal, "CancelError should carry interrupt signal") + } + } + } + assert.True(t, hasCancelError, "Should have received CancelError") + }) +} + +type slowToolWithSignal struct { + name string + delay time.Duration + result string + callCount int32 + startedChan chan struct{} +} + +func (t *slowToolWithSignal) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "A slow tool for testing", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string", Desc: "Input parameter"}, + }), + }, nil +} + +func (t *slowToolWithSignal) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + atomic.AddInt32(&t.callCount, 1) + t.startedChan <- struct{}{} + time.Sleep(t.delay) + return t.result, nil +} + +type simpleChatModel struct { + delay time.Duration + response *schema.Message +} + +func (m *simpleChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + if m.delay > 0 { + select { + case <-time.After(m.delay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + return m.response, nil +} + +func (m *simpleChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + if m.delay > 0 { + select { + case <-time.After(m.delay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + return schema.StreamReaderFromArray([]*schema.Message{m.response}), nil +} + +func (m *simpleChatModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +func TestWithCancel_WithCheckpoint(t *testing.T) { + ctx := context.Background() + + t.Run("CancelWithCheckpoint", func(t *testing.T) { + modelStarted := make(chan struct{}, 1) + st := newSlowTool("slow_tool", 100*time.Millisecond, "tool result") + + slowModel := &cancelTestChatModel{ + delayNs: int64(1 * time.Second), + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt, WithCheckPointID("cancel-1")) + + <-modelStarted + + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + hasCancelError = true + continue + } + } + + assert.True(t, hasCancelError, "Should have CancelError event after cancel") + }) +} + +func TestAgentCancelFuncMultipleCalls(t *testing.T) { + ctx := context.Background() + + t.Run("SecondCancelReturnsErrAgentFinished", func(t *testing.T) { + modelStarted := make(chan struct{}, 1) + st := newSlowTool("slow_tool", 100*time.Millisecond, "tool result") + + slowModel := &cancelTestChatModel{ + delayNs: int64(1 * time.Second), + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt) + + <-modelStarted + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + }) +} + +func TestWithCancel_Streaming(t *testing.T) { + ctx := context.Background() + + t.Run("CancelImmediate_DuringModelStream", func(t *testing.T) { + modelStarted := make(chan struct{}, 1) + st := newSlowTool("slow_tool", 100*time.Millisecond, "tool result") + + slowModel := &cancelTestChatModel{ + delayNs: int64(2 * time.Second), + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt) + assert.NotNil(t, iter) + assert.NotNil(t, cancelFn) + + eventsCh := make(chan []*AgentEvent, 1) + go func() { + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + eventsCh <- events + }() + + select { + case <-modelStarted: + case <-time.After(5 * time.Second): + t.Fatal("Model did not start within 5 seconds") + } + + time.Sleep(100 * time.Millisecond) + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + var events []*AgentEvent + select { + case events = <-eventsCh: + case <-time.After(5 * time.Second): + t.Fatal("Timed out waiting for events") + } + + assert.NotEmpty(t, events) + + assertHasCancelError(t, events) + }) + + t.Run("CancelAfterToolCalls_Streaming", func(t *testing.T) { + toolStarted := make(chan struct{}, 1) + st := &slowToolWithSignal{ + name: "slow_tool", + delay: 500 * time.Millisecond, + result: "tool result", + startedChan: toolStarted, + } + + modelWithToolCall := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: modelWithToolCall, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt) + assert.NotNil(t, iter) + assert.NotNil(t, cancelFn) + + <-toolStarted + + time.Sleep(100 * time.Millisecond) + + handle, _ := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + continue + } + events = append(events, event) + } + + assert.NotEmpty(t, events) + assert.True(t, atomic.LoadInt32(&st.callCount) >= 1, "Tool should have been called") + }) +} + +// TestWithCancel_Resume tests the workflow of Cancel followed by Resume. +// +// To avoid data races, we create new agent and runner instances for the Resume phase +// instead of reusing and modifying the original model instance. +func TestWithCancel_Resume(t *testing.T) { + ctx := context.Background() + + t.Run("Cancel_ThenResume", func(t *testing.T) { + modelStarted := make(chan struct{}, 1) + modelCallCount := int32(0) + st := newSlowTool("slow_tool", 100*time.Millisecond, "tool result") + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + checkpointID := "resume-cancel-test-1" + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt, WithCheckPointID(checkpointID)) + + <-modelStarted + atomic.AddInt32(&modelCallCount, 1) + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + hasCancelErr := false + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + var ce *CancelError + if errors.As(event.Err, &ce) { + hasCancelErr = true + continue + } + t.Fatalf("unexpected error: %v", event.Err) + } + } + assert.True(t, hasCancelErr, "Should have CancelError event after cancel") + + newModelStarted := make(chan struct{}, 1) + slowModel2 := &cancelTestChatModel{ + delayNs: int64(100 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Final response after resume", + }, + startedChan: newModelStarted, + doneChan: make(chan struct{}, 1), + } + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel2, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: agent2, + EnableStreaming: false, + CheckPointStore: store, + }) + + resumeCancelOpt, _ := WithCancel() + resumeIter, err := runner2.Resume(ctx, checkpointID, resumeCancelOpt) + assert.NoError(t, err) + assert.NotNil(t, resumeIter) + + var resumeEvents []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + assert.Nil(t, event.Err, "Should not have error event during resume") + resumeEvents = append(resumeEvents, event) + } + + assert.NotEmpty(t, resumeEvents, "Resume should produce events") + }) + + t.Run("Resume_ThenCancel", func(t *testing.T) { + firstModelStarted := make(chan struct{}, 1) + modelCallCount := int32(0) + st := newSlowTool("slow_tool", 100*time.Millisecond, "tool result") + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + startedChan: firstModelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + checkpointID := "resume-then-cancel-test-1" + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt, WithCheckPointID(checkpointID)) + + <-firstModelStarted + atomic.AddInt32(&modelCallCount, 1) + + handle, _ := cancelFn() + cancelErr := handle.Wait() + assert.NoError(t, cancelErr) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + slowModel2 := newBlockingChatModel(toolCallMsg(toolCall("call_1", "slow_tool", `{"input": "test"}`))) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tool", + Instruction: "You are a test assistant", + Model: slowModel2, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: agent2, + EnableStreaming: false, + CheckPointStore: store, + }) + + resumeCancelOpt, resumeCancelFn := WithCancel() + resumeIter, err := runner2.Resume(ctx, checkpointID, resumeCancelOpt) + assert.NoError(t, err) + + resumeEventsCh := make(chan []*AgentEvent, 1) + go func() { + var events []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + events = append(events, event) + } + resumeEventsCh <- events + }() + + <-slowModel2.started + atomic.AddInt32(&modelCallCount, 1) + + cancelHandle, _ := resumeCancelFn() + close(slowModel2.unblockCh) + err = cancelHandle.Wait() + assert.True(t, err == nil || errors.Is(err, ErrExecutionEnded), "unexpected cancel wait error: %v", err) + + start := time.Now() + resumeEvents := <-resumeEventsCh + elapsed := time.Since(start) + + assert.True(t, elapsed < 1*time.Second, "Resume should return quickly after cancel, elapsed: %v", elapsed) + assert.NotEmpty(t, resumeEvents) + + hasCancelError := false + for _, e := range resumeEvents { + var ce *CancelError + if e.Err != nil && errors.As(e.Err, &ce) { + hasCancelError = true + } + } + executionCompletedBeforeCancel := errors.Is(err, ErrExecutionEnded) + assert.True(t, hasCancelError || executionCompletedBeforeCancel, "Resume should have CancelError event after cancel, or execution completed before cancel") + }) +} + +func TestCancelMonitoredToolHandler_StreamableToolCall(t *testing.T) { + t.Run("NoCancelContext_PassesThrough", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + + // Create a stream with some data + r, w := schema.Pipe[string](1) + go func() { + w.Send("chunk1", nil) + w.Send("chunk2", nil) + w.Close() + }() + + next := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return &compose.StreamToolOutput{Result: r}, nil + } + + wrapped := handler.WrapStreamableToolCall(next) + // No cancelContext in the Go context + output, err := wrapped(context.Background(), &compose.ToolInput{Name: "test"}) + assert.NoError(t, err) + + // Should get the original stream unchanged + chunk1, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, "chunk1", chunk1) + + chunk2, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, "chunk2", chunk2) + + _, err = output.Result.Recv() + assert.ErrorIs(t, err, io.EOF) + }) + + t.Run("WithCancelContext_NoCancel_StreamsNormally", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + cc := newCancelContext() + + r, w := schema.Pipe[string](1) + go func() { + w.Send("data1", nil) + w.Send("data2", nil) + w.Close() + }() + + next := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return &compose.StreamToolOutput{Result: r}, nil + } + + wrapped := handler.WrapStreamableToolCall(next) + ctx := withCancelContext(context.Background(), cc) + output, err := wrapped(ctx, &compose.ToolInput{Name: "test"}) + assert.NoError(t, err) + + chunk1, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, "data1", chunk1) + + chunk2, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, "data2", chunk2) + + _, err = output.Result.Recv() + assert.ErrorIs(t, err, io.EOF) + }) + + t.Run("WithCancelContext_ImmediateCancel_TerminatesStream", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + cc := newCancelContext() + + // Create a slow stream that we'll cancel mid-way + r, w := schema.Pipe[string](1) + go func() { + defer w.Close() + w.Send("chunk1", nil) + time.Sleep(200 * time.Millisecond) + w.Send("chunk2", nil) + }() + + next := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return &compose.StreamToolOutput{Result: r}, nil + } + + wrapped := handler.WrapStreamableToolCall(next) + ctx := withCancelContext(context.Background(), cc) + output, err := wrapped(ctx, &compose.ToolInput{Name: "test"}) + assert.NoError(t, err) + + // Read first chunk + chunk1, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, "chunk1", chunk1) + + // Fire immediate cancel + close(cc.immediateChan) + + // Next recv should get ErrStreamCanceled + _, err = output.Result.Recv() + assert.ErrorIs(t, err, ErrStreamCanceled) + }) + + t.Run("WithCancelContext_AlreadyCancelled_TerminatesImmediately", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + cc := newCancelContext() + close(cc.immediateChan) // Already canceled + + r, w := schema.Pipe[string](1) + go func() { + w.Send("should-not-see", nil) + w.Close() + }() + + next := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return &compose.StreamToolOutput{Result: r}, nil + } + + wrapped := handler.WrapStreamableToolCall(next) + ctx := withCancelContext(context.Background(), cc) + output, err := wrapped(ctx, &compose.ToolInput{Name: "test"}) + assert.NoError(t, err) + + _, err = output.Result.Recv() + assert.ErrorIs(t, err, ErrStreamCanceled) + }) + + t.Run("NextReturnsError_PropagatesError", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + cc := newCancelContext() + + nextErr := errors.New("tool execution failed") + next := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return nil, nextErr + } + + wrapped := handler.WrapStreamableToolCall(next) + ctx := withCancelContext(context.Background(), cc) + _, err := wrapped(ctx, &compose.ToolInput{Name: "test"}) + assert.ErrorIs(t, err, nextErr) + }) +} + +func TestCancelMonitoredToolHandler_EnhancedStreamableToolCall(t *testing.T) { + t.Run("NoCancelContext_PassesThrough", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + + tr1 := &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "chunk1"}}} + r, w := schema.Pipe[*schema.ToolResult](1) + go func() { + w.Send(tr1, nil) + w.Close() + }() + + next := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{Result: r}, nil + } + + wrapped := handler.WrapEnhancedStreamableToolCall(next) + output, err := wrapped(context.Background(), &compose.ToolInput{Name: "test"}) + assert.NoError(t, err) + + result, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, tr1, result) + + _, err = output.Result.Recv() + assert.ErrorIs(t, err, io.EOF) + }) + + t.Run("WithCancelContext_ImmediateCancel_TerminatesStream", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + cc := newCancelContext() + + tr1 := &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "chunk1"}}} + tr2 := &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "chunk2"}}} + r, w := schema.Pipe[*schema.ToolResult](1) + go func() { + defer w.Close() + w.Send(tr1, nil) + time.Sleep(200 * time.Millisecond) + w.Send(tr2, nil) + }() + + next := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{Result: r}, nil + } + + wrapped := handler.WrapEnhancedStreamableToolCall(next) + ctx := withCancelContext(context.Background(), cc) + output, err := wrapped(ctx, &compose.ToolInput{Name: "test"}) + assert.NoError(t, err) + + result, err := output.Result.Recv() + assert.NoError(t, err) + assert.Equal(t, tr1, result) + + close(cc.immediateChan) + + _, err = output.Result.Recv() + assert.ErrorIs(t, err, ErrStreamCanceled) + }) + + t.Run("NextReturnsError_PropagatesError", func(t *testing.T) { + handler := &cancelMonitoredToolHandler{} + cc := newCancelContext() + + nextErr := errors.New("enhanced tool failed") + next := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return nil, nextErr + } + + wrapped := handler.WrapEnhancedStreamableToolCall(next) + ctx := withCancelContext(context.Background(), cc) + _, err := wrapped(ctx, &compose.ToolInput{Name: "test"}) + assert.ErrorIs(t, err, nextErr) + }) +} + +func TestCancelContextKey(t *testing.T) { + t.Run("WithAndGet_RoundTrips", func(t *testing.T) { + cc := newCancelContext() + ctx := withCancelContext(context.Background(), cc) + got := getCancelContext(ctx) + assert.Equal(t, cc, got) + }) + + t.Run("Get_NoValue_ReturnsNil", func(t *testing.T) { + got := getCancelContext(context.Background()) + assert.Nil(t, got) + }) + + t.Run("With_NilCancelContext_ReturnsOriginalCtx", func(t *testing.T) { + ctx := context.Background() + result := withCancelContext(ctx, nil) + assert.Equal(t, ctx, result) + }) +} + +// -- Tests for cancel support across all agent types -- + +// cancelTestAgent is a ChatModelAgent-based agent where the model blocks until +// signalled, allowing tests to control exactly when to issue a cancel. +func newCancelTestAgent(t *testing.T, name string, modelDelay time.Duration, modelStarted chan struct{}) *ChatModelAgent { + t.Helper() + slowModel := &cancelTestChatModel{ + delayNs: int64(modelDelay), + response: &schema.Message{ + Role: schema.Assistant, + Content: "response from " + name, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(context.Background(), &ChatModelAgentConfig{ + Name: name, + Description: "Test agent " + name, + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + return agent +} + +func newCancelTestAgentWithTools(t *testing.T, name string, modelDelay time.Duration, modelStarted chan struct{}) *ChatModelAgent { + t.Helper() + toolName := name + "_tool" + slowModel := &cancelTestChatModel{ + delayNs: int64(modelDelay), + response: &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call_1", Type: "function", + Function: schema.FunctionCall{ + Name: toolName, + Arguments: `{"input": "test"}`, + }, + }}, + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + st := newSlowTool(toolName, 10*time.Millisecond, "tool result") + + agent, err := NewChatModelAgent(context.Background(), &ChatModelAgentConfig{ + Name: name, + Description: "Test agent " + name, + Instruction: "You are a test assistant", + Model: slowModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + return agent +} + +func newCancelTestAgentWithToolsFinalAnswer(t *testing.T, name string) *ChatModelAgent { + t.Helper() + toolName := name + "_tool" + finalModel := &cancelTestChatModel{ + delayNs: int64(10 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "final response from " + name, + }, + startedChan: make(chan struct{}, 1), + doneChan: make(chan struct{}, 1), + } + + st := newSlowTool(toolName, 10*time.Millisecond, "tool result") + + agent, err := NewChatModelAgent(context.Background(), &ChatModelAgentConfig{ + Name: name, + Description: "Test agent " + name, + Instruction: "You are a test assistant", + Model: finalModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + return agent +} + +func TestWithCancel_SequentialAgent(t *testing.T) { + ctx := context.Background() + + t.Run("CancelImmediate_DuringSecondAgent", func(t *testing.T) { + // The first agent completes quickly. The second agent takes a long time. + // Cancel during the second agent's model call. + agent1Started := make(chan struct{}, 1) + agent2Started := make(chan struct{}, 1) + + agent1 := newCancelTestAgent(t, "fast_agent", 50*time.Millisecond, agent1Started) + agent2 := newCancelTestAgent(t, "slow_agent", 5*time.Second, agent2Started) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq_agent", + Description: "Sequential test", + SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + // Wait for second agent to start + select { + case <-agent2Started: + case <-time.After(10 * time.Second): + t.Fatal("Second agent did not start") + } + + time.Sleep(50 * time.Millisecond) + + // Cancel should NOT return ErrExecutionEnded (the bug before the fix) + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err, "Cancel during second agent should succeed, not return ErrExecutionEnded") + + drainEventsAndAssertCancelError(t, iter) + }) +} + +func TestWithCancel_LoopAgent(t *testing.T) { + ctx := context.Background() + + t.Run("CancelImmediate_DuringIteration", func(t *testing.T) { + // Agent in a loop. Cancel during second iteration's model call. + modelStarted := make(chan struct{}, 10) + + slowModel := &cancelTestChatModel{ + delayNs: int64(3 * time.Second), + response: &schema.Message{ + Role: schema.Assistant, + Content: "loop response", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 10), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "loop_inner", + Description: "Inner loop agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop_agent", + Description: "Loop test", + SubAgents: []Agent{agent}, + MaxIterations: 10, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: loopAgent, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + // Wait for first iteration's model call to start + select { + case <-modelStarted: + case <-time.After(10 * time.Second): + t.Fatal("Model did not start") + } + + time.Sleep(50 * time.Millisecond) + + // Cancel should succeed + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err, "Cancel during loop iteration should succeed") + + drainAndAssertCancelError(t, iter) + }) +} + +func TestWithCancel_ParallelAgent(t *testing.T) { + ctx := context.Background() + + t.Run("CancelImmediate_InterruptsAllBranches", func(t *testing.T) { + agent1Started := make(chan struct{}, 1) + agent2Started := make(chan struct{}, 1) + + // Both agents have long delays, so cancel should interrupt both. + agent1 := newCancelTestAgent(t, "par_agent1", 5*time.Second, agent1Started) + agent2 := newCancelTestAgent(t, "par_agent2", 5*time.Second, agent2Started) + + parAgent, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "par_agent", + Description: "Parallel test", + SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: parAgent, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + // Wait for both agents to start + for i := 0; i < 2; i++ { + select { + case <-agent1Started: + case <-agent2Started: + case <-time.After(10 * time.Second): + t.Fatal("Parallel agents did not start") + } + } + + time.Sleep(50 * time.Millisecond) + + start := time.Now() + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err, "Cancel during parallel agents should succeed") + + events := drainEventsAndAssertCancelError(t, iter) + elapsed := time.Since(start) + + _ = events + assert.True(t, elapsed < 3*time.Second, "Should complete quickly after cancel, elapsed: %v", elapsed) + }) +} + +func TestWithCancel_SupervisorAgent(t *testing.T) { + ctx := context.Background() + + t.Run("CancelImmediate_DuringSubAgent", func(t *testing.T) { + // Supervisor delegates to a slow sub-agent via transfer. + // Cancel during the sub-agent's model call. + supervisorModelStarted := make(chan struct{}, 1) + subAgentModelStarted := make(chan struct{}, 1) + + // The supervisor model returns a transfer_to_agent tool call + supervisorModel := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "slow_sub"}`, + }, + }, + }, + }, + } + + supervisorAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "supervisor", + Description: "Supervisor agent", + Instruction: "You are a supervisor", + Model: supervisorModel, + }) + assert.NoError(t, err) + + subAgent := newCancelTestAgent(t, "slow_sub", 5*time.Second, subAgentModelStarted) + + agentWithSubAgents, err := SetSubAgents(ctx, supervisorAgent, []Agent{subAgent}) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agentWithSubAgents, + EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + // Ignore the supervisor model start, wait for the sub-agent model + // The supervisor model is fast (simpleChatModel), so it will start and finish quickly + _ = supervisorModelStarted + select { + case <-subAgentModelStarted: + case <-time.After(10 * time.Second): + t.Fatal("Sub-agent model did not start") + } + + time.Sleep(50 * time.Millisecond) + + start := time.Now() + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err, "Cancel during sub-agent should succeed") + + drainAndAssertCancelError(t, iter) + elapsed := time.Since(start) + + assert.True(t, elapsed < 3*time.Second, "Should complete quickly after cancel, elapsed: %v", elapsed) + }) +} + +func TestFilterCancelOption(t *testing.T) { + t.Run("RemovesCancelOption", func(t *testing.T) { + cancelOpt, _ := WithCancel() + sessionOpt := WithSessionValues(map[string]any{"key": "value"}) + opts := []AgentRunOption{cancelOpt, sessionOpt} + + filtered := filterCancelOption(opts) + assert.Len(t, filtered, 1, "Should have removed the cancel option") + + // Verify the remaining option is the session option + testOpt := &options{} + filtered[0].implSpecificOptFn.(func(*options))(testOpt) + assert.NotNil(t, testOpt.sessionValues) + assert.Nil(t, testOpt.cancelCtx) + }) + + t.Run("KeepsNonCancelOptions", func(t *testing.T) { + sessionOpt := WithSessionValues(map[string]any{"key": "value"}) + callbackOpt := WithCallbacks() + opts := []AgentRunOption{sessionOpt, callbackOpt} + + filtered := filterCancelOption(opts) + assert.Len(t, filtered, 2, "Should keep all non-cancel options") + }) + + t.Run("EmptyInput", func(t *testing.T) { + filtered := filterCancelOption(nil) + assert.Nil(t, filtered) + }) +} + +func wrapIterWithMarkDone(iter *AsyncIterator[*AgentEvent], cc *cancelContext) *AsyncIterator[*AgentEvent] { + if cc == nil { + return iter + } + outIter, outGen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer cc.markDone() + defer outGen.Close() + for { + event, ok := iter.Next() + if !ok { + return + } + outGen.Send(event) + } + }() + return outIter +} + +func TestWrapIterWithMarkDone(t *testing.T) { + t.Run("MarksDoneAfterDrain", func(t *testing.T) { + cc := newCancelContext() + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + gen.Send(&AgentEvent{AgentName: "test"}) + gen.Close() + }() + + wrapped := wrapIterWithMarkDone(iter, cc) + + event, ok := wrapped.Next() + assert.True(t, ok) + assert.Equal(t, "test", event.AgentName) + + _, ok = wrapped.Next() + assert.False(t, ok) + + // markDone should have been called, so doneChan should be closed + select { + case <-cc.doneChan: + // good + case <-time.After(time.Second): + t.Fatal("doneChan was not closed after drain") + } + }) + + t.Run("NilCancelContext_PassesThrough", func(t *testing.T) { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + gen.Send(&AgentEvent{AgentName: "test"}) + gen.Close() + }() + + wrapped := wrapIterWithMarkDone(iter, nil) + assert.Equal(t, iter, wrapped, "Should return same iter when cc is nil") + }) +} + +func TestGraphInterruptFuncs_Parallel(t *testing.T) { + t.Run("MultipleGraphInterruptFuncsAllCalled", func(t *testing.T) { + cc := newCancelContext() + + var called1, called2 int32 + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&called1, 1) + }) + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&called2, 1) + }) + + // Simulate immediate cancel + cc.setMode(CancelImmediate) + atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) + close(cc.cancelChan) + cc.sendImmediateInterrupt() + + assert.Equal(t, int32(1), atomic.LoadInt32(&called1), "First graph interrupt func should be called") + assert.Equal(t, int32(1), atomic.LoadInt32(&called2), "Second graph interrupt func should be called") + }) + + t.Run("RetroactiveFire_OnSetAfterCancel", func(t *testing.T) { + cc := newCancelContext() + + // First set up cancel state with immediate interrupt + cc.setMode(CancelImmediate) + atomic.CompareAndSwapInt32(&cc.state, stateRunning, stateCancelling) + close(cc.cancelChan) + close(cc.immediateChan) + atomic.StoreInt32(&cc.interruptSent, interruptImmediate) + + // Now register a new function - it should be retroactively fired + var called int32 + cc.setGraphInterruptFunc(func(opts ...compose.GraphInterruptOption) { + atomic.AddInt32(&called, 1) + }) + + assert.Equal(t, int32(1), atomic.LoadInt32(&called), "setGraphInterruptFunc should retroactively fire new func") + }) +} + +// -- Tests for transition-point cancel (cancel between sub-agents) -- + +// gatedChatModel is a model that: +// - Signals doneChan when Generate completes +// - Optionally blocks on gateChan before returning (nil gateChan = no blocking) +// - Tracks call count via callCount +type gatedChatModel struct { + response *schema.Message + gateChan chan struct{} // if non-nil, blocks until closed before returning + doneChan chan struct{} // signalled after Generate completes + callCount int32 +} + +func (m *gatedChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m.callCount, 1) + if m.gateChan != nil { + select { + case <-m.gateChan: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + select { + case m.doneChan <- struct{}{}: + default: + } + return m.response, nil +} + +func (m *gatedChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *gatedChatModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +func TestCheckCancel_Sequential_BetweenSubAgents(t *testing.T) { + ctx := context.Background() + + // CancelAfterToolCalls fires at transition boundaries between sub-agents. + // At a transition boundary, the completed sub-agent's entire execution + // (including any tool calls) is done, satisfying the CancelAfterToolCalls + // contract — even if this particular sub-agent had no tools. + model1 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent1 done"}, + gateChan: make(chan struct{}), + doneChan: make(chan struct{}, 1), + } + model2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent2 done"}, + doneChan: make(chan struct{}, 1), + } + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", Model: model1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", Model: model2, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "sequential test", SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt) + + for atomic.LoadInt32(&model1.callCount) == 0 { + runtime.Gosched() + } + + cancelCalled, result := cancelAsync(cancelFn, WithAgentCancelMode(CancelAfterToolCalls)) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(model1.gateChan) + + assert.NoError(t, result.waitDone(t), "CancelAfterToolCalls should succeed at transition boundary") + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, int32(1), atomic.LoadInt32(&model1.callCount), "Agent1 model should be invoked") + assert.Equal(t, int32(0), atomic.LoadInt32(&model2.callCount), + "Agent2 model should NOT be invoked (CancelAfterToolCalls caught at transition)") +} + +func TestCheckCancel_Loop_BetweenIterations(t *testing.T) { + ctx := context.Background() + + // CancelAfterToolCalls fires at loop iteration boundaries. + // After the first iteration completes, any tool calls it made are done, + // satisfying the CancelAfterToolCalls contract. + mdl := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "loop iter"}, + gateChan: make(chan struct{}), + doneChan: make(chan struct{}, 10), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "loop_inner", Description: "inner", Instruction: "test", Model: mdl, + }) + assert.NoError(t, err) + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop", Description: "loop test", SubAgents: []Agent{agent}, MaxIterations: 3, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: loopAgent, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt) + + for atomic.LoadInt32(&mdl.callCount) == 0 { + runtime.Gosched() + } + + cancelCalled, result := cancelAsync(cancelFn, WithAgentCancelMode(CancelAfterToolCalls)) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(mdl.gateChan) + + assert.NoError(t, result.waitDone(t), "CancelAfterToolCalls should succeed at loop transition boundary") + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, int32(1), atomic.LoadInt32(&mdl.callCount), + "Model should be called once; second iteration caught at transition") +} + +func TestCheckCancel_Parallel_PreSpawn(t *testing.T) { + ctx := context.Background() + + // Cancel fires before Run is called. Neither model should be invoked. + model1 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "par1"}, + doneChan: make(chan struct{}, 1), + } + model2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "par2"}, + doneChan: make(chan struct{}, 1), + } + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "par1", Description: "first", Instruction: "test", Model: model1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "par2", Description: "second", Instruction: "test", Model: model2, + }) + assert.NoError(t, err) + + parAgent, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "par", Description: "parallel test", SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + // Fire cancel in goroutine (cancelFn blocks until handled) + cancelOpt, cancelFn := WithCancel() + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn() + cancelDone <- handle.Wait() + }() + // Wait for cancelChan to be closed (happens synchronously before the blocking doneChan wait) + time.Sleep(20 * time.Millisecond) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: parAgent, EnableStreaming: false, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt) + + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + cancelErr = ce + } + } + + // cancelFn should have completed + select { + case err = <-cancelDone: + assert.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("cancelFn did not return") + } + + assert.NotNil(t, cancelErr, "Should have CancelError") + assert.Equal(t, int32(0), atomic.LoadInt32(&model1.callCount), "First model should never be invoked") + assert.Equal(t, int32(0), atomic.LoadInt32(&model2.callCount), "Second model should never be invoked") +} + +func TestCheckCancel_Transfer_BeforeTarget(t *testing.T) { + ctx := context.Background() + + // Supervisor CMA returns a transfer action (instantly). + // Cancel fires after transfer action but before target runs. + // Target model should never be invoked. + supervisorModel := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call_1", Type: "function", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "target"}`, + }, + }}, + }, + } + targetModel := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "target done"}, + doneChan: make(chan struct{}, 1), + } + + supervisorAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "supervisor", Description: "supervisor", Instruction: "test", Model: supervisorModel, + }) + assert.NoError(t, err) + + targetAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "target", Description: "target", Instruction: "test", Model: targetModel, + }) + assert.NoError(t, err) + + agentWithSub, err := SetSubAgents(ctx, supervisorAgent, []Agent{targetAgent}) + assert.NoError(t, err) + + // Fire cancel in goroutine (cancelFn blocks until handled) + cancelOpt, cancelFn := WithCancel() + cancelDone := make(chan error, 1) + go func() { + handle, _ := cancelFn() + cancelDone <- handle.Wait() + }() + time.Sleep(20 * time.Millisecond) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agentWithSub, EnableStreaming: false, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + cancelErr = ce + } + } + + select { + case err = <-cancelDone: + assert.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("cancelFn did not return") + } + + assert.NotNil(t, cancelErr, "Should have CancelError") + assert.Equal(t, int32(0), atomic.LoadInt32(&targetModel.callCount), "Target model should never be invoked") +} + +func TestCheckCancel_AlreadyHandled_NoDuplicate(t *testing.T) { + ctx := context.Background() + + // In a sequential agent, if the first CMA handles the cancel (graph interrupt), + // the workflow's transition check should NOT emit a duplicate CancelError. + // Use a slow model so cancel fires during its execution (handled by CMA). + modelStarted := make(chan struct{}, 1) + model1 := &cancelTestChatModel{ + delayNs: int64(2 * time.Second), + response: &schema.Message{Role: schema.Assistant, Content: "agent1"}, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + model2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent2"}, + doneChan: make(chan struct{}, 1), + } + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", Model: model1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", Model: model2, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "sequential", SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + // Wait for model to start, then cancel during model execution + select { + case <-modelStarted: + case <-time.After(5 * time.Second): + t.Fatal("Model did not start") + } + time.Sleep(50 * time.Millisecond) + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err) + + cancelCount := 0 + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + cancelCount++ + } + } + + assert.Equal(t, 1, cancelCount, "Should have exactly one CancelError, no duplicate from workflow transition") + assert.Equal(t, int32(0), atomic.LoadInt32(&model2.callCount), "Second agent should not run") +} + +// Tests for CancelAfterChatModel/CancelAfterToolCalls in nested workflow structures. +// These verify that safe-point cancel modes propagate through the entire agent hierarchy +// and fire at whichever nested level reaches the safe-point first. + +func TestCancel_SequentialWorkflow_CancelAfterChatModel(t *testing.T) { + ctx := context.Background() + agent1Started := make(chan struct{}, 1) + + agent1 := newCancelTestAgentWithTools(t, "seq_slow", 500*time.Millisecond, agent1Started) + agent2 := newCancelTestAgentWithTools(t, "seq_fast", 50*time.Millisecond, make(chan struct{}, 1)) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq_agent", + Description: "Sequential workflow", + SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("seq-cancel-1")) + + select { + case <-agent1Started: + case <-time.After(10 * time.Second): + t.Fatal("First agent did not start") + } + + handle, contributed := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + assert.True(t, contributed, "Cancel should contribute") + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil && errors.As(event.Err, &cancelErr) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "Should have CancelError") + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + assert.NotNil(t, cancelErr.interruptSignal, "CancelError should have interrupt signal for checkpoint") + + resumeAgent1 := newCancelTestAgentWithToolsFinalAnswer(t, "seq_slow") + resumeAgent2 := newCancelTestAgentWithToolsFinalAnswer(t, "seq_fast") + + resumeSeq, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq_agent", + Description: "Sequential workflow", + SubAgents: []Agent{resumeAgent1, resumeAgent2}, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumeSeq, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, "seq-cancel-1") + assert.NoError(t, err) + assert.NotNil(t, resumeIter) + + var resumeEvents []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + assert.Nil(t, event.Err, "Should not have error during resume") + resumeEvents = append(resumeEvents, event) + } + assert.NotEmpty(t, resumeEvents, "Resume should produce events") +} + +func TestCancelImmediate_OrphanedToolGoroutine_NoPanic(t *testing.T) { + t.Run("unit_send_after_close", func(t *testing.T) { + _, gen := NewAsyncIteratorPair[*AgentEvent]() + + cc := newCancelContext() + cc.setMode(CancelImmediate) + close(cc.cancelChan) + close(cc.immediateChan) + + gen.Close() + + execCtx := &chatModelAgentExecCtx{ + generator: gen, + cancelCtx: cc, + } + + assert.NotPanics(t, func() { + execCtx.send(&AgentEvent{AgentName: "test"}) + }, "send after generator.Close must not panic") + }) + + t.Run("unit_send_after_close_without_cancel_ctx", func(t *testing.T) { + _, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Close() + + execCtx := &chatModelAgentExecCtx{ + generator: gen, + } + + assert.NotPanics(t, func() { + execCtx.send(&AgentEvent{AgentName: "test"}) + }, "send after generator.Close must not panic even without cancelCtx (trySend safety net)") + }) + + t.Run("unit_send_nil_execCtx", func(t *testing.T) { + var execCtx *chatModelAgentExecCtx + assert.NotPanics(t, func() { + execCtx.send(&AgentEvent{AgentName: "test"}) + }, "send on nil execCtx must not panic") + }) + + t.Run("unit_send_nil_generator", func(t *testing.T) { + execCtx := &chatModelAgentExecCtx{} + assert.NotPanics(t, func() { + execCtx.send(&AgentEvent{AgentName: "test"}) + }, "send with nil generator must not panic") + }) + + t.Run("unit_isImmediateCancelled_nil_cancelContext", func(t *testing.T) { + var cc *cancelContext + assert.False(t, cc.isImmediateCancelled(), "nil cancelContext should return false") + }) + + t.Run("unit_trySend_race_window", func(t *testing.T) { + _, gen := NewAsyncIteratorPair[*AgentEvent]() + cc := newCancelContext() + + gen.Close() + + execCtx := &chatModelAgentExecCtx{ + generator: gen, + cancelCtx: cc, + } + + assert.NotPanics(t, func() { + execCtx.send(&AgentEvent{AgentName: "test"}) + }, "trySend must handle the case where isImmediateCancelled is false but generator is closed") + }) + + t.Run("unit_SendEvent_after_close", func(t *testing.T) { + _, gen := NewAsyncIteratorPair[*AgentEvent]() + + cc := newCancelContext() + cc.setMode(CancelImmediate) + close(cc.cancelChan) + close(cc.immediateChan) + + gen.Close() + + execCtx := &chatModelAgentExecCtx{ + generator: gen, + cancelCtx: cc, + } + + ctx := withTypedChatModelAgentExecCtx(context.Background(), execCtx) + + assert.NotPanics(t, func() { + err := SendEvent(ctx, &AgentEvent{AgentName: "test"}) + assert.NoError(t, err) + }, "SendEvent after generator.Close must not panic") + }) + + t.Run("unit_SendEvent_no_execCtx", func(t *testing.T) { + err := SendEvent(context.Background(), &AgentEvent{AgentName: "test"}) + assert.Error(t, err, "SendEvent without execCtx should return error") + }) + + t.Run("integration_cancel_escalation_orphans_tool", func(t *testing.T) { + ctx := context.Background() + + toolStarted := make(chan struct{}, 1) + toolDone := make(chan struct{}, 1) + st := &slowToolWithSignal{ + name: "orphan_tool", + delay: 2 * time.Second, + result: "tool result", + startedChan: toolStarted, + } + + mdl := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_orphan_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "orphan_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OrphanTestAgent", + Description: "Test agent for orphaned tool goroutine panic", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + cancelOpt, cancelFn := WithCancel() + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("Use the tool")}, + }, cancelOpt) + assert.NotNil(t, iter) + + select { + case <-toolStarted: + case <-time.After(10 * time.Second): + t.Fatal("Tool did not start") + } + + timeout := 50 * time.Millisecond + handle, contributed := cancelFn( + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(timeout), + ) + assert.True(t, contributed, "Cancel should contribute") + + err = handle.Wait() + assert.True(t, err == nil || errors.Is(err, ErrCancelTimeout), + "handle.Wait should return nil or ErrCancelTimeout, got: %v", err) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + go func() { + time.Sleep(3 * time.Second) + select { + case toolDone <- struct{}{}: + default: + } + }() + + runtime.Gosched() + time.Sleep(3 * time.Second) + + select { + case <-toolDone: + default: + } + }) +} + +// -- Tests for CancelImmediate in nested agent structures -- + +func newTestChatModel(response *schema.Message, delay time.Duration) *cancelTestChatModel { + m := &cancelTestChatModel{ + response: response, + startedChan: make(chan struct{}, 1), + doneChan: make(chan struct{}, 1), + } + if delay > 0 { + m.setDelay(delay) + } + return m +} + +func newToolCallResponse(toolName string) *schema.Message { + return &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Type: "function", Function: schema.FunctionCall{Name: toolName, Arguments: `{}`}}, + }, + } +} + +func newAgentWithTool(t *testing.T, ctx context.Context, name string, mdl model.BaseChatModel, subAgent Agent) (Agent, error) { + t.Helper() + return NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: name, + Description: name, + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, subAgent)}, + }, + }, + }) +} + +func waitForChan(t *testing.T, ch <-chan struct{}, msg string) { + t.Helper() + select { + case <-ch: + case <-time.After(10 * time.Second): + t.Fatal(msg) + } +} + +func drainCancelError(t *testing.T, iter *AsyncIterator[*AgentEvent]) *CancelError { + t.Helper() + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + errors.As(event.Err, &cancelErr) + } + } + return cancelErr +} + +func drainResumeErrors(t *testing.T, iter *AsyncIterator[*AgentEvent]) []error { + t.Helper() + var errs []error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + errs = append(errs, event.Err) + } + } + return errs +} + +type cancelResult struct { + err error + contributed bool + done chan struct{} +} + +func cancelAsync(cancelFn AgentCancelFunc, opts ...AgentCancelOption) (cancelCalled chan struct{}, result *cancelResult) { + cancelCalled = make(chan struct{}) + result = &cancelResult{done: make(chan struct{})} + go func() { + handle, contributed := cancelFn(opts...) + result.contributed = contributed + close(cancelCalled) + result.err = handle.Wait() + close(result.done) + }() + return +} + +func (r *cancelResult) waitDone(t *testing.T) error { + t.Helper() + select { + case <-r.done: + return r.err + case <-time.After(10 * time.Second): + t.Fatal("cancel did not complete") + return nil + } +} + +func TestCancelImmediate_AgentTool_PreservesChildCheckpoint(t *testing.T) { + ctx := context.Background() + + leafModel := newTestChatModel( + &schema.Message{Role: schema.Assistant, Content: "leaf response"}, 2*time.Second) + leafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "leaf_agent", Description: "Leaf agent in agentTool", Model: leafModel, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "inner_seq", Description: "Inner sequential workflow", SubAgents: []Agent{leafAgent}, + }) + assert.NoError(t, err) + + rootModel := newTestChatModel(newToolCallResponse("inner_seq"), 0) + rootAgent, err := newAgentWithTool(t, ctx, "root_agent", rootModel, seqAgent) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{Agent: rootAgent, CheckPointStore: store}) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("immediate-agent-tool-1")) + + waitForChan(t, leafModel.startedChan, "Leaf agent model did not start") + + handle, contributed := cancelFn(WithRecursive()) + assert.True(t, contributed) + assert.NoError(t, handle.Wait()) + + cancelErr := drainCancelError(t, iter) + assert.NotNil(t, cancelErr, "Should have CancelError from CancelImmediate through agentTool") + assert.NotNil(t, cancelErr.interruptSignal) + + resumeLeaf, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "leaf_agent", Description: "Leaf agent in agentTool", + Model: newTestChatModel(&schema.Message{Role: schema.Assistant, Content: "resumed leaf"}, 0), + }) + assert.NoError(t, err) + resumeSeq, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "inner_seq", Description: "Inner sequential workflow", SubAgents: []Agent{resumeLeaf}, + }) + assert.NoError(t, err) + resumeRoot, err := newAgentWithTool(t, ctx, "root_agent", + newTestChatModel(&schema.Message{Role: schema.Assistant, Content: "resumed root"}, 0), resumeSeq) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{Agent: resumeRoot, CheckPointStore: store}) + resumeIter, err := runner2.Resume(ctx, "immediate-agent-tool-1") + assert.NoError(t, err) + assert.Empty(t, drainResumeErrors(t, resumeIter), "Resume should complete without errors") +} + +func TestCancelImmediate_ParallelWorkflow_WithAgentTool(t *testing.T) { + ctx := context.Background() + + leafModel := newTestChatModel( + &schema.Message{Role: schema.Assistant, Content: "leaf response"}, 2*time.Second) + leafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "leaf_agent", Description: "Leaf agent in agentTool", Model: leafModel, + }) + assert.NoError(t, err) + + agentWithTool, err := newAgentWithTool(t, ctx, "agent_with_tool", + newTestChatModel(newToolCallResponse("leaf_agent"), 0), leafAgent) + assert.NoError(t, err) + + simpleStarted := make(chan struct{}, 1) + simpleAgent := newCancelTestAgent(t, "simple_agent", 2*time.Second, simpleStarted) + + parAgent, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "par_agent", Description: "Parallel with agentTool and simple agent", + SubAgents: []Agent{agentWithTool, simpleAgent}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: parAgent, EnableStreaming: false}) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + waitForChan(t, leafModel.startedChan, "Leaf agent did not start") + waitForChan(t, simpleStarted, "Simple agent did not start") + + start := time.Now() + handle, _ := cancelFn() + assert.NoError(t, handle.Wait()) + + cancelErr := drainCancelError(t, iter) + elapsed := time.Since(start) + + assert.NotNil(t, cancelErr, "Should have CancelError from parallel with agentTool") + assert.True(t, elapsed < 5*time.Second, "Should complete quickly after cancel, elapsed: %v", elapsed) +} + +type cancelUnawareAgent struct { + name string + desc string + delay time.Duration + response string +} + +type multiResponseGatedModel struct { + responses []*schema.Message + gateChan chan struct{} + gateOnce bool + gated int32 + doneChan chan struct{} + callCount int32 +} + +func (m *multiResponseGatedModel) Generate(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + idx := atomic.AddInt32(&m.callCount, 1) + if m.gateChan != nil && (!m.gateOnce || atomic.CompareAndSwapInt32(&m.gated, 0, 1)) { + select { + case <-m.gateChan: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if len(m.responses) == 0 { + return nil, fmt.Errorf("multiResponseGatedModel: no responses configured") + } + resp := m.responses[(int(idx)-1)%len(m.responses)] + if m.doneChan != nil { + select { + case m.doneChan <- struct{}{}: + default: + } + } + return resp, nil +} + +func (m *multiResponseGatedModel) Stream(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + resp, err := m.Generate(ctx, msgs, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{resp}), nil +} + +func (m *multiResponseGatedModel) BindTools(tools []*schema.ToolInfo) error { return nil } + +func (a *cancelUnawareAgent) Name(_ context.Context) string { return a.name } +func (a *cancelUnawareAgent) Description(_ context.Context) string { return a.desc } + +func (a *cancelUnawareAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + // Intentionally ignores ctx.Done() — simulates a custom agent that + // does not participate in the cancel protocol at all. + // Delay is kept short (relative to grace period) to avoid goroutine + // leak lasting long after the test completes. + time.Sleep(a.delay) + }() + return iter +} + +func TestCancelImmediate_CustomAgent_GracePeriodFallback(t *testing.T) { + ctx := context.Background() + + customAgent := &cancelUnawareAgent{ + name: "custom_slow", desc: "A custom agent that ignores cancel", + delay: 5 * time.Second, response: "custom response", + } + + rootModel := newTestChatModel(newToolCallResponse("custom_slow"), 0) + rootAgent, err := newAgentWithTool(t, ctx, "root_agent", rootModel, customAgent) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{Agent: rootAgent, EnableStreaming: false}) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt) + + waitForChan(t, rootModel.startedChan, "Root model did not start") + waitForChan(t, rootModel.doneChan, "Root model did not finish") + + start := time.Now() + handle, _ := cancelFn() + assert.NoError(t, handle.Wait()) + + cancelErr := drainCancelError(t, iter) + elapsed := time.Since(start) + + assert.NotNil(t, cancelErr, "Should have CancelError (from grace period fallback)") + assert.True(t, elapsed < 5*time.Second, + "Should complete within grace period + overhead, elapsed: %v", elapsed) +} + +func TestCancelImmediate_MultiLevelNesting(t *testing.T) { + ctx := context.Background() + + innerLeafModel := newTestChatModel( + &schema.Message{Role: schema.Assistant, Content: "inner leaf response"}, 2*time.Second) + innerLeafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner_leaf", Description: "Innermost leaf agent", Model: innerLeafModel, + }) + assert.NoError(t, err) + + middleAgent, err := newAgentWithTool(t, ctx, "middle_agent", + newTestChatModel(newToolCallResponse("inner_leaf"), 0), innerLeafAgent) + assert.NoError(t, err) + + rootAgent, err := newAgentWithTool(t, ctx, "root_agent", + newTestChatModel(newToolCallResponse("middle_agent"), 0), middleAgent) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{Agent: rootAgent, CheckPointStore: store}) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("multi-level-1")) + + waitForChan(t, innerLeafModel.startedChan, "Inner leaf model did not start") + + start := time.Now() + handle, contributed := cancelFn() + assert.True(t, contributed) + assert.NoError(t, handle.Wait()) + + cancelErr := drainCancelError(t, iter) + elapsed := time.Since(start) + + assert.NotNil(t, cancelErr, "Should have CancelError from multi-level nesting") + assert.NotNil(t, cancelErr.interruptSignal) + assert.True(t, elapsed < 5*time.Second, "Should complete quickly, elapsed: %v", elapsed) + + resumeInnerLeaf, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "inner_leaf", Description: "Innermost leaf agent", + Model: newTestChatModel(&schema.Message{Role: schema.Assistant, Content: "resumed inner leaf"}, 0), + }) + assert.NoError(t, err) + resumeMiddle, err := newAgentWithTool(t, ctx, "middle_agent", + newTestChatModel(&schema.Message{Role: schema.Assistant, Content: "resumed middle"}, 0), resumeInnerLeaf) + assert.NoError(t, err) + resumeRoot, err := newAgentWithTool(t, ctx, "root_agent", + newTestChatModel(&schema.Message{Role: schema.Assistant, Content: "resumed root"}, 0), resumeMiddle) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{Agent: resumeRoot, CheckPointStore: store}) + resumeIter, err := runner2.Resume(ctx, "multi-level-1") + assert.NoError(t, err) + assert.Empty(t, drainResumeErrors(t, resumeIter), "Resume should complete without errors") +} + +func TestCancelImmediate_SequentialTransitionBoundary(t *testing.T) { + ctx := context.Background() + + model1 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent1 done"}, + gateChan: make(chan struct{}), + doneChan: make(chan struct{}, 1), + } + model2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent2 done"}, + doneChan: make(chan struct{}, 1), + } + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", Model: model1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", Model: model2, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "sequential test", SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt) + + for atomic.LoadInt32(&model1.callCount) == 0 { + runtime.Gosched() + } + + cancelCalled, result := cancelAsync(cancelFn) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(model1.gateChan) + + assert.NoError(t, result.waitDone(t), "CancelImmediate should succeed at transition") + + cancelErr := drainCancelError(t, iter) + + assert.NotNil(t, cancelErr, "Should have CancelError at transition boundary") + assert.Equal(t, int32(1), atomic.LoadInt32(&model1.callCount), "Agent1 model should be invoked") + assert.Equal(t, int32(0), atomic.LoadInt32(&model2.callCount), "Agent2 model should NOT be invoked (caught at transition)") +} + +func TestCancelImmediate_LoopTransitionBoundary(t *testing.T) { + ctx := context.Background() + + mdl := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "loop iter"}, + gateChan: make(chan struct{}), + doneChan: make(chan struct{}, 10), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "loop_inner", Description: "inner", Instruction: "test", Model: mdl, + }) + assert.NoError(t, err) + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop", Description: "loop test", SubAgents: []Agent{agent}, MaxIterations: 5, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: loopAgent, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt) + + for atomic.LoadInt32(&mdl.callCount) == 0 { + runtime.Gosched() + } + + cancelCalled, result := cancelAsync(cancelFn) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(mdl.gateChan) + + assert.NoError(t, result.waitDone(t), "CancelImmediate should succeed at loop transition") + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, int32(1), atomic.LoadInt32(&mdl.callCount), + "Model should be called once; second iteration caught at transition") +} + +func TestCancelAfterChatModel_SequentialTransitionBoundary(t *testing.T) { + ctx := context.Background() + + model1 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent1 done"}, + gateChan: make(chan struct{}), + doneChan: make(chan struct{}, 1), + } + model2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent2 done"}, + doneChan: make(chan struct{}, 1), + } + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", Model: model1, + }) + assert.NoError(t, err) + + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", Model: model2, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "sequential test", SubAgents: []Agent{agent1, agent2}, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, + EnableStreaming: false, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt, WithCheckPointID("chatmodel-transition-1")) + + for atomic.LoadInt32(&model1.callCount) == 0 { + runtime.Gosched() + } + + cancelCalled, result := cancelAsync(cancelFn, WithAgentCancelMode(CancelAfterChatModel)) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(model1.gateChan) + + assert.NoError(t, result.waitDone(t), "CancelAfterChatModel should succeed at transition boundary") + + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + cancelErr = ce + } + } + + assert.NotNil(t, cancelErr, "Should have CancelError at transition boundary") + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + assert.Equal(t, int32(1), atomic.LoadInt32(&model1.callCount), "Agent1 model should be invoked") + assert.Equal(t, int32(0), atomic.LoadInt32(&model2.callCount), + "Agent2 model should NOT be invoked (CancelAfterChatModel caught at transition)") +} + +func TestCancelAfterChatModel_Sequential_Agent1CompletesCancelBeforeAgent2Resume(t *testing.T) { + ctx := context.Background() + + model1 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent1 done"}, + gateChan: make(chan struct{}), + doneChan: make(chan struct{}, 1), + } + model2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent2 done"}, + doneChan: make(chan struct{}, 1), + } + model3 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "agent3 done"}, + doneChan: make(chan struct{}, 1), + } + + agent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", Model: model1, + }) + assert.NoError(t, err) + agent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", Model: model2, + }) + assert.NoError(t, err) + agent3, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent3", Description: "third", Instruction: "test", Model: model3, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "3-agent sequential", SubAgents: []Agent{agent1, agent2, agent3}, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, CheckPointStore: store, EnableStreaming: false, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt, + WithCheckPointID("seq-transition-resume-1")) + + for atomic.LoadInt32(&model1.callCount) == 0 { + runtime.Gosched() + } + + cancelCalled, result := cancelAsync(cancelFn, WithAgentCancelMode(CancelAfterChatModel)) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(model1.gateChan) + + assert.NoError(t, result.waitDone(t)) + + cancelErr := drainCancelError(t, iter) + assert.NotNil(t, cancelErr, "Should have CancelError") + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + assert.Equal(t, int32(1), atomic.LoadInt32(&model1.callCount)) + assert.Equal(t, int32(0), atomic.LoadInt32(&model2.callCount), + "Agent2 should NOT run (cancel caught at transition after agent1)") + assert.Equal(t, int32(0), atomic.LoadInt32(&model3.callCount)) + + resumeModel2 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "resumed agent2"}, + doneChan: make(chan struct{}, 1), + } + resumeModel3 := &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "resumed agent3"}, + doneChan: make(chan struct{}, 1), + } + + resumeAgent1, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent1", Description: "first", Instruction: "test", + Model: &gatedChatModel{ + response: &schema.Message{Role: schema.Assistant, Content: "should not run"}, + doneChan: make(chan struct{}, 1), + }, + }) + assert.NoError(t, err) + resumeAgent2, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent2", Description: "second", Instruction: "test", Model: resumeModel2, + }) + assert.NoError(t, err) + resumeAgent3, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent3", Description: "third", Instruction: "test", Model: resumeModel3, + }) + assert.NoError(t, err) + + resumeSeq, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq", Description: "3-agent sequential", + SubAgents: []Agent{resumeAgent1, resumeAgent2, resumeAgent3}, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumeSeq, CheckPointStore: store, EnableStreaming: false, + }) + resumeIter, err := runner2.Resume(ctx, "seq-transition-resume-1") + assert.NoError(t, err) + assert.Empty(t, drainResumeErrors(t, resumeIter), "Resume should complete without errors") + + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeModel2.callCount), + "Agent2 should run on resume") + assert.Equal(t, int32(1), atomic.LoadInt32(&resumeModel3.callCount), + "Agent3 should run on resume") +} + +func TestCancelAfterToolCalls_LoopTransitionBoundary(t *testing.T) { + ctx := context.Background() + + // Model that returns tool calls on odd calls and no tools on even calls. + // This completes one ReAct cycle per pair of calls: + // call 1 (gated): returns tool call → tool runs → call 2: returns no tools → END + // The gate only blocks the very first call. After that, all calls proceed instantly. + mdl := &multiResponseGatedModel{ + responses: []*schema.Message{ + {Role: schema.Assistant, ToolCalls: []schema.ToolCall{{ + ID: "call_1", Type: "function", + Function: schema.FunctionCall{Name: "loop_tool", Arguments: `{"input": "test"}`}, + }}}, + {Role: schema.Assistant, Content: "iteration done"}, + }, + gateChan: make(chan struct{}), + gateOnce: true, + doneChan: make(chan struct{}, 10), + } + + st := &slowTool{ + name: "loop_tool", + delay: 10 * time.Millisecond, + result: "tool done", + startedChan: make(chan struct{}, 10), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "loop_inner", Description: "inner", Instruction: "test", Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop", Description: "loop test", SubAgents: []Agent{agent}, MaxIterations: 10, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{Agent: loopAgent, CheckPointStore: store}) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("toolcalls-loop-1")) + + // Wait for the model to be entered (blocked on gate) + for atomic.LoadInt32(&mdl.callCount) == 0 { + runtime.Gosched() + } + + // Fire cancel, wait for it to be registered, then release the gate + cancelCalled, result := cancelAsync(cancelFn, WithAgentCancelMode(CancelAfterToolCalls)) + waitForChan(t, cancelCalled, "cancelFn was not called") + close(mdl.gateChan) + + // Iteration 1 completes fully (model→tool→model-no-tools→END). + // The CancelAfterToolCalls safe-point inside ReAct fires after tool calls, + // OR the transition boundary catches it before iteration 2. + // Note: this test doesn't deterministically distinguish which path fires — + // both are semantically correct for CancelAfterToolCalls. The transition- + // boundary code path for CancelAfterToolCalls in loops is not definitively + // covered here because the ReAct safe-point may handle it first. + assert.NoError(t, result.waitDone(t)) + + cancelErr := drainCancelError(t, iter) + assert.NotNil(t, cancelErr, "Should have CancelError from CancelAfterToolCalls in loop") + assert.Equal(t, CancelAfterToolCalls, cancelErr.Info.Mode) +} + +func TestCancelContext_RecursiveGraceBoundary(t *testing.T) { + t.Run("AgentToolDescendantMarker_PropagatesToParents", func(t *testing.T) { + parent := newCancelContext() + ctx := context.Background() + child := parent.deriveAgentToolCancelContext(ctx) + grandchild := child.deriveAgentToolCancelContext(ctx) + t.Cleanup(func() { + grandchild.markDone() + child.markDone() + }) + + assert.False(t, parent.hasAgentToolDescendant()) + assert.False(t, child.hasAgentToolDescendant()) + + grandchild.markAgentToolDescendant() + + assert.True(t, grandchild.hasAgentToolDescendant()) + assert.True(t, child.hasAgentToolDescendant()) + assert.True(t, parent.hasAgentToolDescendant()) + }) +} + +func TestCancel_ParallelWorkflow_CancelAfterChatModel(t *testing.T) { + ctx := context.Background() + slowStarted := make(chan struct{}, 1) + + slowAgent := newCancelTestAgentWithTools(t, "par_slow", 1*time.Second, slowStarted) + fastAgent := newCancelTestAgentWithTools(t, "par_fast", 50*time.Millisecond, make(chan struct{}, 1)) + + parAgent, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "par_agent", + Description: "Parallel workflow", + SubAgents: []Agent{slowAgent, fastAgent}, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: parAgent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("par-cancel-1")) + + select { + case <-slowStarted: + case <-time.After(10 * time.Second): + t.Fatal("Slow agent did not start") + } + + handle, contributed := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + assert.True(t, contributed, "Cancel should contribute") + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil && errors.As(event.Err, &cancelErr) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "Should have CancelError from parallel workflow") + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + + resumeSlow := newCancelTestAgentWithToolsFinalAnswer(t, "par_slow") + resumeFast := newCancelTestAgentWithToolsFinalAnswer(t, "par_fast") + + resumePar, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "par_agent", + Description: "Parallel workflow", + SubAgents: []Agent{resumeSlow, resumeFast}, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumePar, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, "par-cancel-1") + assert.NoError(t, err) + assert.NotNil(t, resumeIter) + + var resumeErrors []error + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Err != nil { + resumeErrors = append(resumeErrors, event.Err) + } + } + assert.Empty(t, resumeErrors, "Resume should complete without errors") +} + +func TestCancel_LoopWorkflow_CancelAfterChatModel(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 10) + + agent := newCancelTestAgentWithTools(t, "loop_inner", 500*time.Millisecond, modelStarted) + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop_agent", + Description: "Loop workflow", + SubAgents: []Agent{agent}, + MaxIterations: 10, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: loopAgent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("loop-cancel-1")) + + select { + case <-modelStarted: + case <-time.After(10 * time.Second): + t.Fatal("Model did not start") + } + + handle, contributed := cancelFn(WithAgentCancelMode(CancelAfterChatModel)) + assert.True(t, contributed, "Cancel should contribute") + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil && errors.As(event.Err, &cancelErr) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "Should have CancelError from loop workflow") + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + + resumeAgent := newCancelTestAgentWithToolsFinalAnswer(t, "loop_inner") + + resumeLoop, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop_agent", + Description: "Loop workflow", + SubAgents: []Agent{resumeAgent}, + MaxIterations: 10, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumeLoop, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, "loop-cancel-1") + assert.NoError(t, err) + assert.NotNil(t, resumeIter) + + var resumeEvents []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + assert.Nil(t, event.Err, "Should not have error during resume") + resumeEvents = append(resumeEvents, event) + } + assert.NotEmpty(t, resumeEvents, "Resume should produce events") +} + +func TestCancel_NestedWorkflow_AgentTool_CancelAfterChatModel(t *testing.T) { + // Structure: Runner -> RootCMA (with tools) -> agentTool -> flowAgent -> seqWorkflow -> LeafCMA + ctx := context.Background() + leafStarted := make(chan struct{}, 1) + + leafModel := &cancelTestChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "leaf response", + }, + startedChan: leafStarted, + doneChan: make(chan struct{}, 1), + } + leafModel.setDelay(500 * time.Millisecond) + + leafAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "leaf_agent", + Description: "Leaf agent in workflow", + Model: leafModel, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "inner_seq", + Description: "Inner sequential workflow", + SubAgents: []Agent{leafAgent}, + }) + assert.NoError(t, err) + + rootModel := &cancelTestChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "inner_seq", + Arguments: `{}`, + }, + }, + }, + }, + startedChan: make(chan struct{}, 1), + doneChan: make(chan struct{}, 1), + } + rootAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "root_agent", + Description: "Root agent", + Model: rootModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, seqAgent)}, + }, + }, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: rootAgent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, cancelOpt, WithCheckPointID("nested-cancel-1")) + + select { + case <-leafStarted: + case <-time.After(10 * time.Second): + t.Fatal("Leaf agent model did not start") + } + + handle, contributed := cancelFn(WithAgentCancelMode(CancelAfterChatModel), WithRecursive()) + assert.True(t, contributed, "Cancel should contribute") + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil && errors.As(event.Err, &cancelErr) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "Should have CancelError from deeply nested workflow") + assert.Equal(t, CancelAfterChatModel, cancelErr.Info.Mode) + assert.NotNil(t, cancelErr.interruptSignal, "CancelError should carry interrupt signal through agent tree") + + // Phase 2: Resume from checkpoint — new instances to avoid data races + resumeLeafModel := &cancelTestChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "resumed leaf response", + }, + startedChan: make(chan struct{}, 1), + doneChan: make(chan struct{}, 1), + } + resumeLeaf, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "leaf_agent", + Description: "Leaf agent in workflow", + Model: resumeLeafModel, + }) + assert.NoError(t, err) + + resumeSeq, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "inner_seq", + Description: "Inner sequential workflow", + SubAgents: []Agent{resumeLeaf}, + }) + assert.NoError(t, err) + + resumeRootModel := &cancelTestChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "resumed root response", + }, + startedChan: make(chan struct{}, 1), + doneChan: make(chan struct{}, 1), + } + resumeRoot, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "root_agent", + Description: "Root agent", + Model: resumeRootModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, resumeSeq)}, + }, + }, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumeRoot, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, "nested-cancel-1") + assert.NoError(t, err) + assert.NotNil(t, resumeIter) + + var resumeErrors []error + for { + event, ok := resumeIter.Next() + if !ok { + break + } + if event.Err != nil { + resumeErrors = append(resumeErrors, event.Err) + } + } + assert.Empty(t, resumeErrors, "Resume should complete without errors") +} + +func TestCancel_CancelAfterToolCalls_InSequentialWorkflow(t *testing.T) { + ctx := context.Background() + toolStarted := make(chan struct{}, 1) + + st := &slowTool{ + name: "slow_tool", + delay: 200 * time.Millisecond, + result: "tool done", + startedChan: toolStarted, + } + + modelWithToolCall := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "slow_tool", + Arguments: `{"input": "test"}`, + }, + }, + }, + }, + } + + agentWithTools, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent_with_tools", + Description: "Agent with slow tool", + Model: modelWithToolCall, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{st}, + }, + }, + }) + assert.NoError(t, err) + + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq_agent", + Description: "Sequential workflow with tool agent", + SubAgents: []Agent{agentWithTools}, + }) + assert.NoError(t, err) + + store := newCancelTestStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: seqAgent, + CheckPointStore: store, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("Use the tool")}, cancelOpt, WithCheckPointID("tool-cancel-1")) + + select { + case <-toolStarted: + case <-time.After(10 * time.Second): + t.Fatal("Tool did not start") + } + + // Cancel after tool calls — should wait for the tool to finish, then cancel + handle, contributed := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + assert.True(t, contributed, "Cancel should contribute") + err = handle.Wait() + assert.NoError(t, err) + + hasCancelError := false + var cancelErr *CancelError + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil && errors.As(event.Err, &cancelErr) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "Should have CancelError after tool calls complete") + assert.Equal(t, CancelAfterToolCalls, cancelErr.Info.Mode) + + // Phase 2: Resume from checkpoint — new instances + resumeTool := &slowTool{ + name: "slow_tool", + delay: 50 * time.Millisecond, + result: "resumed tool done", + startedChan: make(chan struct{}, 1), + } + + resumeModel := &simpleChatModel{ + response: &schema.Message{ + Role: schema.Assistant, + Content: "resumed response after tool", + }, + } + + resumeAgentWithTools, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "agent_with_tools", + Description: "Agent with slow tool", + Model: resumeModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{resumeTool}, + }, + }, + }) + assert.NoError(t, err) + + resumeSeq, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "seq_agent", + Description: "Sequential workflow with tool agent", + SubAgents: []Agent{resumeAgentWithTools}, + }) + assert.NoError(t, err) + + runner2 := NewRunner(ctx, RunnerConfig{ + Agent: resumeSeq, + CheckPointStore: store, + }) + + resumeIter, err := runner2.Resume(ctx, "tool-cancel-1") + assert.NoError(t, err) + assert.NotNil(t, resumeIter) + + var resumeEvents []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + assert.Nil(t, event.Err, "Should not have error during resume") + resumeEvents = append(resumeEvents, event) + } + assert.NotEmpty(t, resumeEvents, "Resume should produce events") +} + +// TestCancel_SafePointNeverFires_ErrExecutionEnded verifies the waitForCompletion +// path where a safe-point cancel is submitted while the agent is running, but +// the agent finishes without hitting the requested safe-point (e.g. +// CancelAfterToolCalls on an agent with no tool calls). The cancel CAS succeeds +// (stateRunning → stateCancelling), but the agent completes normally (markDone → +// stateDone), so waitForCompletion returns ErrExecutionEnded. +func TestCancel_SafePointNeverFires_ErrExecutionEnded(t *testing.T) { + ctx := context.Background() + + gate := make(chan struct{}) + done := make(chan struct{}, 1) + + m := &gatedChatModel{ + gateChan: gate, + doneChan: done, + response: &schema.Message{ + Role: schema.Assistant, + Content: "Final answer, no tool calls", + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "NoToolAgent", + Description: "Agent with no tools", + Instruction: "You are a test assistant", + Model: m, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + }) + + cancelOpt, cancelFn := WithCancel() + iter := runner.Run(ctx, []Message{schema.UserMessage("hello")}, cancelOpt) + + // Wait a moment for the agent to enter Generate and block on gateChan. + runtime.Gosched() + time.Sleep(50 * time.Millisecond) + + // Submit a safe-point cancel for tool calls. The agent has no tools, + // so this safe-point will never fire. + handle, submitted := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + assert.True(t, submitted) + + // Let the model complete. The agent finishes without hitting the tool + // calls safe-point → markDone → stateDone → waitForCompletion returns + // ErrExecutionEnded. + close(gate) + + waitErr := handle.Wait() + assert.ErrorIs(t, waitErr, ErrExecutionEnded) + + for { + _, ok := iter.Next() + if !ok { + break + } + } +} + +// TestBuildCancelFunc_StateDoneUnderLock exercises the race-condition path +// in buildCancelFunc where the state transitions to stateDone between the +// lockless check and the locked check (cancel.go L732-734). +func TestBuildCancelFunc_StateDoneUnderLock(t *testing.T) { + cc := newCancelContext() + cancelFn := cc.buildCancelFunc() + + // Hold cancelMu so the cancel func blocks when it tries to acquire the lock. + cc.cancelMu.Lock() + + type result struct { + handle *CancelHandle + ok bool + } + ch := make(chan result, 1) + + go func() { + h, ok := cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + ch <- result{h, ok} + }() + + // Give the goroutine time to reach the Lock() call. + runtime.Gosched() + time.Sleep(20 * time.Millisecond) + + // Transition to stateDone while the cancel goroutine is blocked on the lock. + cc.markDone() + + // Release the lock. The cancel func resumes and finds stateDone. + cc.cancelMu.Unlock() + + r := <-ch + assert.False(t, r.ok, "cancel should not be accepted when execution already done") + assert.ErrorIs(t, r.handle.Wait(), ErrExecutionEnded) +} + +// TestBuildCancelFunc_CASFailStateDone exercises the race-condition path +// in buildCancelFunc where the CAS on stateRunning→stateCancelling fails +// because markDone transitioned stateRunning→stateDone concurrently +// (cancel.go L742-743). +func TestBuildCancelFunc_CASFailStateDone(t *testing.T) { + // Exercises cancel.go L742-743: CAS(stateRunning→stateCancelling) fails + // because markDone transitions stateRunning→stateDone concurrently. + // + // The window between the state check (L738) and CAS (L739) is extremely + // tight. We maximize the chance by having the cancel goroutine block on + // cancelMu, then racing markDone with the lock release. + hit := false + for i := 0; i < 100000 && !hit; i++ { + cc := newCancelContext() + cancelFn := cc.buildCancelFunc() + + // Hold cancelMu so the cancel goroutine blocks at L725. + cc.cancelMu.Lock() + + cancelDone := make(chan struct{}) + var h *CancelHandle + var ok bool + + go func() { + defer close(cancelDone) + h, ok = cancelFn(WithAgentCancelMode(CancelAfterToolCalls)) + }() + + // Let the cancel goroutine reach the Lock() call. + runtime.Gosched() + + // Release lock and fire markDone concurrently. The cancel goroutine + // will acquire the lock and race with markDone on the CAS. + go cc.markDone() + cc.cancelMu.Unlock() + + <-cancelDone + + if !ok && errors.Is(h.Wait(), ErrExecutionEnded) { + hit = true + } + } + if hit { + t.Log("Successfully hit CAS-fail → stateDone path") + } else { + t.Log("CAS race path not triggered (L743 remains a theoretical race edge)") + } +} diff --git a/adk/chatmodel.go b/adk/chatmodel.go new file mode 100644 index 0000000..37a55f1 --- /dev/null +++ b/adk/chatmodel.go @@ -0,0 +1,1714 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "fmt" + "math" + "runtime/debug" + "strings" + "sync" + "sync/atomic" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +var _ ResumableAgent = &TypedChatModelAgent[*schema.Message]{} +var _ TypedResumableAgent[*schema.AgenticMessage] = &TypedChatModelAgent[*schema.AgenticMessage]{} + +type typedChatModelAgentExecCtx[M MessageType] struct { + runtimeReturnDirectly map[string]bool + generator *AsyncGenerator[*TypedAgentEvent[M]] + cancelCtx *cancelContext + + failoverLastSuccessModel model.BaseModel[M] + + // suppressEventSend prevents eventSenderModel from emitting AgentEvents for the current + // Generate call. Set to true before each rejected retry attempt and reset to false after. + // Invariant: any code path that emits model output events MUST check this flag. + suppressEventSend bool + retryVerdictSignal *retryVerdictSignal + + afterToolCallsHook func(ctx context.Context) error +} + +func (e *typedChatModelAgentExecCtx[M]) send(event *TypedAgentEvent[M]) { + if e == nil || e.generator == nil { + return + } + if e.cancelCtx != nil && e.cancelCtx.isImmediateCancelled() { + return + } + e.generator.trySend(event) +} + +type chatModelAgentExecCtx = typedChatModelAgentExecCtx[*schema.Message] + +type typedChatModelAgentExecCtxKey[M MessageType] struct{} + +func withTypedChatModelAgentExecCtx[M MessageType](ctx context.Context, execCtx *typedChatModelAgentExecCtx[M]) context.Context { + return context.WithValue(ctx, typedChatModelAgentExecCtxKey[M]{}, execCtx) +} + +func getTypedChatModelAgentExecCtx[M MessageType](ctx context.Context) *typedChatModelAgentExecCtx[M] { + if v := ctx.Value(typedChatModelAgentExecCtxKey[M]{}); v != nil { + return v.(*typedChatModelAgentExecCtx[M]) + } + return nil +} + +type chatModelAgentRunOptions struct { + chatModelOptions []model.Option + toolOptions []tool.Option + agentToolOptions map[string][]AgentRunOption + + historyModifier func(context.Context, []Message) []Message + + afterToolCallsHook func(ctx context.Context) error +} + +// WithChatModelOptions sets options for the underlying chat model. +func WithChatModelOptions(opts []model.Option) AgentRunOption { + return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) { + t.chatModelOptions = opts + }) +} + +// WithToolOptions sets options for tools used by the chat model agent. +func WithToolOptions(opts []tool.Option) AgentRunOption { + return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) { + t.toolOptions = opts + }) +} + +// WithAgentToolRunOptions specifies per-tool run options for the agent. +func WithAgentToolRunOptions(opts map[string][]AgentRunOption) AgentRunOption { + return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) { + t.agentToolOptions = opts + }) +} + +// WithHistoryModifier sets a function to modify history during resume. +// Deprecated: use ResumeWithData and ChatModelAgentResumeData instead. +func WithHistoryModifier(f func(context.Context, []Message) []Message) AgentRunOption { + return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) { + t.historyModifier = f + }) +} + +// WithAfterToolCallsHook registers a per-run hook that fires synchronously after +// all tool calls in a react iteration complete, before the next ChatModel call. +// +// This is suitable for TurnLoop Push+Preempt patterns where the pushed item +// must be visible to the next turn's GenInput. +func WithAfterToolCallsHook(fn func(ctx context.Context) error) AgentRunOption { + return WrapImplSpecificOptFn(func(t *chatModelAgentRunOptions) { + t.afterToolCallsHook = fn + }) +} + +type ToolsConfig struct { + compose.ToolsNodeConfig + + // ReturnDirectly specifies tools that cause the agent to return immediately when called. + // The map keys are tool names indicate whether the tool should trigger immediate return. + ReturnDirectly map[string]bool + + // EmitInternalEvents indicates whether internal events from agentTool should be emitted + // to the parent agent's AsyncGenerator, allowing real-time streaming of nested agent output + // to the end-user via Runner. + // + // Note that these forwarded events are NOT recorded in the parent agent's runSession. + // They are only emitted to the end-user and have no effect on the parent agent's state + // or checkpoint. + // + // Action Scoping: + // Actions emitted by the inner agent are scoped to the agent tool boundary: + // - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume + // - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool + EmitInternalEvents bool +} + +// TypedGenModelInput transforms the agent's system instruction and user input into model input +// messages ([]M). This is the primary customization point for controlling what the model sees. +// The default implementation prepends a system message (if instruction is non-empty), +// followed by the user's input messages. +type TypedGenModelInput[M MessageType] func(ctx context.Context, instruction string, input *TypedAgentInput[M]) ([]M, error) + +// GenModelInput transforms agent instructions and input into a format suitable for the model. +type GenModelInput = TypedGenModelInput[*schema.Message] + +func defaultGenModelInput(ctx context.Context, instruction string, input *AgentInput) ([]Message, error) { + msgs := make([]Message, 0, len(input.Messages)+1) + + if instruction != "" { + sp := schema.SystemMessage(instruction) + + vs := GetSessionValues(ctx) + if len(vs) > 0 { + ct := prompt.FromMessages(schema.FString, sp) + ms, err := ct.Format(ctx, vs) + if err != nil { + return nil, fmt.Errorf("defaultGenModelInput: failed to format instruction using FString template. "+ + "This formatting is triggered automatically when SessionValues are present. "+ + "If your instruction contains literal curly braces (e.g., JSON), provide a custom GenModelInput that uses another format. If you are using "+ + "SessionValues for purposes other than instruction formatting, provide a custom GenModelInput that does no formatting at all: %w", err) + } + + sp = ms[0] + } + + msgs = append(msgs, sp) + } + + msgs = append(msgs, input.Messages...) + + return msgs, nil +} + +func newDefaultGenModelInput[M MessageType]() TypedGenModelInput[M] { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(GenModelInput(defaultGenModelInput)).(TypedGenModelInput[M]) + case *schema.AgenticMessage: + return any(TypedGenModelInput[*schema.AgenticMessage](func(_ context.Context, instruction string, input *TypedAgentInput[*schema.AgenticMessage]) ([]*schema.AgenticMessage, error) { + msgs := make([]*schema.AgenticMessage, 0, len(input.Messages)+1) + if instruction != "" { + msgs = append(msgs, schema.SystemAgenticMessage(instruction)) + } + msgs = append(msgs, input.Messages...) + return msgs, nil + })).(TypedGenModelInput[M]) + default: + panic("unreachable: unknown MessageType") + } +} + +// TypedChatModelAgentState represents the state of a chat model agent during conversation. +// This is the primary state type for both TypedChatModelAgentMiddleware and AgentMiddleware callbacks. +type TypedChatModelAgentState[M MessageType] struct { + // Messages contains all messages in the current conversation session. + Messages []M + + // ToolInfos contains the tool definitions passed to the model via model.WithTools. + // BeforeModelRewriteState handlers can read and modify this field to control which tools + // the model sees on each call. + ToolInfos []*schema.ToolInfo + + // DeferredToolInfos contains tool definitions for server-side deferred retrieval, + // passed to the model via model.WithDeferredTools. These tools are not included in the + // immediate tool list but can be discovered by the model through its native search capability. + // Nil when not in use. + DeferredToolInfos []*schema.ToolInfo +} + +// ChatModelAgentState is the default state type using *schema.Message. +type ChatModelAgentState = TypedChatModelAgentState[*schema.Message] + +// Deprecated: Use ChatModelAgentMiddleware (interface-based Handlers) instead. +// AgentMiddleware will be removed in a future release. +// +// AgentMiddleware provides hooks to customize agent behavior at various stages of execution. +type AgentMiddleware struct { + // AdditionalInstruction adds supplementary text to the agent's system instruction. + // This instruction is concatenated with the base instruction before each chat model call. + AdditionalInstruction string + + // AdditionalTools adds supplementary tools to the agent's available toolset. + // These tools are combined with the tools configured for the agent. + AdditionalTools []tool.BaseTool + + // BeforeChatModel is called before each ChatModel invocation, allowing modification of the agent state. + BeforeChatModel func(context.Context, *ChatModelAgentState) error + + // AfterChatModel is called after each ChatModel invocation, allowing modification of the agent state. + AfterChatModel func(context.Context, *ChatModelAgentState) error + + // WrapToolCall wraps tool calls with custom middleware logic. + // Each middleware contains Invokable and/or Streamable functions for tool calls. + WrapToolCall compose.ToolMiddleware +} + +// TypedChatModelAgentConfig is the generic configuration for ChatModelAgent. +type TypedChatModelAgentConfig[M MessageType] struct { + // Name of the agent. Better be unique across all agents. + // Optional. If empty, the agent can still run standalone but cannot be used as + // a sub-agent tool via NewAgentTool (which requires a non-empty Name). + Name string + // Description of the agent's capabilities. + // Helps other agents determine whether to transfer tasks to this agent. + // Optional. If empty, the agent can still run standalone but cannot be used as + // a sub-agent tool via NewAgentTool (which requires a non-empty Description). + Description string + // Instruction used as the system prompt for this agent. + // Optional. If empty, no system prompt will be used. + // Supports f-string placeholders for session values in default GenModelInput, for example: + // "You are a helpful assistant. The current time is {Time}. The current user is {User}." + // These placeholders will be replaced with session values for "Time" and "User". + Instruction string + + // Model is the chat model used by the agent. + // If your ChatModelAgent uses any tools, this model must support the model.WithTools + // call option, as that's how ChatModelAgent configures the model with tool information. + Model model.BaseModel[M] + + ToolsConfig ToolsConfig + + // GenModelInput transforms instructions and input messages into the model's input format. + // Optional. Defaults to defaultGenModelInput which combines instruction and messages. + GenModelInput TypedGenModelInput[M] + + // Exit defines the tool used to terminate the agent process. + // Optional. If nil, no Exit Action will be generated. + // You can use the provided 'ExitTool' implementation directly. + // + // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven + // to be more effective empirically. Consider using ChatModelAgent with AgentTool + // or DeepAgent instead for most multi-agent scenarios. + Exit tool.BaseTool + + // OutputKey stores the agent's response in the session. + // Optional. When set, stores output via AddSessionValue(ctx, outputKey, msg.Content). + // + // NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven + // to be more effective empirically. Consider using ChatModelAgent with AgentTool + // or DeepAgent instead for most multi-agent scenarios. + OutputKey string + + // MaxIterations defines the upper limit of ChatModel generation cycles. + // The agent will terminate with an error if this limit is exceeded. + // Optional. Defaults to 20. + MaxIterations int + + // Deprecated: Use Handlers instead. Middlewares will be removed in a future release. + // Handlers provides a more flexible interface-based approach for extending agent behavior. + Middlewares []AgentMiddleware + + // Handlers configures interface-based handlers for extending agent behavior. + // Unlike Middlewares (struct-based), Handlers allow users to: + // - Add custom methods to their handler implementations + // - Return modified context from handler methods + // - Centralize configuration in struct fields instead of closures + // + // Handlers are processed after Middlewares, in registration order. + // See ChatModelAgentMiddleware documentation for when to use Handlers vs Middlewares. + // + // Execution Order (relative to AgentMiddleware and ToolsConfig): + // + // Model call lifecycle (outermost to innermost wrapper chain): + // 1. AgentMiddleware.BeforeChatModel (hook, runs before model call) + // 2. ChatModelAgentMiddleware.BeforeModelRewriteState (hook, can modify state before model call) + // 3. failoverModelWrapper (internal - failover between models, if configured) + // 4. retryModelWrapper (internal - retries on failure, if configured) + // 5. eventSenderModelWrapper (internal - sends model response events) + // 6. ChatModelAgentMiddleware.WrapModel (wrapper, first registered is outermost) + // 7. callbackInjectionModelWrapper (internal - injects callbacks if not enabled; when failover is enabled, this is handled per-model inside failoverProxyModel instead) + // 8. failoverProxyModel (internal - dispatches to selected failover model, if configured) / Model.Generate/Stream + // 9. ChatModelAgentMiddleware.AfterModelRewriteState (hook, can modify state after model call) + // 10. AgentMiddleware.AfterChatModel (hook, runs after model call) + // + // Custom Event Sender Position: + // By default, events are sent after all user middlewares (WrapModel) have processed the output, + // containing the modified messages. To send events with original (unmodified) output, pass + // NewEventSenderModelWrapper as a Handler after the modifying middleware: + // + // agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + // Handlers: []adk.ChatModelAgentMiddleware{ + // myCustomHandler, // First registered = outermost wrapper + // adk.NewEventSenderModelWrapper(), // Last registered = innermost, events sent with original output + // }, + // }) + // + // Handler order: first registered is outermost. So [A, B, C] becomes A(B(C(model))). + // EventSenderModelWrapper sends events in post-processing, so placing it innermost + // means it receives the original model output before outer handlers modify it. + // + // When EventSenderModelWrapper is detected in Handlers, the framework skips + // the default event sender to avoid duplicate events. + // + // Tool call lifecycle (outermost to innermost): + // 1. eventSenderToolWrapper (internal ToolMiddleware - sends tool result events after all processing) + // 2. ToolsConfig.ToolCallMiddlewares (ToolMiddleware) + // 3. AgentMiddleware.WrapToolCall (ToolMiddleware) + // 4. ChatModelAgentMiddleware.WrapToolCall (wrapper, first registered is outermost) + // 5. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them) + // 6. Tool.InvokableRun/StreamableRun + // + // Custom Tool Event Sender Position: + // By default, tool result events are emitted by an internal event sender placed before + // all user middlewares (outermost), so events reflect the fully processed tool output. + // To control exactly where in the handler chain tool events are emitted, pass + // NewEventSenderToolWrapper() as one of the Handlers. Its position determines which + // middlewares' effects are visible in the emitted event: + // + // agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + // Handlers: []adk.ChatModelAgentMiddleware{ + // loggingHandler, // Outermost: sees event-sender output + // adk.NewEventSenderToolWrapper(), // Events reflect output from handlers below + // sanitizationHandler, // Innermost: runs first, modifies tool output + // }, + // }) + // + // Handler order: first registered is outermost. So [A, B, C] wraps as A(B(C(tool))). + // The event sender captures tool output in post-processing, so its position controls + // which handlers' modifications are included in the emitted events. + // + // When NewEventSenderToolWrapper is detected in Handlers, the framework skips + // the default event sender to avoid duplicate events. + // + // Tool List Modification: + // + // There are two ways to modify the tool list: + // + // 1. In BeforeAgent: Modify ChatModelAgentContext.Tools ([]tool.BaseTool) directly. This affects + // both the tool info list passed to ChatModel AND the actual tools available for + // execution. Changes persist for the entire agent run. + // + // 2. In BeforeModelRewriteState: Modify state.ToolInfos and state.DeferredToolInfos directly. + // This affects the tool info list passed to ChatModel for this and all subsequent model + // calls (changes are persisted in state). This is the recommended approach for dynamic + // tool filtering/selection based on conversation context. + // + // Modifying tools in WrapModel (e.g. via model.WithTools) is discouraged: changes there + // are NOT persisted in state, only affect a single model call, and break prompt cache. + Handlers []TypedChatModelAgentMiddleware[M] + + // ModelRetryConfig configures retry behavior for the ChatModel. + // When set, the agent will automatically retry failed ChatModel calls + // based on the configured policy. + // Optional. If nil, no retry will be performed. + ModelRetryConfig *TypedModelRetryConfig[M] + + // ModelFailoverConfig configures failover behavior for the ChatModel. + // When set, the agent will first try the last successful model (initially the configured Model), + // and on failure, call GetFailoverModel to select alternate models. + // Model field is still required as it serves as the initial model. + // Optional. If nil, no failover will be performed. + ModelFailoverConfig *ModelFailoverConfig[M] +} + +type ChatModelAgentConfig = TypedChatModelAgentConfig[*schema.Message] + +// TypedChatModelAgent is a chat model-backed agent parameterized by message type. +// +// For M = *schema.Message, the full ReAct loop (model → tool calls → model) is used. +// For M = *schema.AgenticMessage, a single-shot chain is used since agentic models +// handle tool calling internally. Cancel monitoring and retry on the model stream +// are not yet supported for agentic models. +type TypedChatModelAgent[M MessageType] struct { + name string + description string + instruction string + + model model.BaseModel[M] + toolsConfig ToolsConfig + + genModelInput TypedGenModelInput[M] + + outputKey string + maxIterations int + + subAgents []TypedAgent[M] + parentAgent TypedAgent[M] + + disallowTransferToParent bool + + exit tool.BaseTool + + handlers []TypedChatModelAgentMiddleware[M] + middlewares []AgentMiddleware + + modelRetryConfig *TypedModelRetryConfig[M] + modelFailoverConfig *ModelFailoverConfig[M] + + once sync.Once + run typedRunFunc[M] + frozen uint32 + exeCtx *execContext +} + +type ChatModelAgent = TypedChatModelAgent[*schema.Message] + +// typedRunParams holds the parameters for a typedRunFunc invocation. +type typedRunParams[M MessageType] struct { + input *TypedAgentInput[M] + generator *AsyncGenerator[*TypedAgentEvent[M]] + store *bridgeStore + instruction string + returnDirectly map[string]bool + cancelCtx *cancelContext + cancelCtxOwned bool + composeOpts []compose.Option + + afterToolCallsHook func(ctx context.Context) error +} + +type typedRunFunc[M MessageType] func(ctx context.Context, p *typedRunParams[M]) + +func resolveRunCancelContext(ctx context.Context, o *options) (*cancelContext, bool) { + inherited := getCancelContext(ctx) + if o.cancelCtx != nil { + return o.cancelCtx, o.cancelCtx != inherited + } + return inherited, false +} + +// NewChatModelAgent creates a new ChatModelAgent with the given config. +func NewChatModelAgent(ctx context.Context, config *ChatModelAgentConfig) (*ChatModelAgent, error) { + return NewTypedChatModelAgent[*schema.Message](ctx, config) +} + +// NewTypedChatModelAgent creates a new TypedChatModelAgent with the given config. +func NewTypedChatModelAgent[M MessageType](ctx context.Context, config *TypedChatModelAgentConfig[M]) (*TypedChatModelAgent[M], error) { + if config.ModelFailoverConfig != nil { + if config.ModelFailoverConfig.GetFailoverModel == nil { + return nil, errors.New("ModelFailoverConfig.GetFailoverModel is required when ModelFailoverConfig is set") + } + + // ShouldFailover is required when ModelFailoverConfig is set + if config.ModelFailoverConfig.ShouldFailover == nil { + return nil, errors.New("ModelFailoverConfig.ShouldFailover is required when ModelFailoverConfig is set") + } + } + + if config.Model == nil { + return nil, errors.New("agent 'Model' is required") + } + + var genInput TypedGenModelInput[M] + if config.GenModelInput != nil { + genInput = config.GenModelInput + } else { + genInput = newDefaultGenModelInput[M]() + } + + tc := config.ToolsConfig + + // Tool call middleware execution order (outermost to innermost): + // 1. eventSenderToolWrapper (internal - sends tool result events after all modifications) + // 2. User-provided ToolsConfig.ToolCallMiddlewares (original order preserved) + // 3. Middlewares' WrapToolCall (in registration order) + // 4. ChatModelAgentMiddleware.WrapToolCall (in registration order) + // 5. callbackInjectedToolCall (internal - injects callbacks if tool doesn't handle them) + if !hasUserEventSenderToolWrapper(config.Handlers) { + defaultToolEventSender := handlersToToolMiddlewares([]TypedChatModelAgentMiddleware[M]{newTypedEventSenderToolWrapper[M]()}) + tc.ToolCallMiddlewares = append(defaultToolEventSender, tc.ToolCallMiddlewares...) + } + tc.ToolCallMiddlewares = append(tc.ToolCallMiddlewares, collectToolMiddlewaresFromMiddlewares(config.Middlewares)...) + + // Cancel monitoring middleware (innermost — close to the tool endpoint). + // This allows early abort of the raw tool result stream when immediateChan fires + // (CancelImmediate or timeout escalation), while requiring outer wrappers to + // propagate stream errors such as ErrStreamCanceled without swallowing them. + cancelToolHandler := &cancelMonitoredToolHandler{} + tc.ToolCallMiddlewares = append(tc.ToolCallMiddlewares, compose.ToolMiddleware{ + Streamable: cancelToolHandler.WrapStreamableToolCall, + EnhancedStreamable: cancelToolHandler.WrapEnhancedStreamableToolCall, + }) + + return &TypedChatModelAgent[M]{ + name: config.Name, + description: config.Description, + instruction: config.Instruction, + model: config.Model, + toolsConfig: tc, + genModelInput: genInput, + exit: config.Exit, + outputKey: config.OutputKey, + maxIterations: config.MaxIterations, + handlers: config.Handlers, + middlewares: config.Middlewares, + modelRetryConfig: config.ModelRetryConfig, + modelFailoverConfig: config.ModelFailoverConfig, + }, nil +} + +func collectToolMiddlewaresFromMiddlewares(mws []AgentMiddleware) []compose.ToolMiddleware { + var middlewares []compose.ToolMiddleware + for _, m := range mws { + if m.WrapToolCall.Invokable == nil && m.WrapToolCall.Streamable == nil && m.WrapToolCall.EnhancedStreamable == nil && m.WrapToolCall.EnhancedInvokable == nil { + continue + } + middlewares = append(middlewares, m.WrapToolCall) + } + return middlewares +} + +const ( + TransferToAgentToolName = "transfer_to_agent" + TransferToAgentToolDesc = "Transfer the question to another agent." + TransferToAgentToolDescChinese = "将问题移交给其他 Agent。" +) + +var ( + toolInfoTransferToAgent = &schema.ToolInfo{ + Name: TransferToAgentToolName, + + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "agent_name": { + Desc: "the name of the agent to transfer to", + Required: true, + Type: schema.String, + }, + }), + } + + ToolInfoExit = &schema.ToolInfo{ + Name: "exit", + Desc: "Exit the agent process and return the final result.", + + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "final_result": { + Desc: "the final result to return", + Required: true, + Type: schema.String, + }, + }), + } +) + +type ExitTool struct{} + +func (et ExitTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return ToolInfoExit, nil +} + +func (et ExitTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + type exitParams struct { + FinalResult string `json:"final_result"` + } + + params := &exitParams{} + err := sonic.UnmarshalString(argumentsInJSON, params) + if err != nil { + return "", err + } + + err = SendToolGenAction(ctx, "exit", NewExitAction()) + if err != nil { + return "", err + } + + return params.FinalResult, nil +} + +type transferToAgent struct{} + +func (tta transferToAgent) Info(_ context.Context) (*schema.ToolInfo, error) { + desc := internal.SelectPrompt(internal.I18nPrompts{ + English: TransferToAgentToolDesc, + Chinese: TransferToAgentToolDescChinese, + }) + info := *toolInfoTransferToAgent + info.Desc = desc + return &info, nil +} + +func transferToAgentToolOutput(destName string) string { + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: "successfully transferred to agent [%s]", + Chinese: "成功移交任务至 agent [%s]", + }) + return fmt.Sprintf(tpl, destName) +} + +func (tta transferToAgent) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + type transferParams struct { + AgentName string `json:"agent_name"` + } + + params := &transferParams{} + err := sonic.UnmarshalString(argumentsInJSON, params) + if err != nil { + return "", err + } + + err = SendToolGenAction(ctx, TransferToAgentToolName, NewTransferToAgentAction(params.AgentName)) + if err != nil { + return "", err + } + + return transferToAgentToolOutput(params.AgentName), nil +} + +func (a *TypedChatModelAgent[M]) Name(_ context.Context) string { + return a.name +} + +func (a *TypedChatModelAgent[M]) Description(_ context.Context) string { + return a.description +} + +func (a *TypedChatModelAgent[M]) GetType() string { + return "ChatModel" +} + +// OnSetSubAgents implements OnSubAgents. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func (a *TypedChatModelAgent[M]) OnSetSubAgents(_ context.Context, subAgents []TypedAgent[M]) error { + if atomic.LoadUint32(&a.frozen) == 1 { + return errors.New("agent has been frozen after run") + } + + if len(a.subAgents) > 0 { + return errors.New("agent's sub-agents has already been set") + } + + a.subAgents = subAgents + return nil +} + +// OnSetAsSubAgent implements OnSubAgents. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func (a *TypedChatModelAgent[M]) OnSetAsSubAgent(_ context.Context, parent TypedAgent[M]) error { + if atomic.LoadUint32(&a.frozen) == 1 { + return errors.New("agent has been frozen after run") + } + + if a.parentAgent != nil { + return errors.New("agent has already been set as a sub-agent of another agent") + } + + a.parentAgent = parent + return nil +} + +// OnDisallowTransferToParent implements OnSubAgents. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func (a *TypedChatModelAgent[M]) OnDisallowTransferToParent(_ context.Context) error { + if atomic.LoadUint32(&a.frozen) == 1 { + return errors.New("agent has been frozen after run") + } + + a.disallowTransferToParent = true + + return nil +} + +type ChatModelAgentInterruptInfo struct { + Info *compose.InterruptInfo + Data []byte +} + +func init() { + schema.RegisterName[*ChatModelAgentInterruptInfo]("_eino_adk_chat_model_agent_interrupt_info") +} + +func extractTextContent[M MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return v.Content + case *schema.AgenticMessage: + var texts []string + for _, block := range v.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeAssistantGenText && block.AssistantGenText != nil { + texts = append(texts, block.AssistantGenText.Text) + } + } + return strings.Join(texts, "\n") + default: + return "" + } +} + +func setOutputToSession[M MessageType](ctx context.Context, msg M, msgStream *schema.StreamReader[M], outputKey string) error { + if !isNilMessage(msg) { + AddSessionValue(ctx, outputKey, extractTextContent(msg)) + return nil + } + + concatenated, err := concatMessageStream(msgStream) + if err != nil { + return err + } + + AddSessionValue(ctx, outputKey, extractTextContent(concatenated)) + return nil +} + +func typedErrFunc[M MessageType](err error) typedRunFunc[M] { + return func(ctx context.Context, p *typedRunParams[M]) { + p.generator.Send(&TypedAgentEvent[M]{Err: err}) + } +} + +// ChatModelAgentResumeData holds data that can be provided to a ChatModelAgent during a resume operation +// to modify its behavior. It is provided via the adk.ResumeWithData function. +type ChatModelAgentResumeData struct { + // HistoryModifier is a function that can transform the agent's message history before it is sent to the model. + // This allows for adding new information or context upon resumption. + HistoryModifier func(ctx context.Context, history []Message) []Message +} + +type execContext struct { + instruction string + toolsNodeConf compose.ToolsNodeConfig + returnDirectly map[string]bool + + toolInfos []*schema.ToolInfo + unwrappedTools []tool.BaseTool + + toolSearchTool *schema.ToolInfo // set by BeforeAgent when the model supports native tool search + + rebuildGraph bool // whether needs to instantiate a new graph because of topology changes due to tool modifications + toolUpdated bool // whether needs to pass a compose.WithToolList option to ToolsNode due to tool list change +} + +func (a *TypedChatModelAgent[M]) applyBeforeAgent(ctx context.Context, ec *execContext) (context.Context, *execContext, error) { + runCtx := &ChatModelAgentContext{ + Instruction: ec.instruction, + Tools: cloneSlice(ec.unwrappedTools), + ReturnDirectly: copyMap(ec.returnDirectly), + } + + var err error + for i, handler := range a.handlers { + ctx, runCtx, err = handler.BeforeAgent(ctx, runCtx) + if err != nil { + return ctx, nil, fmt.Errorf("handler[%d] (%T) BeforeAgent failed: %w", i, handler, err) + } + } + + toolsNodeConf := ec.toolsNodeConf + toolsNodeConf.Tools = runCtx.Tools + toolsNodeConf.ToolCallMiddlewares = cloneSlice(ec.toolsNodeConf.ToolCallMiddlewares) + + runtimeEC := &execContext{ + instruction: runCtx.Instruction, + toolsNodeConf: toolsNodeConf, + returnDirectly: runCtx.ReturnDirectly, + toolSearchTool: runCtx.ToolSearchTool, + toolUpdated: true, + rebuildGraph: (len(ec.toolsNodeConf.Tools) == 0 && len(runCtx.Tools) > 0) || + (len(ec.returnDirectly) == 0 && len(runCtx.ReturnDirectly) > 0), + } + + toolInfos, err := genToolInfos(ctx, &runtimeEC.toolsNodeConf) + if err != nil { + return ctx, nil, err + } + + runtimeEC.toolInfos = toolInfos + + return ctx, runtimeEC, nil +} + +func (a *TypedChatModelAgent[M]) applyAfterAgent(ctx context.Context) (context.Context, error) { + if len(a.handlers) == 0 { + return ctx, nil + } + + var state TypedChatModelAgentState[M] + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + state.Messages = st.Messages + state.ToolInfos = st.ToolInfos + state.DeferredToolInfos = st.DeferredToolInfos + return nil + }) + + var err error + for i, handler := range a.handlers { + ctx, err = handler.AfterAgent(ctx, &state) + if err != nil { + return ctx, fmt.Errorf("handler[%d] (%T) AfterAgent failed: %w", i, handler, err) + } + } + return ctx, nil +} + +func (a *TypedChatModelAgent[M]) prepareExecContext(ctx context.Context) (*execContext, error) { + instruction := a.instruction + toolsNodeConf := a.toolsConfig.ToolsNodeConfig + toolsNodeConf.Tools = cloneSlice(a.toolsConfig.Tools) + toolsNodeConf.ToolCallMiddlewares = cloneSlice(a.toolsConfig.ToolCallMiddlewares) + returnDirectly := copyMap(a.toolsConfig.ReturnDirectly) + + transferToAgents := a.subAgents + if a.parentAgent != nil && !a.disallowTransferToParent { + transferToAgents = append(transferToAgents, a.parentAgent) + } + + if len(transferToAgents) > 0 { + transferInstruction := genTransferToAgentInstruction(ctx, transferToAgents) + instruction = concatInstructions(instruction, transferInstruction) + + toolsNodeConf.Tools = append(toolsNodeConf.Tools, &transferToAgent{}) + returnDirectly[TransferToAgentToolName] = true + } + + if a.exit != nil { + toolsNodeConf.Tools = append(toolsNodeConf.Tools, a.exit) + exitInfo, err := a.exit.Info(ctx) + if err != nil { + return nil, err + } + returnDirectly[exitInfo.Name] = true + } + + for _, m := range a.middlewares { + if m.AdditionalInstruction != "" { + instruction = concatInstructions(instruction, m.AdditionalInstruction) + } + toolsNodeConf.Tools = append(toolsNodeConf.Tools, m.AdditionalTools...) + } + + unwrappedTools := cloneSlice(toolsNodeConf.Tools) + + handlerMiddlewares := handlersToToolMiddlewares(a.handlers) + toolsNodeConf.ToolCallMiddlewares = append(toolsNodeConf.ToolCallMiddlewares, handlerMiddlewares...) + + toolInfos, err := genToolInfos(ctx, &toolsNodeConf) + if err != nil { + return nil, err + } + + return &execContext{ + instruction: instruction, + toolsNodeConf: toolsNodeConf, + returnDirectly: returnDirectly, + toolInfos: toolInfos, + unwrappedTools: unwrappedTools, + }, nil +} + +// handleRunFuncError is the common error handler for buildNoToolsRunFunc and buildReActRunFunc. +// It handles compose interrupts (both cancel-triggered and business) +// and generic errors, sending the appropriate event to the generator. +func (a *TypedChatModelAgent[M]) handleRunFuncError( + ctx context.Context, + err error, + cancelCtx *cancelContext, + cancelCtxOwned bool, + store *bridgeStore, + generator *AsyncGenerator[*TypedAgentEvent[M]], +) { + info, ok := compose.ExtractInterruptInfo(err) + if ok { + if cancelCtx != nil { + if !cancelCtx.shouldCancel() { + // Note: there is a benign TOCTOU window here. Between shouldCancel() + // returning false and markDone() executing, a concurrent cancel could + // transition stateRunning→stateCancelling. markDone() then does + // stateCancelling→stateDone, and the cancel func receives + // ErrExecutionEnded (execution finished before cancel took effect). + cancelCtx.markDone() + } + } + + data, existed, sErr := store.Get(ctx, bridgeCheckpointID) + if sErr != nil { + generator.Send(&TypedAgentEvent[M]{AgentName: a.name, Err: fmt.Errorf("failed to get interrupt info: %w", sErr)}) + return + } + if !existed { + generator.Send(&TypedAgentEvent[M]{AgentName: a.name, Err: fmt.Errorf("interrupt occurred but checkpoint data is missing")}) + return + } + + is := FromInterruptContexts(info.InterruptContexts) + event := TypedCompositeInterrupt[M](ctx, info, data, is) + event.Action.Interrupted.Data = &ChatModelAgentInterruptInfo{ + Info: info, + Data: data, + } + event.AgentName = a.name + generator.Send(event) + return + } + + if cancelCtxOwned && cancelCtx != nil { + cancelCtx.markDone() + } + generator.Send(&TypedAgentEvent[M]{Err: err}) +} + +type typedNoToolsInput[M MessageType] struct { + input *TypedAgentInput[M] + instruction string +} + +func appendModelToChain[I, O any, M MessageType](chain *compose.Chain[I, O], m model.BaseModel[M]) { + var zero M + switch any(zero).(type) { + case *schema.Message: + chain.AppendChatModel(any(m).(model.BaseChatModel)) + case *schema.AgenticMessage: + chain.AppendAgenticModel(any(m).(model.AgenticModel)) + } +} + +func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRunFunc[M], error) { + return func(ctx context.Context, p *typedRunParams[M]) { + cancelCtx := p.cancelCtx + ctx = withCancelContext(ctx, cancelCtx) + + wrappedModel := buildModelWrappers(a.model, &typedModelWrapperConfig[M]{ + handlers: a.handlers, + middlewares: a.middlewares, + retryConfig: a.modelRetryConfig, + failoverConfig: a.modelFailoverConfig, + cancelContext: cancelCtx, + }) + + chain := compose.NewChain[typedNoToolsInput[M], M]( + compose.WithGenLocalState(func(ctx context.Context) (state *typedState[M]) { + return &typedState[M]{} + })) + + chain.AppendLambda(compose.InvokableLambda(func(ctx context.Context, in typedNoToolsInput[M]) ([]M, error) { + messages, err := a.genModelInput(ctx, in.instruction, in.input) + if err != nil { + return nil, err + } + if err := compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.Messages = append(st.Messages, messages...) + return nil + }); err != nil { + return nil, err + } + return messages, nil + })) + + appendModelToChain(chain, wrappedModel) + + if len(a.handlers) > 0 { + chain.AppendLambda(compose.InvokableLambda(func(ctx context.Context, msg M) (M, error) { + _, err := a.applyAfterAgent(ctx) + return msg, err + })) + } + + var compileOptions []compose.GraphCompileOption + compileOptions = append(compileOptions, + compose.WithGraphName(a.name), + compose.WithCheckPointStore(p.store), + compose.WithSerializer(&gobSerializer{})) + + if cancelCtx != nil { + var interrupt func(...compose.GraphInterruptOption) + ctx, interrupt = compose.WithGraphInterrupt(ctx) + cancelCtx.setGraphInterruptFunc(interrupt) + } + + r, err := chain.Compile(ctx, compileOptions...) + if err != nil { + p.generator.Send(&TypedAgentEvent[M]{Err: err}) + return + } + + ctx = withTypedChatModelAgentExecCtx(ctx, &typedChatModelAgentExecCtx[M]{ + generator: p.generator, + cancelCtx: cancelCtx, + failoverLastSuccessModel: a.model, + }) + + // Pre-execution cancel check + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode() == CancelImmediate || atomic.LoadInt32(&cancelCtx.escalated) == 1 { + cancelErr, ok := cancelCtx.createAndMarkCancelHandled() + if !ok { + return + } + p.generator.Send(&TypedAgentEvent[M]{Err: cancelErr}) + return + } + } + + in := typedNoToolsInput[M]{input: p.input, instruction: p.instruction} + + var msg M + var msgStream *schema.StreamReader[M] + if p.input.EnableStreaming { + msgStream, err = r.Stream(ctx, in, p.composeOpts...) + } else { + msg, err = r.Invoke(ctx, in, p.composeOpts...) + } + + if err == nil { + if a.outputKey != "" { + err = setOutputToSession(ctx, msg, msgStream, a.outputKey) + if err != nil { + p.generator.Send(&TypedAgentEvent[M]{Err: err}) + } + } else if msgStream != nil { + msgStream.Close() + } + return + } + + a.handleRunFuncError(ctx, err, cancelCtx, p.cancelCtxOwned, p.store, p.generator) + }, nil +} + +func (a *TypedChatModelAgent[M]) buildReActRunFunc(ctx context.Context, bc *execContext) (typedRunFunc[M], error) { + var zero M + switch any(zero).(type) { + case *schema.Message: + return a.buildMessageReActRunFunc(ctx, bc) + case *schema.AgenticMessage: + // single-shot: agentic models handle tool calling internally + return a.buildAgenticReActRunFunc(ctx, bc) + default: + return nil, fmt.Errorf("unsupported message type %T for ReAct run mode", zero) + } +} + +type reactRunInput struct { + input *AgentInput + instruction string +} + +func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc *execContext) (typedRunFunc[M], error) { + // safe: only called when M = *schema.Message (guarded by type switch in buildReActRunFunc) + msgModel := any(a.model).(model.BaseChatModel) + msgHandlers := any(a.handlers).([]ChatModelAgentMiddleware) + genModelInputFn := any(a.genModelInput).(GenModelInput) + msgConf := &reactConfig{ + model: msgModel, + toolsConfig: &bc.toolsNodeConf, + modelWrapperConf: &modelWrapperConfig{ + handlers: msgHandlers, + middlewares: a.middlewares, + retryConfig: any(a.modelRetryConfig).(*ModelRetryConfig), + failoverConfig: any(a.modelFailoverConfig).(*ModelFailoverConfig[*schema.Message]), + toolInfos: bc.toolInfos, + }, + toolsReturnDirectly: bc.returnDirectly, + agentName: a.name, + maxIterations: a.maxIterations, + } + if len(a.handlers) > 0 { + msgAgent := any(a).(*TypedChatModelAgent[*schema.Message]) + msgConf.afterAgentFunc = func(ctx context.Context, msg *schema.Message) (*schema.Message, error) { + _, err := msgAgent.applyAfterAgent(ctx) + return msg, err + } + } + + return func(ctx context.Context, p *typedRunParams[M]) { + mp := any(p).(*typedRunParams[*schema.Message]) + cancelCtx := mp.cancelCtx + msgConf.cancelCtx = cancelCtx + if msgConf.modelWrapperConf != nil { + msgConf.modelWrapperConf.cancelContext = cancelCtx + } + ctx = withCancelContext(ctx, cancelCtx) + + g, err := newReact(ctx, msgConf) + if err != nil { + mp.generator.Send(&AgentEvent{Err: err}) + return + } + + chain := compose.NewChain[reactRunInput, Message](). + AppendLambda( + compose.InvokableLambda(func(ctx context.Context, in reactRunInput) (*reactInput, error) { + messages, genErr := genModelInputFn(ctx, in.instruction, in.input) + if genErr != nil { + return nil, genErr + } + return &reactInput{ + Messages: messages, + }, nil + }), + ). + AppendGraph(g, compose.WithNodeName("ReAct"), compose.WithGraphCompileOptions(compose.WithMaxRunSteps(math.MaxInt))) + + var compileOptions []compose.GraphCompileOption + compileOptions = append(compileOptions, + compose.WithGraphName(a.name), + compose.WithCheckPointStore(mp.store), + compose.WithSerializer(&gobSerializer{}), + compose.WithMaxRunSteps(math.MaxInt)) + + if cancelCtx != nil { + var interrupt func(...compose.GraphInterruptOption) + ctx, interrupt = compose.WithGraphInterrupt(ctx) + cancelCtx.setGraphInterruptFunc(interrupt) + } + + runnable, err_ := chain.Compile(ctx, compileOptions...) + if err_ != nil { + mp.generator.Send(&AgentEvent{Err: err_}) + return + } + + ctx = withTypedChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{ + runtimeReturnDirectly: mp.returnDirectly, + generator: mp.generator, + cancelCtx: cancelCtx, + failoverLastSuccessModel: msgModel, + afterToolCallsHook: mp.afterToolCallsHook, + }) + + // Pre-execution cancel check + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode() == CancelImmediate || atomic.LoadInt32(&cancelCtx.escalated) == 1 { + cancelErr, ok := cancelCtx.createAndMarkCancelHandled() + if !ok { + return + } + mp.generator.Send(&AgentEvent{Err: cancelErr}) + return + } + } + + in := reactRunInput{ + input: mp.input, + instruction: mp.instruction, + } + + var runOpts []compose.Option + runOpts = append(runOpts, mp.composeOpts...) + if a.toolsConfig.EmitInternalEvents { + runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEventGenerator(mp.generator)))) + } + if mp.input.EnableStreaming { + runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEnableStreaming(true)))) + } + + var msg Message + var msgStream MessageStream + if mp.input.EnableStreaming { + msgStream, err_ = runnable.Stream(ctx, in, runOpts...) + } else { + msg, err_ = runnable.Invoke(ctx, in, runOpts...) + } + + if err_ == nil { + if a.outputKey != "" { + err_ = setOutputToSession[*schema.Message](ctx, msg, msgStream, a.outputKey) + if err_ != nil { + mp.generator.Send(&AgentEvent{Err: err_}) + } + } else if msgStream != nil { + msgStream.Close() + } + + return + } + + a.handleRunFuncError(ctx, err_, cancelCtx, mp.cancelCtxOwned, mp.store, p.generator) + }, nil +} + +type agenticReactRunInput struct { + input *TypedAgentInput[*schema.AgenticMessage] + instruction string +} + +func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc *execContext) (typedRunFunc[M], error) { + agenticModel := any(a.model).(model.AgenticModel) + agenticHandlers := any(a.handlers).([]TypedChatModelAgentMiddleware[*schema.AgenticMessage]) + genModelInputFn := any(a.genModelInput).(TypedGenModelInput[*schema.AgenticMessage]) + agenticConf := &agenticReactConfig{ + model: agenticModel, + toolsConfig: &bc.toolsNodeConf, + modelWrapperConf: &typedModelWrapperConfig[*schema.AgenticMessage]{ + handlers: agenticHandlers, + middlewares: a.middlewares, + retryConfig: any(a.modelRetryConfig).(*TypedModelRetryConfig[*schema.AgenticMessage]), + failoverConfig: any(a.modelFailoverConfig).(*ModelFailoverConfig[*schema.AgenticMessage]), + toolInfos: bc.toolInfos, + }, + toolsReturnDirectly: bc.returnDirectly, + agentName: a.name, + maxIterations: a.maxIterations, + } + if len(a.handlers) > 0 { + agenticAgent := any(a).(*TypedChatModelAgent[*schema.AgenticMessage]) + agenticConf.afterAgentFunc = func(ctx context.Context, msg *schema.AgenticMessage) (*schema.AgenticMessage, error) { + _, err := agenticAgent.applyAfterAgent(ctx) + return msg, err + } + } + + return func(ctx context.Context, p *typedRunParams[M]) { + ap := any(p).(*typedRunParams[*schema.AgenticMessage]) + cancelCtx := ap.cancelCtx + agenticConf.cancelCtx = cancelCtx + if agenticConf.modelWrapperConf != nil { + agenticConf.modelWrapperConf.cancelContext = cancelCtx + } + ctx = withCancelContext(ctx, cancelCtx) + + g, err := newAgenticReact(ctx, agenticConf) + if err != nil { + ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err}) + return + } + + chain := compose.NewChain[agenticReactRunInput, *schema.AgenticMessage](). + AppendLambda( + compose.InvokableLambda(func(ctx context.Context, in agenticReactRunInput) (*agenticReactInput, error) { + messages, genErr := genModelInputFn(ctx, in.instruction, in.input) + if genErr != nil { + return nil, genErr + } + return &agenticReactInput{ + Messages: messages, + }, nil + }), + ). + AppendGraph(g, compose.WithNodeName("ReAct"), compose.WithGraphCompileOptions(compose.WithMaxRunSteps(math.MaxInt))) + + var compileOptions []compose.GraphCompileOption + compileOptions = append(compileOptions, + compose.WithGraphName(a.name), + compose.WithCheckPointStore(ap.store), + compose.WithSerializer(&gobSerializer{}), + compose.WithMaxRunSteps(math.MaxInt)) + + if cancelCtx != nil { + var interrupt func(...compose.GraphInterruptOption) + ctx, interrupt = compose.WithGraphInterrupt(ctx) + cancelCtx.setGraphInterruptFunc(interrupt) + } + + runnable, err_ := chain.Compile(ctx, compileOptions...) + if err_ != nil { + ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err_}) + return + } + + ctx = withTypedChatModelAgentExecCtx(ctx, &typedChatModelAgentExecCtx[*schema.AgenticMessage]{ + runtimeReturnDirectly: ap.returnDirectly, + generator: ap.generator, + cancelCtx: cancelCtx, + failoverLastSuccessModel: agenticModel, + afterToolCallsHook: ap.afterToolCallsHook, + }) + + // Pre-execution cancel check + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode() == CancelImmediate || atomic.LoadInt32(&cancelCtx.escalated) == 1 { + cancelErr, ok := cancelCtx.createAndMarkCancelHandled() + if !ok { + return + } + ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: cancelErr}) + return + } + } + + in := agenticReactRunInput{input: ap.input, instruction: ap.instruction} + + var runOpts []compose.Option + runOpts = append(runOpts, ap.composeOpts...) + if a.toolsConfig.EmitInternalEvents { + runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withTypedAgentToolEventGenerator[*schema.AgenticMessage](ap.generator)))) + } + if ap.input.EnableStreaming { + runOpts = append(runOpts, compose.WithToolsNodeOption(compose.WithToolOption(withAgentToolEnableStreaming(true)))) + } + + var msg *schema.AgenticMessage + var msgStream *schema.StreamReader[*schema.AgenticMessage] + if ap.input.EnableStreaming { + msgStream, err_ = runnable.Stream(ctx, in, runOpts...) + } else { + msg, err_ = runnable.Invoke(ctx, in, runOpts...) + } + + if err_ == nil { + if a.outputKey != "" { + err_ = setOutputToSession(ctx, msg, msgStream, a.outputKey) + if err_ != nil { + ap.generator.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err_}) + } + } else if msgStream != nil { + msgStream.Close() + } + + return + } + + a.handleRunFuncError(ctx, err_, cancelCtx, ap.cancelCtxOwned, ap.store, p.generator) + }, nil +} + +func (a *TypedChatModelAgent[M]) buildRunFunc(ctx context.Context) typedRunFunc[M] { + a.once.Do(func() { + ec, err := a.prepareExecContext(ctx) + if err != nil { + a.run = typedErrFunc[M](err) + return + } + + a.exeCtx = ec + + if len(ec.toolsNodeConf.Tools) == 0 { + var run typedRunFunc[M] + run, err = a.buildNoToolsRunFunc(ctx) + if err != nil { + a.run = typedErrFunc[M](err) + return + } + a.run = run + return + } + + var run typedRunFunc[M] + run, err = a.buildReActRunFunc(ctx, ec) + if err != nil { + a.run = typedErrFunc[M](err) + return + } + a.run = run + }) + + atomic.StoreUint32(&a.frozen, 1) + + return a.run +} + +func (a *TypedChatModelAgent[M]) getRunFunc(ctx context.Context) (context.Context, typedRunFunc[M], *execContext, error) { + defaultRun := a.buildRunFunc(ctx) + bc := a.exeCtx + + if bc == nil { + return ctx, defaultRun, bc, nil + } + + if len(a.handlers) == 0 { + runtimeBC := &execContext{ + instruction: bc.instruction, + toolsNodeConf: bc.toolsNodeConf, + returnDirectly: bc.returnDirectly, + toolInfos: bc.toolInfos, + } + return ctx, defaultRun, runtimeBC, nil + } + + ctx, runtimeBC, err := a.applyBeforeAgent(ctx, bc) + if err != nil { + return ctx, nil, nil, err + } + + if !runtimeBC.rebuildGraph { + return ctx, defaultRun, runtimeBC, nil + } + + var tempRun typedRunFunc[M] + if len(runtimeBC.toolsNodeConf.Tools) == 0 { + tempRun, err = a.buildNoToolsRunFunc(ctx) + if err != nil { + return ctx, nil, nil, err + } + } else { + tempRun, err = a.buildReActRunFunc(ctx, runtimeBC) + if err != nil { + return ctx, nil, nil, err + } + } + + return ctx, tempRun, runtimeBC, nil +} + +func (a *TypedChatModelAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + + o := getCommonOptions(nil, opts...) + cancelCtx, cancelCtxOwned := resolveRunCancelContext(ctx, o) + + ctx, run, bc, err := a.getRunFunc(ctx) + if err != nil { + go func() { + if cancelCtxOwned && cancelCtx != nil { + defer cancelCtx.markDone() + } + generator.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("ChatModelAgent getRunFunc error: %w", err)}) + generator.Close() + }() + return iterator + } + + co := getComposeOptions(opts) + co = append(co, compose.WithCheckPointID(bridgeCheckpointID)) + runOps := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...) + + if bc != nil { + if len(bc.toolInfos) > 0 { + co = append(co, compose.WithChatModelOption(model.WithTools(bc.toolInfos))) + } + if bc.toolSearchTool != nil { + co = append(co, compose.WithChatModelOption(model.WithToolSearchTool(bc.toolSearchTool))) + } + if bc.toolUpdated { + co = append(co, compose.WithToolsNodeOption(compose.WithToolList(bc.toolsNodeConf.Tools...))) + } + } + + go func() { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&TypedAgentEvent[M]{Err: e}) + } + + generator.Close() + }() + + var ( + instruction string + returnDirectly map[string]bool + ) + + if bc != nil { + instruction = bc.instruction + returnDirectly = bc.returnDirectly + } + + run(ctx, &typedRunParams[M]{ + input: input, + generator: generator, + store: newBridgeStore(), + instruction: instruction, + returnDirectly: returnDirectly, + cancelCtx: cancelCtx, + cancelCtxOwned: cancelCtxOwned, + composeOpts: co, + afterToolCallsHook: runOps.afterToolCallsHook, + }) + }() + + if cancelCtxOwned { + return wrapIterWithCancelCtx(iterator, cancelCtx) + } + return iterator +} + +func (a *TypedChatModelAgent[M]) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + + o := getCommonOptions(nil, opts...) + cancelCtx, cancelCtxOwned := resolveRunCancelContext(ctx, o) + + ctx, run, bc, err := a.getRunFunc(ctx) + if err != nil { + go func() { + if cancelCtxOwned && cancelCtx != nil { + defer cancelCtx.markDone() + } + generator.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("ChatModelAgent getRunFunc error: %w", err)}) + generator.Close() + }() + return iterator + } + + co := getComposeOptions(opts) + co = append(co, compose.WithCheckPointID(bridgeCheckpointID)) + resumeRunOps := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...) + + if bc != nil { + if len(bc.toolInfos) > 0 { + co = append(co, compose.WithChatModelOption(model.WithTools(bc.toolInfos))) + } + if bc.toolSearchTool != nil { + co = append(co, compose.WithChatModelOption(model.WithToolSearchTool(bc.toolSearchTool))) + } + if bc.toolUpdated { + co = append(co, compose.WithToolsNodeOption(compose.WithToolList(bc.toolsNodeConf.Tools...))) + } + } + + if info == nil { + panic(fmt.Sprintf("ChatModelAgent.Resume: agent '%s' was asked to resume but info is nil", a.Name(ctx))) + } + + if info.InterruptState == nil { + panic(fmt.Sprintf("ChatModelAgent.Resume: agent '%s' was asked to resume but has no state", a.Name(ctx))) + } + + stateByte, ok := info.InterruptState.([]byte) + if !ok { + panic(fmt.Sprintf("ChatModelAgent.Resume: agent '%s' was asked to resume but has invalid interrupt state type: %T", + a.Name(ctx), info.InterruptState)) + } + + // Migrate legacy checkpoints before resume. + // This covers both: + // - v0.7.*: state is stored as a struct wire type (stateV07) under the legacy name. + // - v0.8.0-v0.8.3: state is stored as a GobEncoder payload under the same legacy name and must + // be routed to a GobDecode-compatible compat type via byte-patching. + // The result is re-encoded so the resume path always operates on the current *State. + stateByte, err = preprocessComposeCheckpoint(stateByte) + if err != nil { + go func() { + generator.Send(&TypedAgentEvent[M]{Err: err}) + generator.Close() + }() + return iterator + } + + var historyModifier func(ctx context.Context, history []Message) []Message + if info.ResumeData != nil { + resumeData, ok := info.ResumeData.(*ChatModelAgentResumeData) + if !ok { + panic(fmt.Sprintf("ChatModelAgent.Resume: agent '%s' was asked to resume but has invalid resume data type: %T", + a.Name(ctx), info.ResumeData)) + } + historyModifier = resumeData.HistoryModifier + } + + if historyModifier != nil { + co = append(co, compose.WithStateModifier(func(ctx context.Context, path compose.NodePath, state any) error { + s, ok := state.(*State) + if !ok { + return nil + } + s.Messages = historyModifier(ctx, s.Messages) + return nil + })) + } + + go func() { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&TypedAgentEvent[M]{Err: e}) + } + + generator.Close() + }() + + var ( + instruction string + returnDirectly map[string]bool + ) + + if bc != nil { + instruction = bc.instruction + returnDirectly = bc.returnDirectly + } + + run(ctx, &typedRunParams[M]{ + input: &TypedAgentInput[M]{EnableStreaming: info.EnableStreaming}, + generator: generator, + store: newResumeBridgeStore(bridgeCheckpointID, stateByte), + instruction: instruction, + returnDirectly: returnDirectly, + cancelCtx: cancelCtx, + cancelCtxOwned: cancelCtxOwned, + composeOpts: co, + afterToolCallsHook: resumeRunOps.afterToolCallsHook, + }) + }() + + if cancelCtxOwned { + return wrapIterWithCancelCtx(iterator, cancelCtx) + } + return iterator +} + +func getComposeOptions(opts []AgentRunOption) []compose.Option { + o := GetImplSpecificOptions[chatModelAgentRunOptions](nil, opts...) + var co []compose.Option + if len(o.chatModelOptions) > 0 { + co = append(co, compose.WithChatModelOption(o.chatModelOptions...)) + } + var to []tool.Option + if len(o.toolOptions) > 0 { + to = append(to, o.toolOptions...) + } + for toolName, atos := range o.agentToolOptions { + to = append(to, withAgentToolOptions(toolName, atos)) + } + if len(to) > 0 { + co = append(co, compose.WithToolsNodeOption(compose.WithToolOption(to...))) + } + if o.historyModifier != nil { + co = append(co, compose.WithStateModifier(func(ctx context.Context, path compose.NodePath, state any) error { + s, ok := state.(*State) + if !ok { + return fmt.Errorf("unexpected state type: %T, expected: %T", state, &State{}) + } + s.Messages = o.historyModifier(ctx, s.Messages) + return nil + })) + } + return co +} + +type gobSerializer struct{} + +func (g *gobSerializer) Marshal(v any) ([]byte, error) { + buf := new(bytes.Buffer) + err := gob.NewEncoder(buf).Encode(v) + if err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func (g *gobSerializer) Unmarshal(data []byte, v any) error { + buf := bytes.NewBuffer(data) + return gob.NewDecoder(buf).Decode(v) +} + +// preprocessComposeCheckpoint migrates legacy compose checkpoints to the current format. +// It handles the v0.8.0-v0.8.3 format: +// - gob name "_eino_adk_state_v080_" (already byte-patched by preprocessADKCheckpoint +// from "_eino_adk_react_state"), opaque-bytes wire format → decoded as *stateV080 +// +// v0.7 checkpoints need no migration — State is now a plain struct registered under the +// same gob name, and gob handles missing fields gracefully. +// +// Fast path: if the legacy name is not present, skip entirely. +func preprocessComposeCheckpoint(data []byte) ([]byte, error) { + const lenPrefixedCompatName = "\x15" + stateGobNameV080 + if bytes.Contains(data, []byte(lenPrefixedCompatName)) { + // v0.8.0-v0.8.3: already byte-patched by preprocessADKCheckpoint; decode as *stateV080. + migrated, err := compose.MigrateCheckpointState(data, &gobSerializer{}, func(state any) (any, bool, error) { + sc, ok := state.(*stateV080) + if !ok { + return state, false, nil + } + return stateV080ToState(sc), true, nil + }) + if err != nil { + return nil, fmt.Errorf("failed to migrate v0.8.0-v0.8.3 compose checkpoint: %w", err) + } + return migrated, nil + } + + return data, nil +} diff --git a/adk/chatmodel_retry_test.go b/adk/chatmodel_retry_test.go new file mode 100644 index 0000000..e7ef159 --- /dev/null +++ b/adk/chatmodel_retry_test.go @@ -0,0 +1,3208 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "io" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +var errRetryAble = errors.New("retry-able error") +var errNonRetryAble = errors.New("non-retry-able error") + +var instantBackoff = func(_ context.Context, _ int) time.Duration { return time.Millisecond } + +type agentEvent struct { + Err error + Output *AgentOutput + StreamContent string +} + +func drainAgentEvents(t *testing.T, iterator *AsyncIterator[*AgentEvent]) []agentEvent { + t.Helper() + var events []agentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, agentEvent{Err: event.Err, Output: event.Output}) + } + return events +} + +func drainStreamingAgentEvents(t *testing.T, iterator *AsyncIterator[*AgentEvent]) (events []agentEvent, streamTermErrs []error) { + t.Helper() + for { + event, ok := iterator.Next() + if !ok { + break + } + ae := agentEvent{Err: event.Err, Output: event.Output} + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + var chunks []string + for { + msg, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + streamTermErrs = append(streamTermErrs, recvErr) + break + } + if msg != nil { + chunks = append(chunks, msg.Content) + } + } + ae.StreamContent = strings.Join(chunks, "") + } + } + events = append(events, ae) + } + return events, streamTermErrs +} + +func TestChatModelAgentRetry_NoTools_DirectError_Generate(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 3 { + return nil, errRetryAble + } + return schema.AssistantMessage("Success after retry", nil), nil + }).Times(3) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.Equal(t, "Success after retry", event.Output.MessageOutput.Message.Content) + + _, ok = iterator.Next() + assert.False(t, ok) + assert.Equal(t, int32(3), atomic.LoadInt32(&callCount)) +} + +func TestChatModelAgentRetry_NoTools_DirectError_Stream(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + if count < 2 { + return nil, errRetryAble + } + return schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Success", nil), + }), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + _, ok = iterator.Next() + assert.False(t, ok) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) +} + +type streamErrorModel struct { + callCount int32 + failAtChunk int + maxFailures int + tools []*schema.ToolInfo + returnTool bool +} + +func (m *streamErrorModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("Generated", nil), nil +} + +func (m *streamErrorModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&m.callCount, 1) + + sr, sw := schema.Pipe[*schema.Message](10) + go func() { + defer sw.Close() + for i := 0; i < 5; i++ { + if i == m.failAtChunk && int(count) <= m.maxFailures { + sw.Send(nil, errRetryAble) + return + } + if m.returnTool && i == 0 { + sw.Send(schema.AssistantMessage("", []schema.ToolCall{{ + ID: "call-1", + Function: schema.FunctionCall{Name: "test_tool", Arguments: "{}"}, + }}), nil) + } else { + sw.Send(schema.AssistantMessage("chunk", nil), nil) + } + } + }() + return sr, nil +} + +func (m *streamErrorModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + m.tools = tools + return m, nil +} + +func TestChatModelAgentRetry_StreamError(t *testing.T) { + t.Run("WithTools", func(t *testing.T) { + ctx := context.Background() + + m := &streamErrorModel{ + failAtChunk: 2, + maxFailures: 2, + returnTool: false, + } + + config := &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: m, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&fakeToolForTest{tarCount: 0}}, + }, + }, + } + + agent, err := NewChatModelAgent(ctx, config) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 3, len(events)) + + var streamErrEventCount int + var errs []error + for i, event := range events { + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + sr := event.Output.MessageOutput.MessageStream + for { + msg, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + streamErrEventCount++ + errs = append(errs, err) + t.Logf("event %d: err: %v", i, err) + break + } + t.Logf("event %d: %v", i, msg.Content) + } + } + } + + assert.Equal(t, 2, streamErrEventCount) + assert.Equal(t, 2, len(errs)) + var willRetryErr *WillRetryError + assert.True(t, errors.As(errs[0], &willRetryErr)) + assert.True(t, errors.As(errs[1], &willRetryErr)) + assert.Equal(t, int32(3), atomic.LoadInt32(&m.callCount)) + }) + + t.Run("NoTools", func(t *testing.T) { + ctx := context.Background() + + m := &streamErrorModel{ + failAtChunk: 2, + maxFailures: 2, + returnTool: false, + } + + config := &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: m, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + } + + agent, err := NewChatModelAgent(ctx, config) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 3, len(events)) + + var streamErrEventCount int + var errs []error + for i, event := range events { + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + sr := event.Output.MessageOutput.MessageStream + for { + msg, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + streamErrEventCount++ + errs = append(errs, err) + t.Logf("event %d: err: %v", i, err) + break + } + t.Logf("event %d: %v", i, msg.Content) + } + } + } + + assert.Equal(t, 2, streamErrEventCount) + assert.Equal(t, 2, len(errs)) + var willRetryErr *WillRetryError + assert.True(t, errors.As(errs[0], &willRetryErr)) + assert.True(t, errors.As(errs[1], &willRetryErr)) + assert.Equal(t, int32(3), atomic.LoadInt32(&m.callCount)) + }) +} + +func TestChatModelAgentRetry_WithTools_DirectError_Generate(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 2 { + return nil, errRetryAble + } + return schema.AssistantMessage("Success after retry", nil), nil + }).Times(2) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + fakeTool := &fakeToolForTest{tarCount: 0} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.Equal(t, "Success after retry", event.Output.MessageOutput.Message.Content) + + _, ok = iterator.Next() + assert.False(t, ok) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) +} + +func TestChatModelAgentRetry_NonRetryableError(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, errNonRetryAble).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.True(t, errors.Is(event.Err, errNonRetryAble)) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +type inputCapturingModel struct { + capturedInputs [][]Message +} + +func (m *inputCapturingModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.capturedInputs = append(m.capturedInputs, input) + return schema.AssistantMessage("Response from capturing model", nil), nil +} + +func (m *inputCapturingModel) Stream(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + m.capturedInputs = append(m.capturedInputs, input) + return schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Response from capturing model", nil), + }), nil +} + +func (m *inputCapturingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func TestChatModelAgentRetry_MaxRetriesExhausted(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, errRetryAble).Times(4) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.True(t, errors.Is(event.Err, ErrExceedMaxRetries)) + var retryErr *RetryExhaustedError + assert.True(t, errors.As(event.Err, &retryErr)) + assert.True(t, errors.Is(retryErr.LastErr, errRetryAble)) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestChatModelAgentRetry_BackoffFunction(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var backoffCalls []int + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 3 { + return nil, errRetryAble + } + return schema.AssistantMessage("Success", nil), nil + }).Times(3) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + BackoffFunc: func(ctx context.Context, attempt int) time.Duration { + backoffCalls = append(backoffCalls, attempt) + return time.Millisecond + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + _, ok = iterator.Next() + assert.False(t, ok) + + assert.Equal(t, []int{1, 2}, backoffCalls) +} + +func TestChatModelAgentRetry_NoRetryConfig(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, errRetryAble).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent without retry config", + Instruction: "You are a helpful assistant.", + Model: cm, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.True(t, errors.Is(event.Err, errRetryAble)) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestChatModelAgentRetry_WithTools_NonRetryAbleStreamError(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, errNonRetryAble).Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + fakeTool := &fakeToolForTest{tarCount: 0} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.True(t, errors.Is(event.Err, errNonRetryAble)) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +type nonRetryAbleStreamErrorModel struct { + tools []*schema.ToolInfo +} + +func (m *nonRetryAbleStreamErrorModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("Generated", nil), nil +} + +func (m *nonRetryAbleStreamErrorModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("chunk1", nil), nil) + sw.Send(nil, errNonRetryAble) + }() + return sr, nil +} + +func (m *nonRetryAbleStreamErrorModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + m.tools = tools + return m, nil +} + +func TestChatModelAgentRetry_NoTools_NonRetryAbleStreamError(t *testing.T) { + ctx := context.Background() + + m := &nonRetryAbleStreamErrorModel{} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent for retry functionality", + Instruction: "You are a helpful assistant.", + Model: m, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 2, len(events)) + + event0 := events[0] + assert.NotNil(t, event0.Output) + assert.NotNil(t, event0.Output.MessageOutput) + assert.True(t, event0.Output.MessageOutput.IsStreaming) + sr := event0.Output.MessageOutput.MessageStream + var streamErr error + for { + _, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + streamErr = err + break + } + } + assert.NotNil(t, streamErr) + assert.True(t, errors.Is(streamErr, errNonRetryAble), "Stream error should be the original error") + + event1 := events[1] + assert.NotNil(t, event1.Err) + assert.True(t, errors.Is(event1.Err, errNonRetryAble)) +} + +func TestDefaultBackoff(t *testing.T) { + ctx := context.Background() + + d1 := defaultBackoff(ctx, 1) + d2 := defaultBackoff(ctx, 2) + d3 := defaultBackoff(ctx, 3) + + t.Logf("Backoff delays: d1=%v, d2=%v, d3=%v", d1, d2, d3) + + assert.True(t, d1 >= 100*time.Millisecond && d1 < 150*time.Millisecond, + "First retry should be ~100ms + jitter (0-50ms), got %v", d1) + assert.True(t, d2 >= 200*time.Millisecond && d2 < 300*time.Millisecond, + "Second retry should be ~200ms + jitter (0-100ms), got %v", d2) + assert.True(t, d3 >= 400*time.Millisecond && d3 < 600*time.Millisecond, + "Third retry should be ~400ms + jitter (0-200ms), got %v", d3) + + d10 := defaultBackoff(ctx, 10) + t.Logf("Backoff delay for attempt 10: %v", d10) + assert.True(t, d10 >= 10*time.Second && d10 <= 15*time.Second, + "Delay should be capped at 10s + jitter (0-5s), got %v", d10) + + d100 := defaultBackoff(ctx, 100) + t.Logf("Backoff delay for attempt 100: %v", d100) + assert.True(t, d100 >= 10*time.Second && d100 <= 15*time.Second, + "Delay should still be capped at 10s + jitter for very high attempts, got %v", d100) +} + +type customError struct { + code int + msg string +} + +func (e *customError) Error() string { + return e.msg +} + +func TestWillRetryError_Unwrap(t *testing.T) { + originalErr := &customError{code: 500, msg: "internal error"} + willRetry := &WillRetryError{ErrStr: originalErr.Error(), RetryAttempt: 1, err: originalErr} + + assert.True(t, errors.Is(willRetry, originalErr)) + + var targetErr *customError + assert.True(t, errors.As(willRetry, &targetErr)) + assert.Equal(t, 500, targetErr.code) + assert.Equal(t, "internal error", targetErr.msg) +} + +func TestChatModelAgentRetry_DefaultIsRetryAble(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 2 { + return nil, errors.New("any error should be retried") + } + return schema.AssistantMessage("Success", nil), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTestAgent", + Description: "Test agent with default IsRetryAble", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.Equal(t, "Success", event.Output.MessageOutput.Message.Content) + + _, ok = iterator.Next() + assert.False(t, ok) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) +} + +func TestSequentialWorkflow_RetryAbleStreamError_SuccessfulRetry(t *testing.T) { + ctx := context.Background() + + retryModel := &streamErrorModel{ + failAtChunk: 2, + maxFailures: 2, + } + + agentA, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AgentA", + Description: "Agent A with retry that emits stream errors then succeeds", + Instruction: "You are agent A.", + Model: retryModel, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + capturingModel := &inputCapturingModel{} + agentB, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AgentB", + Description: "Agent B that captures input", + Instruction: "You are agent B.", + Model: capturingModel, + }) + assert.NoError(t, err) + + sequentialAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "SequentialAgent", + Description: "Sequential agent A->B", + SubAgents: []Agent{agentA, agentB}, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + ctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input) + iterator := sequentialAgent.Run(ctx, input) + + var willRetryErrCount int + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + sr := event.Output.MessageOutput.MessageStream + for { + _, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + var retryErr *WillRetryError + if errors.As(err, &retryErr) { + willRetryErrCount++ + } + break + } + } + } + } + + assert.Equal(t, 2, willRetryErrCount, "End-user should receive 2 WillRetryError events") + assert.Equal(t, 1, len(capturingModel.capturedInputs), "Agent B should be called exactly once") + + successorInput := capturingModel.capturedInputs[0] + var hasSuccessfulMessage bool + for _, msg := range successorInput { + if strings.Contains(msg.Content, "chunkchunkchunkchunkchunk") { + hasSuccessfulMessage = true + break + } + } + assert.True(t, hasSuccessfulMessage, "Agent B should receive the successful message from Agent A") + + for _, msg := range successorInput { + assert.NotContains(t, msg.Content, "retry-able error", "Agent B should not receive failed stream messages") + } +} + +type streamErrorModelNoRetry struct { + callCount int32 +} + +func (m *streamErrorModelNoRetry) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("Generated", nil), nil +} + +func (m *streamErrorModelNoRetry) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m.callCount, 1) + sr, sw := schema.Pipe[*schema.Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("chunk1", nil), nil) + sw.Send(schema.AssistantMessage("chunk2", nil), nil) + sw.Send(nil, errRetryAble) + }() + return sr, nil +} + +func (m *streamErrorModelNoRetry) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func TestSequentialWorkflow_NonRetryAbleStreamError_StopsFlow(t *testing.T) { + ctx := context.Background() + + nonRetryModel := &nonRetryAbleStreamErrorModel{} + + agentA, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AgentA", + Description: "Agent A that emits non-retryable stream error", + Instruction: "You are agent A.", + Model: nonRetryModel, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + }) + assert.NoError(t, err) + + capturingModel := &inputCapturingModel{} + agentB, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AgentB", + Description: "Agent B that captures input", + Instruction: "You are agent B.", + Model: capturingModel, + }) + assert.NoError(t, err) + + sequentialAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "SequentialAgent", + Description: "Sequential agent A->B", + SubAgents: []Agent{agentA, agentB}, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + ctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input) + iterator := sequentialAgent.Run(ctx, input) + + var streamErrFound bool + var finalErrEvent *AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Err != nil { + finalErrEvent = event + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + sr := event.Output.MessageOutput.MessageStream + for { + _, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + streamErrFound = true + assert.True(t, errors.Is(err, errNonRetryAble), "Stream error should be the original error") + break + } + } + } + } + + assert.True(t, streamErrFound, "End-user should receive stream error") + assert.NotNil(t, finalErrEvent, "Should receive a final error event") + assert.True(t, errors.Is(finalErrEvent.Err, errNonRetryAble), "Final error should be the non-retryable error") + assert.Equal(t, 0, len(capturingModel.capturedInputs), "Agent B should NOT be called due to error") +} + +func TestSequentialWorkflow_NoRetryConfig_StreamError_StopsFlow(t *testing.T) { + ctx := context.Background() + + noRetryModel := &streamErrorModelNoRetry{} + + agentA, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AgentA", + Description: "Agent A without retry config that emits stream error", + Instruction: "You are agent A.", + Model: noRetryModel, + }) + assert.NoError(t, err) + + capturingModel := &inputCapturingModel{} + agentB, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AgentB", + Description: "Agent B that captures input", + Instruction: "You are agent B.", + Model: capturingModel, + }) + assert.NoError(t, err) + + sequentialAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "SequentialAgent", + Description: "Sequential agent A->B", + SubAgents: []Agent{agentA, agentB}, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + ctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input) + iterator := sequentialAgent.Run(ctx, input) + + var streamErrFound bool + var finalErrEvent *AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Err != nil { + finalErrEvent = event + } + if event.Output != nil && event.Output.MessageOutput != nil && event.Output.MessageOutput.IsStreaming { + sr := event.Output.MessageOutput.MessageStream + for { + _, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + streamErrFound = true + assert.True(t, errors.Is(err, errRetryAble), "Stream error should be the original error") + break + } + } + } + } + + assert.True(t, streamErrFound, "End-user should receive stream error") + assert.NotNil(t, finalErrEvent, "Should receive a final error event") + assert.True(t, errors.Is(finalErrEvent.Err, errRetryAble), "Final error should be the original error") + assert.Equal(t, 0, len(capturingModel.capturedInputs), "Agent B should NOT be called due to error") + assert.Equal(t, int32(1), atomic.LoadInt32(&noRetryModel.callCount), "Model should only be called once (no retry)") +} + +// failThenToolCallStreamModel is a ChatModel that: +// - First Stream() call: yields a partial chunk then fails with a retryable error mid-stream. +// - Second Stream() call (retry): yields a tool-call message (success). +// - Third Generate() call (after tool result): yields a final assistant message. +// +// This exercises the path where the eventSenderModel copies the first stream, +// wraps its error as WillRetryError, and sends it as an event to the session. +// The retryModelWrapper then retries, gets a clean stream with a tool call, +// the tool interrupts, and checkpoint save needs to gob-encode the session +// (which still contains the unconsumed WillRetryError event stream). +type failThenToolCallStreamModel struct { + streamCallCount int32 + genCallCount int32 +} + +func (m *failThenToolCallStreamModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m.genCallCount, 1) + return schema.AssistantMessage("final answer", nil), nil +} + +func (m *failThenToolCallStreamModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&m.streamCallCount, 1) + + sr, sw := schema.Pipe[*schema.Message](10) + go func() { + defer sw.Close() + if count == 1 { + // First call: yield a partial chunk then fail. + sw.Send(schema.AssistantMessage("partial", nil), nil) + sw.Send(nil, errRetryAble) + return + } + // Second call (retry): yield a tool-call message. + sw.Send(schema.AssistantMessage("", []schema.ToolCall{{ + ID: "call-1", + Function: schema.FunctionCall{ + Name: "interrupt_tool", + Arguments: `{}`, + }, + }}), nil) + }() + return sr, nil +} + +func (m *failThenToolCallStreamModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +// interruptToolForRetryTest is a tool that always interrupts. +type interruptToolForRetryTest struct{} + +func (t *interruptToolForRetryTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "interrupt_tool", + Desc: "tool that interrupts", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string"}, + }), + }, nil +} + +func (t *interruptToolForRetryTest) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + return "", tool.Interrupt(ctx, "interrupted by tool") +} + +// TestCheckpointSave_WillRetryError_StreamNotConsumed verifies that checkpoint +// saving succeeds when the session contains an event with an unconsumed stream +// that ends with WillRetryError. +// +// Scenario: +// 1. ChatModelAgent with retry (MaxRetries=1) and a tool that always interrupts +// 2. Model.Stream() #1 yields "partial" then errRetryAble mid-stream +// → eventSenderModel copies the stream, wraps the error as WillRetryError, +// sends the event to the session (stream NOT consumed by anyone yet) +// → retryModelWrapper detects error on its copy, retries +// 3. Model.Stream() #2 succeeds with a tool-call message +// 4. Tool executes → interrupts +// 5. Runner.handleIter sees the interrupt → saveCheckPoint → gob encodes runSession +// 6. The session has the WillRetryError event with an unconsumed stream +// → agentEventWrapper.GobEncode proactively consumes the stream via +// getMessageFromWrappedEvent, so MessageVariant.GobEncode sees an error-free +// array and succeeds +func TestCheckpointSave_WillRetryError_StreamNotConsumed(t *testing.T) { + ctx := context.Background() + + mdl := &failThenToolCallStreamModel{} + itool := &interruptToolForRetryTest{} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Agent for checkpoint stream error test", + Instruction: "You are a test agent.", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{itool}, + }, + }, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { + return errors.Is(err, errRetryAble) + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + store := newMyStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, + []Message{schema.UserMessage("hello")}, + WithCheckPointID("ckpt-1"), + ) + + for { + event, ok := iter.Next() + if !ok { + break + } + + if event.Err != nil { + t.Logf("event error: %v", event.Err) + } + } + + // Verify the checkpoint was saved successfully. + _, exists, _ := store.Get(ctx, "ckpt-1") + assert.True(t, exists, "checkpoint should be saved successfully; "+ + "if this fails, the WillRetryError stream in the session caused gob encoding to fail") + + // Sanity: the model should have been called twice for Stream (fail + retry). + assert.Equal(t, int32(2), atomic.LoadInt32(&mdl.streamCallCount), + "model should be called twice: first fail, then retry success") +} + +func TestChatModelAgentRetry_ShouldRetry_RejectMessage_Stream(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + r, w := schema.Pipe[*schema.Message](1) + go func() { + if count < 2 { + _ = w.Send(schema.AssistantMessage("bad stream content", nil), nil) + } else { + _ = w.Send(schema.AssistantMessage("good stream content", nil), nil) + } + w.Close() + }() + return r, nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ShouldRetryStreamTestAgent", + Description: "Test ShouldRetry message rejection in stream mode", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{Retry: true} + } + if retryCtx.OutputMessage != nil && strings.Contains(retryCtx.OutputMessage.Content, "bad") { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events, _ := drainStreamingAgentEvents(t, iterator) + var foundGoodContent bool + for _, e := range events { + if e.StreamContent == "good stream content" { + foundGoodContent = true + } + } + require.True(t, foundGoodContent, "should have received good stream content") + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) +} + +func TestShouldRetry_Generate(t *testing.T) { + t.Run("RetryContext_Fields", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 2 { + return schema.AssistantMessage("bad", nil), nil + } + return schema.AssistantMessage("good", nil), nil + }).Times(2) + + var capturedContexts []*RetryContext + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryContextFieldsAgent", + Description: "Test that RetryContext fields are correctly populated", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + capturedContexts = append(capturedContexts, retryCtx) + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + for { + event, ok := iterator.Next() + if !ok { + break + } + _ = event + } + + assert.Len(t, capturedContexts, 2, "ShouldRetry should be called twice") + + assert.Equal(t, 1, capturedContexts[0].RetryAttempt) + assert.Len(t, capturedContexts[0].InputMessages, 2) + assert.True(t, len(capturedContexts[0].Options) > 0, "should have default options") + assert.Equal(t, "bad", capturedContexts[0].OutputMessage.Content) + assert.Nil(t, capturedContexts[0].Err) + + assert.Equal(t, 2, capturedContexts[1].RetryAttempt) + assert.Equal(t, "good", capturedContexts[1].OutputMessage.Content) + assert.Nil(t, capturedContexts[1].Err) + }) + + t.Run("RewriteError_OnMessage", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("unrecoverable bad message", nil), nil).Times(1) + + fatalErr := errors.New("fatal: unrecoverable model output") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RewriteErrorTestAgent", + Description: "Test ShouldRetry RewriteError on message", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && strings.Contains(retryCtx.OutputMessage.Content, "unrecoverable") { + return &RetryDecision{ + Retry: false, + RewriteError: fatalErr, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + events := drainAgentEvents(t, iterator) + require.NotEmpty(t, events) + foundErr := false + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, fatalErr) { + foundErr = true + } + } + require.True(t, foundErr, "should have received the fatal rewrite error") + }) + + t.Run("RewriteError_OnError", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + origErr := errors.New("original transient error") + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, origErr).Times(1) + + wrappedErr := errors.New("wrapped: original transient error with more context") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RewriteErrorOnErrorTestAgent", + Description: "Test ShouldRetry RewriteError replacing original error", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{ + Retry: false, + RewriteError: wrappedErr, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + events := drainAgentEvents(t, iterator) + require.NotEmpty(t, events) + foundErr := false + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, wrappedErr) { + foundErr = true + } + } + require.True(t, foundErr, "should have received the wrapped rewrite error") + }) + + t.Run("AdditionalOptions", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + var capturedOpts [][]model.Option + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + capturedOpts = append(capturedOpts, opts) + if count < 2 { + return nil, errRetryAble + } + return schema.AssistantMessage("success", nil), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "AdditionalOptionsTestAgent", + Description: "Test ShouldRetry AdditionalOptions", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{ + Retry: true, + AdditionalOptions: []model.Option{model.WithMaxTokens(8192)}, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + assert.Equal(t, 2, len(capturedOpts)) + assert.Equal(t, len(capturedOpts[0])+1, len(capturedOpts[1])) + }) + + t.Run("ModifiedInputMessages_NoPersist", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + var capturedInputs [][]*schema.Message + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + inputCopy := make([]*schema.Message, len(input)) + copy(inputCopy, input) + capturedInputs = append(capturedInputs, inputCopy) + if count < 2 { + return nil, errRetryAble + } + return schema.AssistantMessage("success", nil), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ModifiedInputNoPersistAgent", + Description: "Test ShouldRetry ModifiedInputMessages without persistence", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{ + Retry: true, + ModifiedInputMessages: []*schema.Message{ + schema.SystemMessage("compressed instruction"), + schema.UserMessage("Hello"), + }, + PersistModifiedInputMessages: false, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + assert.Equal(t, 2, len(capturedInputs)) + assert.Equal(t, "compressed instruction", capturedInputs[1][0].Content, "second call should use modified input") + }) + + t.Run("Backoff", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 2 { + return nil, errRetryAble + } + return schema.AssistantMessage("success", nil), nil + }).Times(2) + + customBackoff := 50 * time.Millisecond + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "BackoffTestAgent", + Description: "Test ShouldRetry custom Backoff in decision", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{ + Retry: true, + Backoff: customBackoff, + } + } + return &RetryDecision{Retry: false} + }, + }, + }) + assert.NoError(t, err) + + start := time.Now() + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + elapsed := time.Since(start) + assert.True(t, elapsed >= customBackoff && elapsed < customBackoff+200*time.Millisecond, "expected backoff ~%v, got %v", customBackoff, elapsed) + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + }) + + t.Run("SuppressFlag_Rejected_NoEvent", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count == 1 { + return schema.AssistantMessage("bad", nil), nil + } + return schema.AssistantMessage("good", nil), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "SuppressRejected", + Description: "Test suppress flag rejects first then accepts", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + var msgEvents []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + msgEvents = append(msgEvents, event) + } + } + assert.Equal(t, 1, len(msgEvents), "should have exactly 1 message event (suppressed rejected)") + assert.Equal(t, "good", msgEvents[0].Output.MessageOutput.Message.Content) + }) + + t.Run("SuppressFlag_AllRejected_NoEvents", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("always bad", nil), nil).Times(3) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "SuppressAllRejected", + Description: "Test suppress flag all rejected no events", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: true} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + events := drainAgentEvents(t, iterator) + var msgEventCount int + var foundExhaustedErr bool + for _, e := range events { + if e.Output != nil && e.Output.MessageOutput != nil { + msgEventCount++ + } + if e.Err != nil && errors.Is(e.Err, ErrExceedMaxRetries) { + foundExhaustedErr = true + } + } + assert.Equal(t, 0, msgEventCount, "no message events should be emitted when all are rejected") + require.True(t, foundExhaustedErr, "final event should have RetryExhaustedError") + }) + + t.Run("SuppressFlag_Accepted_FirstAttempt", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("perfect", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "SuppressAcceptedFirst", + Description: "Test suppress flag accepted first attempt", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + var msgEvents []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + msgEvents = append(msgEvents, event) + } + } + assert.Equal(t, 1, len(msgEvents), "should have exactly 1 event") + assert.Equal(t, "perfect", msgEvents[0].Output.MessageOutput.Message.Content) + }) + + t.Run("ContextCanceled_DuringSleep", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&callCount, 1) + return nil, errors.New("transient") + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ContextCancelDuringSleep", + Description: "Test context cancellation during backoff sleep", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 5, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: true} + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 10 * time.Second }, + }, + }) + require.NoError(t, err) + + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + start := time.Now() + iterator := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + }) + events := drainAgentEvents(t, iterator) + elapsed := time.Since(start) + + require.True(t, elapsed < 2*time.Second, "should not block for full backoff; elapsed: %v", elapsed) + assert.Equal(t, int32(1), atomic.LoadInt32(&callCount)) + + var foundCtxErr bool + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, context.Canceled) { + foundCtxErr = true + } + } + require.True(t, foundCtxErr, "should have received context.Canceled error") + }) +} + +func TestShouldRetry_Stream(t *testing.T) { + t.Run("ErrorRetry", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + streamErr := errors.New("stream unavailable") + var callCount int32 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + if count < 2 { + return nil, streamErr + } + r, w := schema.Pipe[*schema.Message](1) + go func() { + _ = w.Send(schema.AssistantMessage("recovered stream", nil), nil) + w.Close() + }() + return r, nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamErrorRetryAgent", + Description: "Test ShouldRetry when Stream returns error (nil stream)", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events, _ := drainStreamingAgentEvents(t, iterator) + var foundContent bool + for _, e := range events { + if e.StreamContent == "recovered stream" { + foundContent = true + } + } + require.True(t, foundContent, "should have received recovered stream content after error retry") + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + }) + + t.Run("ErrorRewrite", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + streamErr := errors.New("model overloaded") + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, streamErr).Times(1) + + fatalErr := errors.New("fatal: model overloaded, aborting") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamErrorRewriteAgent", + Description: "Test ShouldRetry RewriteError when Stream returns error", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil && strings.Contains(retryCtx.Err.Error(), "overloaded") { + return &RetryDecision{ + Retry: false, + RewriteError: fatalErr, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events := drainAgentEvents(t, iterator) + require.NotEmpty(t, events) + foundErr := false + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, fatalErr) { + foundErr = true + } + } + require.True(t, foundErr, "should have received the fatal rewrite error from stream") + }) + + t.Run("RewriteError_OnMessage", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + _ = w.Send(schema.AssistantMessage("hallucinated garbage output", nil), nil) + w.Close() + }() + return r, nil + }).Times(1) + + fatalErr := errors.New("fatal: hallucinated output detected") + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamRewriteOnMessageAgent", + Description: "Test ShouldRetry RewriteError on successful stream with bad content", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && strings.Contains(retryCtx.OutputMessage.Content, "hallucinated") { + return &RetryDecision{ + Retry: false, + RewriteError: fatalErr, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events := drainAgentEvents(t, iterator) + require.NotEmpty(t, events) + foundErr := false + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, fatalErr) { + foundErr = true + } + } + require.True(t, foundErr, "should have received fatal rewrite error from stream message inspection") + }) + + t.Run("PartialStreamError", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + partialErr := errors.New("connection reset mid-stream") + var callCount int32 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + r, w := schema.Pipe[*schema.Message](1) + go func() { + _ = w.Send(schema.AssistantMessage("partial chunk", nil), nil) + if count < 2 { + w.Send(nil, partialErr) + } else { + _ = w.Send(schema.AssistantMessage(" complete", nil), nil) + w.Close() + } + }() + return r, nil + }).Times(2) + + var capturedContexts []*RetryContext + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamPartialErrorAgent", + Description: "Test ShouldRetry when stream has partial content then error", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + capturedContexts = append(capturedContexts, retryCtx) + if retryCtx.Err != nil { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + for { + _, err := mo.MessageStream.Recv() + if err != nil { + break + } + } + } + } + } + + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + assert.Equal(t, 2, len(capturedContexts)) + assert.NotNil(t, capturedContexts[0].Err, "first attempt should have stream error") + assert.NotNil(t, capturedContexts[0].OutputMessage, "first attempt should have partial message despite error") + assert.Equal(t, "partial chunk", capturedContexts[0].OutputMessage.Content) + }) + + t.Run("ModifiedInputsAndOptions_WithPersist", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + var capturedInputs [][]*schema.Message + var capturedOptLens []int + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + inputCopy := make([]*schema.Message, len(input)) + copy(inputCopy, input) + capturedInputs = append(capturedInputs, inputCopy) + capturedOptLens = append(capturedOptLens, len(opts)) + + r, w := schema.Pipe[*schema.Message](1) + go func() { + if count < 2 { + _ = w.Send(schema.AssistantMessage("too long response exceeds limit", nil), nil) + } else { + _ = w.Send(schema.AssistantMessage("good response", nil), nil) + } + w.Close() + }() + return r, nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamModifiedInputsPersistAgent", + Description: "Test ShouldRetry with ModifiedInputMessages (persist) and AdditionalOptions in stream mode", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && strings.Contains(retryCtx.OutputMessage.Content, "too long") { + return &RetryDecision{ + Retry: true, + ModifiedInputMessages: []*schema.Message{ + schema.SystemMessage("compressed instruction"), + schema.UserMessage("summarized history"), + }, + PersistModifiedInputMessages: true, + AdditionalOptions: []model.Option{model.WithMaxTokens(16384)}, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events, _ := drainStreamingAgentEvents(t, iterator) + var foundGood bool + for _, e := range events { + if e.StreamContent == "good response" { + foundGood = true + } + } + + require.True(t, foundGood, "should have received good response after retry with modified inputs") + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + assert.Equal(t, 2, len(capturedInputs)) + assert.Equal(t, "compressed instruction", capturedInputs[1][0].Content, "second call should use modified input") + assert.Equal(t, "summarized history", capturedInputs[1][1].Content) + assert.Equal(t, capturedOptLens[0]+1, capturedOptLens[1]) + }) + + t.Run("VerdictSignal_CleanStream_Rejected", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + if count == 1 { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("bad", nil)}), nil + } + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("good", nil)}), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "VerdictCleanRejected", + Description: "Test verdict signal on clean stream rejected", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + var streamEvents []int + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + idx := len(streamEvents) + streamEvents = append(streamEvents, idx) + var lastErr error + for { + _, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + lastErr = recvErr + break + } + } + if idx == 0 { + var willRetryErr *WillRetryError + assert.True(t, errors.As(lastErr, &willRetryErr), "first stream should end with WillRetryError") + } else { + assert.ErrorIs(t, lastErr, io.EOF, "second stream should end with io.EOF") + } + } + } + } + assert.Equal(t, 2, len(streamEvents), "should have exactly 2 stream events") + assert.Equal(t, int32(2), atomic.LoadInt32(&callCount)) + }) + + t.Run("VerdictSignal_StreamError_Rejected", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + streamErr := errors.New("mid-stream error") + var callCount int32 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&callCount, 1) + if count == 1 { + r, w := schema.Pipe[*schema.Message](1) + go func() { + _ = w.Send(schema.AssistantMessage("partial", nil), nil) + w.Send(nil, streamErr) + }() + return r, nil + } + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("good", nil)}), nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "VerdictStreamErrorRejected", + Description: "Test verdict signal on stream error rejected", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + var firstEventHasWillRetry bool + var eventCount int + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + eventCount++ + for { + _, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + if eventCount == 1 { + var willRetryErr *WillRetryError + if errors.As(recvErr, &willRetryErr) { + firstEventHasWillRetry = true + } + } + break + } + } + } + } + } + assert.True(t, firstEventHasWillRetry, "first event stream should end with WillRetryError via errWrapper path") + assert.Equal(t, 2, eventCount, "should have 2 stream events") + }) + + t.Run("VerdictSignal_Accepted_FirstAttempt", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("perfect", nil)}), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "VerdictAcceptedFirst", + Description: "Test verdict signal accepted first attempt", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + var eventCount int + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + eventCount++ + var lastErr error + for { + _, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + lastErr = recvErr + break + } + } + assert.ErrorIs(t, lastErr, io.EOF, "accepted stream should end with io.EOF") + } + } + } + assert.Equal(t, 1, eventCount, "should have exactly 1 event") + }) + + t.Run("VerdictSignal_AllRejected_Exhausted", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("always bad", nil)}), nil + }).Times(3) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "VerdictAllRejected", + Description: "Test verdict signal all rejected exhausted", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: true} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events, streamTermErrs := drainStreamingAgentEvents(t, iterator) + var willRetryCount int + var foundExhaustedErr bool + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, ErrExceedMaxRetries) { + foundExhaustedErr = true + } + } + for _, termErr := range streamTermErrs { + var willRetryErr *WillRetryError + if errors.As(termErr, &willRetryErr) { + willRetryCount++ + } + } + assert.Equal(t, 3, willRetryCount, "all 3 stream events should end with WillRetryError") + require.True(t, foundExhaustedErr, "final error should be RetryExhaustedError") + }) + + t.Run("ShouldRetry_Panics_VerdictStillSent", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("trigger panic", nil)}), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ShouldRetryPanicsAgent", + Description: "Test that ShouldRetry panic sends verdict signal and does not deadlock", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + panic("deliberate panic in ShouldRetry") + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + + done := make(chan struct{}) + var events []agentEvent + go func() { + defer close(done) + iterator := agent.Run(ctx, input) + events = drainAgentEvents(t, iterator) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("test deadlocked — verdict signal was not sent after ShouldRetry panic") + } + require.NotEmpty(t, events) + var foundPanicErr bool + for _, e := range events { + if e.Err != nil && strings.Contains(e.Err.Error(), "panic") { + foundPanicErr = true + } + } + assert.True(t, foundPanicErr, "should have received a panic error event") + }) +} + +func TestErrStreamCanceled(t *testing.T) { + t.Run("Stream_ShouldRetry_NeverRetried", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + _ = w.Send(schema.AssistantMessage("partial", nil), nil) + w.Send(nil, ErrStreamCanceled) + }() + return r, nil + }).Times(1) + + var shouldRetryCalled int32 + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamCanceledShouldRetry", + Description: "Test ErrStreamCanceled never retried with ShouldRetry", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + atomic.AddInt32(&shouldRetryCalled, 1) + return &RetryDecision{Retry: true} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + for { + _, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + break + } + } + } + } + } + assert.Equal(t, int32(0), atomic.LoadInt32(&shouldRetryCalled), "ShouldRetry should never be called for ErrStreamCanceled") + }) + + t.Run("Stream_LegacyIsRetryAble_NeverRetried", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + _ = w.Send(schema.AssistantMessage("partial", nil), nil) + w.Send(nil, ErrStreamCanceled) + }() + return r, nil + }).Times(1) + + var isRetryAbleCalled int32 + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamCanceledLegacy", + Description: "Test ErrStreamCanceled never retried with legacy IsRetryAble", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(_ context.Context, err error) bool { + atomic.AddInt32(&isRetryAbleCalled, 1) + return true + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + for { + _, recvErr := mo.MessageStream.Recv() + if recvErr != nil { + break + } + } + } + } + } + assert.Equal(t, int32(0), atomic.LoadInt32(&isRetryAbleCalled), "IsRetryAble should never be called for ErrStreamCanceled") + }) + + t.Run("Generate_ShouldRetry_NeverRetried", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, ErrStreamCanceled).Times(1) + + var shouldRetryCalled int32 + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "GenCanceledShouldRetry", + Description: "Test ErrStreamCanceled in Generate never retried", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + atomic.AddInt32(&shouldRetryCalled, 1) + return &RetryDecision{Retry: true} + }, + BackoffFunc: instantBackoff, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + for { + _, ok := iterator.Next() + if !ok { + break + } + } + assert.Equal(t, int32(0), atomic.LoadInt32(&shouldRetryCalled), "ShouldRetry should never be called for ErrStreamCanceled") + }) +} + +func TestAttack_ShouldRetry_NilDecisionOnEveryCall(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("ok", nil), nil).Times(1) + + var shouldRetryCalls int32 + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "NilDecisionAgent", + Description: "ShouldRetry always returns nil — should accept on first call", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + atomic.AddInt32(&shouldRetryCalls, 1) + return nil + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + iterator := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("Hello")}}) + events := drainAgentEvents(t, iterator) + + require.NotEmpty(t, events) + assert.Equal(t, int32(1), atomic.LoadInt32(&shouldRetryCalls)) + var foundOK bool + for _, e := range events { + if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.Message != nil { + if e.Output.MessageOutput.Message.Content == "ok" { + foundOK = true + } + } + } + assert.True(t, foundOK, "nil decision should accept the message as-is") +} + +func TestAttack_ShouldRetry_MaxRetriesZero_RejectFirstAttempt(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("bad", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "MaxZeroRejectAgent", + Description: "MaxRetries=0 with ShouldRetry rejecting — should exhaust immediately", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 0, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: true} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + iterator := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("Hello")}}) + events := drainAgentEvents(t, iterator) + + var foundExhausted bool + for _, e := range events { + if e.Err != nil { + var exhaustedErr *RetryExhaustedError + if errors.As(e.Err, &exhaustedErr) { + foundExhausted = true + } + } + } + assert.True(t, foundExhausted, "MaxRetries=0 with Retry:true should produce RetryExhaustedError") +} + +func TestAttack_ShouldRetry_RetryTrueWithRewriteError_IgnoresRewrite(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count == 1 { + return nil, errors.New("transient") + } + return schema.AssistantMessage("success", nil), nil + }).Times(2) + + rewriteErr := errors.New("this should be ignored") + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryTrueRewriteAgent", + Description: "Retry=true with RewriteError should ignore the rewrite", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{Retry: true, RewriteError: rewriteErr} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + iterator := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("Hello")}}) + events := drainAgentEvents(t, iterator) + + var foundSuccess bool + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, rewriteErr) { + t.Fatal("RewriteError should be ignored when Retry=true") + } + if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.Message != nil { + if e.Output.MessageOutput.Message.Content == "success" { + foundSuccess = true + } + } + } + assert.True(t, foundSuccess, "should eventually succeed after retry, ignoring RewriteError") +} + +func TestAttack_ShouldRetry_OptionsAccumulateAcrossRetries(t *testing.T) { + ctx := context.Background() + + var capturedOpts [][]model.Option + var callCount int32 + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + capturedOpts = append(capturedOpts, opts) + if count <= 2 { + return nil, errors.New("needs retry") + } + return schema.AssistantMessage("done", nil), nil + }).Times(3) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OptsAccumulateAgent", + Description: "Verify options accumulate across retries", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 5, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{ + Retry: true, + AdditionalOptions: []model.Option{model.WithMaxTokens(100 * retryCtx.RetryAttempt)}, + } + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + iterator := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("Hello")}}) + drainAgentEvents(t, iterator) + + require.Len(t, capturedOpts, 3) + assert.True(t, len(capturedOpts[1]) > len(capturedOpts[0]), + "second call should have more options than first (accumulated AdditionalOptions)") + assert.True(t, len(capturedOpts[2]) > len(capturedOpts[1]), + "third call should have more options than second (accumulated AdditionalOptions)") +} + +func TestAttack_ShouldRetry_Stream_NilDecisionAccepts(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("stream ok", nil)}), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamNilDecisionAgent", + Description: "ShouldRetry returns nil in stream mode — should accept", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return nil + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + + events, streamTermErrs := drainStreamingAgentEvents(t, iterator) + var foundStreamContent bool + for _, e := range events { + if e.StreamContent == "stream ok" { + foundStreamContent = true + } + } + assert.True(t, foundStreamContent, "nil decision should accept the stream") + for _, termErr := range streamTermErrs { + assert.Equal(t, io.EOF, termErr, "stream should terminate with clean EOF, not error") + } +} + +func TestAttack_ShouldRetry_Stream_MaxRetriesZero_Exhausted(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("rejected", nil)}), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamMaxZeroAgent", + Description: "Stream mode with MaxRetries=0 rejecting — should exhaust immediately", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 0, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: true} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + + done := make(chan struct{}) + var events []agentEvent + go func() { + defer close(done) + iterator := agent.Run(ctx, input) + events, _ = drainStreamingAgentEvents(t, iterator) + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("test deadlocked — Stream MaxRetries=0 with reject should not hang") + } + + var foundExhausted bool + for _, e := range events { + if e.Err != nil { + var exhaustedErr *RetryExhaustedError + if errors.As(e.Err, &exhaustedErr) { + foundExhausted = true + } + } + } + assert.True(t, foundExhausted, "MaxRetries=0 stream reject should produce RetryExhaustedError") +} + +func TestAttack_ShouldRetry_Stream_RewriteErrorOnCleanStream(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("looks good but bad", nil)}), nil + }).Times(1) + + fatalErr := errors.New("fatal: content policy violation") + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "StreamRewriteCleanAgent", + Description: "Stream returns cleanly but ShouldRetry rewrites to error", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: false, RewriteError: fatalErr} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + + done := make(chan struct{}) + var events []agentEvent + go func() { + defer close(done) + iterator := agent.Run(ctx, input) + events, _ = drainStreamingAgentEvents(t, iterator) + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("test deadlocked") + } + + var foundFatal bool + for _, e := range events { + if e.Err != nil && errors.Is(e.Err, fatalErr) { + foundFatal = true + } + } + assert.True(t, foundFatal, "RewriteError on clean stream should propagate the fatal error") +} + +func TestAttack_ShouldRetry_ConcatMessagesFails_EmptyStream(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + r, w := schema.Pipe[*schema.Message](1) + w.Close() + return r, nil + }).Times(1) + + var capturedCtx *RetryContext + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "EmptyStreamAgent", + Description: "Stream returns zero chunks — both OutputMessage and Err should be nil", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + capturedCtx = retryCtx + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + + done := make(chan struct{}) + go func() { + defer close(done) + iterator := agent.Run(ctx, input) + drainStreamingAgentEvents(t, iterator) + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("test deadlocked on empty stream") + } + + require.NotNil(t, capturedCtx) + assert.NotNil(t, capturedCtx.OutputMessage, "empty stream should have non-nil OutputMessage from ConcatMessages") + assert.Nil(t, capturedCtx.Err, "empty stream should have nil Err") +} + +func TestAttack_ShouldRetry_Stream_MidStreamError_VerdictDoubleRead(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + midStreamErr := errors.New("mid-stream transient error") + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + r, w := schema.Pipe[*schema.Message](1) + go func() { + defer w.Close() + _ = w.Send(schema.AssistantMessage("chunk1", nil), nil) + _ = w.Send(nil, midStreamErr) + }() + return r, nil + }).Times(2) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "DoubleReadBugAgent", + Description: "Reproduces signal.ch double-read when event stream hits mid-stream error then EOF", + Instruction: "You are a helpful assistant.", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: instantBackoff, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + EnableStreaming: true, + } + + done := make(chan struct{}) + go func() { + defer close(done) + iterator := agent.Run(ctx, input) + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil { + mo := event.Output.MessageOutput + if mo.IsStreaming && mo.MessageStream != nil { + for { + _, recvErr := mo.MessageStream.Recv() + if recvErr == io.EOF { + break + } + } + } + } + } + }() + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("goroutine leak: onEOF blocked on signal.ch after errWrapper already drained the verdict") + } +} + +type rejectReasonTestModel struct { + streamFn func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) +} + +func (m *rejectReasonTestModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("generated", nil), nil +} + +func (m *rejectReasonTestModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return m.streamFn(ctx, input, opts...) +} + +func TestRejectReason_StreamPath(t *testing.T) { + ctx := context.Background() + + type rejectInfo struct { + Reason string + Attempt int + } + + streamErr := errors.New("bad output") + var streamCallCount int32 + + m := &rejectReasonTestModel{ + streamFn: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + n := atomic.AddInt32(&streamCallCount, 1) + if n == 1 { + return streamWithMidError( + []*schema.Message{schema.AssistantMessage("rejected chunk", nil)}, + streamErr, + ), nil + } + sr, sw := schema.Pipe[*schema.Message](1) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("accepted", nil), nil) + }() + return sr, nil + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "reject-reason-agent", + Description: "test reject reason", + Instruction: "test", + Model: m, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.Err != nil { + return &RetryDecision{ + Retry: true, + RejectReason: rejectInfo{ + Reason: "output quality too low", + Attempt: retryCtx.RetryAttempt, + }, + } + } + return nil + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return time.Millisecond }, + }, + }) + require.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("hello")}, + EnableStreaming: true, + } + ctx, _ = initRunCtx(ctx, agent.Name(ctx), input) + iter := agent.Run(ctx, input) + + var capturedRejectReasons []any + var finalContent string + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + continue + } + if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.IsStreaming { + sr := ev.Output.MessageOutput.MessageStream + for { + chunk, recvErr := sr.Recv() + if recvErr != nil { + var willRetry *WillRetryError + if errors.As(recvErr, &willRetry) { + capturedRejectReasons = append(capturedRejectReasons, willRetry.RejectReason()) + } + break + } + if chunk != nil { + finalContent = chunk.Content + } + } + } + } + + assert.Contains(t, finalContent, "accepted") + require.NotEmpty(t, capturedRejectReasons, "should have at least one WillRetryError with RejectReason from stream Recv()") + for _, reason := range capturedRejectReasons { + require.NotNil(t, reason) + ri, ok := reason.(rejectInfo) + require.True(t, ok, "RejectReason should be rejectInfo type, got %T", reason) + assert.Equal(t, "output quality too low", ri.Reason) + assert.Equal(t, 1, ri.Attempt) + } +} + +func TestWillRetryError_RejectReason(t *testing.T) { + t.Run("nil when not set", func(t *testing.T) { + wrErr := &WillRetryError{ErrStr: "test", RetryAttempt: 1, err: errors.New("test")} + assert.Nil(t, wrErr.RejectReason(), "RejectReason should be nil when not set") + }) + + t.Run("returns value when set", func(t *testing.T) { + reason := map[string]string{"key": "value"} + wrErr := &WillRetryError{ + ErrStr: "rejected", + RetryAttempt: 2, + rejectReason: reason, + err: errors.New("inner"), + } + assert.Equal(t, reason, wrErr.RejectReason()) + assert.Equal(t, "rejected", wrErr.Error()) + assert.Equal(t, 2, wrErr.RetryAttempt) + }) +} diff --git a/adk/chatmodel_test.go b/adk/chatmodel_test.go new file mode 100644 index 0000000..2c92064 --- /dev/null +++ b/adk/chatmodel_test.go @@ -0,0 +1,2417 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +// TestChatModelAgentRun tests the Run method of ChatModelAgent +func TestChatModelAgentRun(t *testing.T) { + // Basic test for Run method + t.Run("BasicFunctionality", func(t *testing.T) { + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello, I am an AI assistant.", nil), nil). + Times(1) + + // Create a ChatModelAgent + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Hello, who are you?"), + }, + } + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + + // Verify the message content + msg := event.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.Equal(t, "Hello, I am an AI assistant.", msg.Content) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) + + t.Run("BasicChatModelWithAgentMiddleware", func(t *testing.T) { + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello, I am an AI assistant.", nil), nil). + Times(1) + + afterChatModelExecuted := false + + // Create a ChatModelAgent + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + Middlewares: []AgentMiddleware{ + { + BeforeChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + state.Messages = append(state.Messages, schema.UserMessage("m")) + return nil + }, + AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + assert.Len(t, state.Messages, 4) + afterChatModelExecuted = true + return nil + }, + }, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Hello, who are you?"), + }, + } + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + + // Verify the message content + msg := event.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.Equal(t, "Hello, I am an AI assistant.", msg.Content) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + + assert.True(t, afterChatModelExecuted) + }) + + t.Run("AfterChatModel_NoTools_ModifyDoesNotAffectEvent", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("original content", nil), nil). + Times(1) + + var capturedMessages []*schema.Message + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for AfterChatModel NoTools scenario", + Instruction: "You are a helpful assistant.", + Model: cm, + Middlewares: []AgentMiddleware{ + { + AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + capturedMessages = make([]*schema.Message, len(state.Messages)) + copy(capturedMessages, state.Messages) + state.Messages = append(state.Messages, schema.AssistantMessage("appended content", nil)) + return nil + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.Equal(t, "original content", msg.Content) + + _, ok = iterator.Next() + assert.False(t, ok) + + assert.Len(t, capturedMessages, 3) + }) + + t.Run("AfterChatModel_ReAct_ModifyAffectsFlow", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("first response with tool call", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: "test_tool", Arguments: "{}"}}, + }), nil + } + return schema.AssistantMessage("final response", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + toolCalled := false + testTool := &fakeToolForTest{tarCount: 0} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for AfterChatModel ReAct scenario", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + lastMsg := state.Messages[len(state.Messages)-1] + if len(lastMsg.ToolCalls) > 0 { + toolCalled = true + state.Messages[len(state.Messages)-1] = schema.AssistantMessage("modified to remove tool call", nil) + } + return nil + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.True(t, toolCalled) + assert.Equal(t, 1, generateCount) + + assert.Equal(t, 1, len(events)) + event := events[0] + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.Equal(t, "first response with tool call", event.Output.MessageOutput.Message.Content) + assert.Len(t, event.Output.MessageOutput.Message.ToolCalls, 1) + }) + + t.Run("AfterChatModel_ReAct_AppendToolCall_AffectsFlow", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("first response no tool", nil), nil + } + return schema.AssistantMessage("final response", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + testTool := &fakeToolForTest{tarCount: 0} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for AfterChatModel ReAct append tool call", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + if generateCount == 1 { + state.Messages[len(state.Messages)-1] = schema.AssistantMessage("modified with tool call", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: "test_tool", Arguments: "{}"}}, + }) + } + return nil + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 2, generateCount) + + assert.Equal(t, 3, len(events)) + + event0 := events[0] + assert.NotNil(t, event0.Output) + assert.NotNil(t, event0.Output.MessageOutput) + assert.Equal(t, "first response no tool", event0.Output.MessageOutput.Message.Content) + assert.Empty(t, event0.Output.MessageOutput.Message.ToolCalls) + + event2 := events[2] + assert.NotNil(t, event2.Output) + assert.NotNil(t, event2.Output.MessageOutput) + assert.Equal(t, "final response", event2.Output.MessageOutput.Message.Content) + }) + + // Test with streaming output + t.Run("StreamOutput", func(t *testing.T) { + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create a stream reader for the mock response + sr := schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Hello", nil), + schema.AssistantMessage(", I am", nil), + schema.AssistantMessage(" an AI assistant.", nil), + }) + + // Set up expectations for the mock model + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(sr, nil). + Times(1) + + // Create a ChatModelAgent + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent with streaming enabled + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello, who are you?")}, + EnableStreaming: true, + } + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) + + // Test error handling + t.Run("ErrorHandling", func(t *testing.T) { + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model to return an error + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, errors.New("model error")). + Times(1) + + // Create a ChatModelAgent + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello, who are you?")}, + } + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator, should contain an error + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "model error") + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) + + // Test with tools + t.Run("WithTools", func(t *testing.T) { + ctx := context.Background() + + // Create a fake tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 1, + } + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }}), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Task completed", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create a ChatModelAgent with tools + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Use the test tool")}, + } + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get events from the iterator + // First event should be the model output with tool call + event1, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event1) + assert.Nil(t, event1.Err) + assert.NotNil(t, event1.Output) + assert.NotNil(t, event1.Output.MessageOutput) + assert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role) + + // Second event should be the tool output + event2, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event2) + assert.Nil(t, event2.Err) + assert.NotNil(t, event2.Output) + assert.NotNil(t, event2.Output.MessageOutput) + assert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role) + + // Third event should be the final model output + event3, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event3) + assert.Nil(t, event3.Err) + assert.NotNil(t, event3.Output) + assert.NotNil(t, event3.Output.MessageOutput) + assert.Equal(t, schema.Assistant, event3.Output.MessageOutput.Role) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) + + t.Run("WrapToolCall_ToolMiddleware", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + fakeTool := &fakeToolForTest{tarCount: 1} + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"name":"test"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + var capturedToolName string + var capturedArgs string + middlewareInvoked := false + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test WrapToolCall middleware", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + WrapToolCall: compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + middlewareInvoked = true + capturedToolName = input.Name + capturedArgs = input.Arguments + return next(ctx, input) + } + }, + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.True(t, middlewareInvoked, "WrapToolCall middleware should have been invoked") + assert.Equal(t, info.Name, capturedToolName) + assert.Equal(t, `{"name":"test"}`, capturedArgs) + assert.Equal(t, 3, len(events)) + assert.Equal(t, schema.Assistant, events[0].Output.MessageOutput.Role) + assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role) + assert.Equal(t, schema.Assistant, events[2].Output.MessageOutput.Role) + assert.Equal(t, "done", events[2].Output.MessageOutput.Message.Content) + }) + + t.Run("WrapToolCall_MultipleMiddlewares_Order", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + fakeTool := &fakeToolForTest{tarCount: 1} + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + var order []string + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test WrapToolCall middleware order", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + WrapToolCall: compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + order = append(order, "mw1_before") + out, err := next(ctx, input) + order = append(order, "mw1_after") + return out, err + } + }, + }, + }, + { + WrapToolCall: compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + order = append(order, "mw2_before") + out, err := next(ctx, input) + order = append(order, "mw2_after") + return out, err + } + }, + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + } + iterator := agent.Run(ctx, input) + + for { + _, ok := iterator.Next() + if !ok { + break + } + } + + // First registered middleware is outermost + assert.Equal(t, []string{"mw1_before", "mw2_before", "mw2_after", "mw1_after"}, order) + }) + + t.Run("WrapToolCall_ModifyResult", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + fakeTool := &fakeToolForTest{tarCount: 1} + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"name":"test"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test WrapToolCall modify result", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + WrapToolCall: compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + out, err := next(ctx, input) + if err != nil { + return out, err + } + out.Result = "modified: " + out.Result + return out, nil + } + }, + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 3, len(events)) + // Tool result event should have the modified content + toolEvent := events[1] + assert.Equal(t, schema.Tool, toolEvent.Output.MessageOutput.Role) + assert.Contains(t, toolEvent.Output.MessageOutput.Message.Content, "modified:") + }) + + t.Run("WrapToolCall_StreamableMiddleware", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + streamTool := &fakeStreamToolForTest{tarCount: 1} + info, err := streamTool.Info(ctx) + assert.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{"name":"test"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + middlewareInvoked := false + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test Streamable WrapToolCall middleware", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{streamTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + WrapToolCall: compose.ToolMiddleware{ + Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + middlewareInvoked = true + return next(ctx, input) + } + }, + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.True(t, middlewareInvoked, "Streamable WrapToolCall middleware should have been invoked") + assert.Equal(t, 3, len(events)) + assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role) + }) + + t.Run("WrapToolCall_EnhancedInvokableMiddleware", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + enhTool := &fakeEnhancedInvokableToolForTest{} + info, err := enhTool.Info(ctx) + assert.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + middlewareInvoked := false + var capturedToolName string + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test EnhancedInvokable WrapToolCall middleware", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{enhTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + WrapToolCall: compose.ToolMiddleware{ + EnhancedInvokable: func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + middlewareInvoked = true + capturedToolName = input.Name + return next(ctx, input) + } + }, + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.True(t, middlewareInvoked, "EnhancedInvokable WrapToolCall middleware should have been invoked") + assert.Equal(t, info.Name, capturedToolName) + assert.Equal(t, 3, len(events)) + assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role) + }) + + t.Run("WrapToolCall_EnhancedStreamableMiddleware", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + enhStreamTool := &fakeEnhancedStreamableToolForTest{} + info, err := enhStreamTool.Info(ctx) + assert.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: info.Name, Arguments: `{}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + middlewareInvoked := false + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test EnhancedStreamable WrapToolCall middleware", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{enhStreamTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + WrapToolCall: compose.ToolMiddleware{ + EnhancedStreamable: func(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + middlewareInvoked = true + return next(ctx, input) + } + }, + }, + }, + }, + }) + assert.NoError(t, err) + + input := &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + } + iterator := agent.Run(ctx, input) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.True(t, middlewareInvoked, "EnhancedStreamable WrapToolCall middleware should have been invoked") + assert.Equal(t, 3, len(events)) + assert.Equal(t, schema.Tool, events[1].Output.MessageOutput.Role) + }) +} + +// TestExitTool tests the Exit tool functionality +func TestExitTool(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock chat model + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + // First call: model generates a message with Exit tool call + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("I'll exit with a final result", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: "exit", + Arguments: `{"final_result": "This is the final result"}`}, + }, + }), nil). + Times(1) + + // Model should implement WithTools + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create an agent with the Exit tool + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with Exit tool", + Instruction: "You are a helpful assistant.", + Model: cm, + Exit: &ExitTool{}, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Please exit with a final result"), + }, + } + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // First event: model output with tool call + event1, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event1) + assert.Nil(t, event1.Err) + assert.NotNil(t, event1.Output) + assert.NotNil(t, event1.Output.MessageOutput) + assert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role) + + // Second event: tool output (Exit) + event2, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event2) + assert.Nil(t, event2.Err) + assert.NotNil(t, event2.Output) + assert.NotNil(t, event2.Output.MessageOutput) + assert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role) + + // Verify the action is Exit + assert.NotNil(t, event2.Action) + assert.True(t, event2.Action.Exit) + + // Verify the final result + assert.Equal(t, "This is the final result", event2.Output.MessageOutput.Message.Content) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestParallelReturnDirectlyToolCall(t *testing.T) { + ctx := context.Background() + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock chat model + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + // First call: model generates a message with Exit tool call + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("I'll exit with a final result", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{Name: "tool1"}, + }, + { + ID: "tool-call-2", + Function: schema.FunctionCall{Name: "tool2"}, + }, + { + ID: "tool-call-3", + Function: schema.FunctionCall{Name: "tool3"}, + }, + }), nil). + Times(1) + + // Model should implement WithTools + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create an agent with the Exit tool + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with Exit tool", + Instruction: "You are a helpful assistant.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{ + &myTool{name: "tool1", desc: "tool1", waitTime: time.Millisecond}, + &myTool{name: "tool2", desc: "tool2", waitTime: 10 * time.Millisecond}, + &myTool{name: "tool3", desc: "tool3", waitTime: 100 * time.Millisecond}, + }, + }, + ReturnDirectly: map[string]bool{ + "tool1": true, + }, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + }) + iter := r.Query(ctx, "") + times := 0 + for { + e, ok := iter.Next() + if !ok { + assert.Equal(t, 4, times) + break + } + if times == 3 { + assert.Equal(t, "tool1", e.Output.MessageOutput.Message.ToolName) + } + times++ + } +} + +func TestConcurrentSameToolSendToolGenActionUsesToolCallID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("tools", []schema.ToolCall{ + {ID: "id1", Function: schema.FunctionCall{Name: "action_tool", Arguments: "A"}}, + {ID: "id2", Function: schema.FunctionCall{Name: "action_tool", Arguments: "B"}}, + }), nil). + Times(1) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Agent with action tool", + Instruction: "", + Model: cm, + ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{actionTool{}}}}, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}}) + seen := map[string]bool{} + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.Message != nil && e.Output.MessageOutput.Message.Role == schema.Tool { + if e.Action != nil && e.Action.CustomizedAction != nil { + if s, ok := e.Action.CustomizedAction.(string); ok { + seen[s] = true + } + } + } + } + assert.True(t, seen["A"]) + assert.True(t, seen["B"]) +} + +type myTool struct { + name string + desc string + waitTime time.Duration +} + +func (m *myTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: m.name, + Desc: m.desc, + }, nil +} + +func (m *myTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + time.Sleep(m.waitTime) + return "success", nil +} + +type actionTool struct{} + +func (a actionTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "action_tool", Desc: "action tool"}, nil +} + +func (a actionTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + _ = SendToolGenAction(ctx, "action_tool", &AgentAction{CustomizedAction: argumentsInJSON}) + return "ok", nil +} + +type streamActionTool struct{} + +func (s streamActionTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "action_tool_stream", Desc: "action stream tool"}, nil +} + +func (s streamActionTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) { + _ = SendToolGenAction(ctx, "action_tool_stream", &AgentAction{CustomizedAction: argumentsInJSON}) + sr, sw := schema.Pipe[string](1) + go func() { + defer sw.Close() + _ = sw.Send("o", nil) + _ = sw.Send("k", nil) + }() + return sr, nil +} + +type legacyStreamActionTool struct{} + +func (s legacyStreamActionTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "legacy_action_tool_stream", Desc: "legacy action stream tool"}, nil +} + +func (s legacyStreamActionTool) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) { + _ = compose.ProcessState(ctx, func(ctx context.Context, st *State) error { + st.setToolGenAction("legacy_action_tool_stream", &AgentAction{CustomizedAction: argumentsInJSON}) + return nil + }) + sr, sw := schema.Pipe[string](1) + go func() { + defer sw.Close() + _ = sw.Send("o", nil) + _ = sw.Send("k", nil) + }() + return sr, nil +} + +// TestChatModelAgentOutputKey tests the outputKey configuration and setOutputToSession function +func TestChatModelAgentOutputKey(t *testing.T) { + // Test outputKey configuration - stores output in session + t.Run("OutputKeyStoresInSession", func(t *testing.T) { + for i := 0; i < 1000; i++ { + + } + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello, I am an AI assistant.", nil), nil). + Times(1) + + // Create a ChatModelAgent with outputKey configured + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + OutputKey: "agent_output", // This should store output in session + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Initialize a run context to enable session storage + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Hello, who are you?"), + }, + } + ctx, runCtx := initRunCtx(ctx, "TestAgent", input) + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Run the agent + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + + // Verify the message content + msg := event.Output.MessageOutput.Message + assert.Equal(t, "Hello, I am an AI assistant.", msg.Content) + + // Verify that the output was stored in the session + time.AfterFunc(100*time.Millisecond, func() { + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "agent_output") + assert.Equal(t, "Hello, I am an AI assistant.", sessionValues["agent_output"]) + }) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) + + // Test outputKey configuration with streaming - stores concatenated output in session + t.Run("OutputKeyWithStreamingStoresInSession", func(t *testing.T) { + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create a stream reader for the mock response + sr := schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Hello", nil), + schema.AssistantMessage(", I am", nil), + schema.AssistantMessage(" an AI assistant.", nil), + }) + + // Set up expectations for the mock model + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(sr, nil). + Times(1) + + // Create a ChatModelAgent with outputKey configured + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for unit testing", + Instruction: "You are a helpful assistant.", + Model: cm, + OutputKey: "agent_output", // This should store concatenated stream in session + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Initialize a run context to enable session storage + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Hello, who are you?")}, + EnableStreaming: true, + } + ctx, runCtx := initRunCtx(ctx, "TestAgent", input) + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Run the agent + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + time.AfterFunc(100*time.Millisecond, func() { + // Verify that the concatenated output was stored in the session + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "agent_output") + assert.Equal(t, "Hello, I am an AI assistant.", sessionValues["agent_output"]) + }) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) + + // Test setOutputToSession function directly - regular message + t.Run("SetOutputToSessionRegularMessage", func(t *testing.T) { + ctx := context.Background() + + // Initialize a run context to enable session storage + input := &AgentInput{ + Messages: []Message{schema.UserMessage("test")}, + } + ctx, runCtx := initRunCtx(ctx, "TestAgent", input) + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Test with a regular message + msg := schema.AssistantMessage("Test response", nil) + err := setOutputToSession(ctx, msg, nil, "test_output") + assert.NoError(t, err) + + // Verify the message content was stored + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "test_output") + assert.Equal(t, "Test response", sessionValues["test_output"]) + }) + + // Test setOutputToSession function directly - streaming message + t.Run("SetOutputToSessionStreamingMessage", func(t *testing.T) { + ctx := context.Background() + + // Initialize a run context to enable session storage + input := &AgentInput{ + Messages: []Message{schema.UserMessage("test")}, + } + ctx, runCtx := initRunCtx(ctx, "TestAgent", input) + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Test with a streaming message + sr := schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Stream", nil), + schema.AssistantMessage(" response", nil), + schema.AssistantMessage(" content", nil), + }) + err := setOutputToSession(ctx, nil, sr, "test_output") + assert.NoError(t, err) + + // Verify the concatenated stream content was stored + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "test_output") + assert.Equal(t, "Stream response content", sessionValues["test_output"]) + }) + + // Test setOutputToSession function directly - error case + t.Run("SetOutputToSessionErrorCase", func(t *testing.T) { + ctx := context.Background() + + // Initialize a run context to enable session storage + input := &AgentInput{ + Messages: []Message{schema.UserMessage("test")}, + } + ctx, runCtx := initRunCtx(ctx, "TestAgent", input) + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Test with an invalid stream (simulate error) + // Create a stream that will fail when concatenated + sr := schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("test", nil), + }) + // Close the stream to simulate an error condition + sr.Close() + + // This should return an error because the stream is closed + err := setOutputToSession(ctx, nil, sr, "test_output") + // Note: The actual behavior may vary depending on the stream implementation + // Some streams may not error when closed, so we'll accept either outcome + if err != nil { + // If there's an error, verify nothing was stored + sessionValues := GetSessionValues(ctx) + assert.NotContains(t, sessionValues, "test_output") + } else { + // If no error, verify the content was stored + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "test_output") + assert.Equal(t, "test", sessionValues["test_output"]) + } + }) + + // Test outputKey with React workflow (tools enabled) + t.Run("OutputKeyWithReactWorkflow", func(t *testing.T) { + ctx := context.Background() + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create a simple tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 1, + } + + // Set up expectations for the mock model - it will generate a final response + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Final response from React workflow", nil), nil). + Times(1) + // Model should implement WithTools + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create a ChatModelAgent with outputKey and tools configured + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent with tools", + Instruction: "You are a helpful assistant.", + Model: cm, + OutputKey: "agent_output", + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + }) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Initialize a run context to enable session storage + input := &AgentInput{ + Messages: []Message{schema.UserMessage("Use the tool")}, + } + ctx, runCtx := initRunCtx(ctx, "TestAgent", input) + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Run the agent + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.Nil(t, event.Err) + + // Verify the message content + msg := event.Output.MessageOutput.Message + assert.Equal(t, "Final response from React workflow", msg.Content) + + // Verify that the output was stored in the session + time.AfterFunc(time.Millisecond*10, func() { + sessionValues := GetSessionValues(ctx) + assert.Contains(t, sessionValues, "agent_output") + assert.Equal(t, "Final response from React workflow", sessionValues["agent_output"]) + }) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) + }) +} + +func TestConcurrentSameStreamToolSendToolGenActionUsesToolCallID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("tools", []schema.ToolCall{ + {ID: "sid1", Function: schema.FunctionCall{Name: "action_tool_stream", Arguments: "SA"}}, + {ID: "sid2", Function: schema.FunctionCall{Name: "action_tool_stream", Arguments: "SB"}}, + }), nil). + Times(1) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Agent with stream action tool", + Instruction: "", + Model: cm, + ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{streamActionTool{}}}}, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}}) + seen := map[string]bool{} + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Output != nil && e.Output.MessageOutput != nil { + if e.Output.MessageOutput.IsStreaming { + if e.Action != nil && e.Action.CustomizedAction != nil { + if s, ok := e.Action.CustomizedAction.(string); ok { + seen[s] = true + } + } + } + } + } + assert.True(t, seen["SA"]) + assert.True(t, seen["SB"]) +} + +func TestStreamToolLegacyNameKeyFallback(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("tools", []schema.ToolCall{ + {ID: "lsid1", Function: schema.FunctionCall{Name: "legacy_action_tool_stream", Arguments: "LA"}}, + }), nil). + Times(1) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Agent with legacy stream action tool", + Instruction: "", + Model: cm, + ToolsConfig: ToolsConfig{ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{legacyStreamActionTool{}}}}, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("go")}}) + found := false + for { + e, ok := iter.Next() + if !ok { + break + } + if e.Output != nil && e.Output.MessageOutput != nil && e.Output.MessageOutput.IsStreaming { + if e.Action != nil && e.Action.CustomizedAction != nil { + if s, ok := e.Action.CustomizedAction.(string); ok { + found = (s == "LA") + } + } + } + } + assert.True(t, found) +} + +func TestChatModelAgent_ToolResultMiddleware_EmitsFinalResult(t *testing.T) { + originalResult := "original_result" + modifiedResult := "modified_by_middleware" + + resultModifyingMiddleware := compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + output, err := next(ctx, input) + if err != nil { + return nil, err + } + output.Result = modifiedResult + return output, nil + } + }, + Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + output, err := next(ctx, input) + if err != nil { + return nil, err + } + output.Result = schema.StreamReaderFromArray([]string{modifiedResult}) + return output, nil + } + }, + } + + t.Run("Invoke", func(t *testing.T) { + ctx := context.Background() + testTool := &simpleToolForMiddlewareTest{name: "test_tool", result: originalResult} + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + info, err := testTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"input": "test"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "test_agent", + Description: "test agent with middleware", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + ToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware}, + }, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false, CheckPointStore: newBridgeStore()}) + it := r.Run(ctx, []Message{schema.UserMessage("call the tool")}) + + var toolResultEvents []*AgentEvent + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil && + ev.Output.MessageOutput.Message != nil && + ev.Output.MessageOutput.Message.Role == schema.Tool { + toolResultEvents = append(toolResultEvents, ev) + } + } + + assert.NotEmpty(t, toolResultEvents, "should have at least one tool result event") + for _, ev := range toolResultEvents { + assert.Equal(t, modifiedResult, ev.Output.MessageOutput.Message.Content, + "tool result event should contain the middleware-modified result, not the original") + assert.NotEqual(t, originalResult, ev.Output.MessageOutput.Message.Content, + "tool result event should NOT contain the original result") + } + }) + + t.Run("Stream", func(t *testing.T) { + ctx := context.Background() + testTool := &simpleToolForMiddlewareTest{name: "test_tool_stream", result: originalResult} + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + info, err := testTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"input": "test"}`, + }, + }, + }), + }), nil). + Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("final response", nil), + }), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "test_agent", + Description: "test agent with middleware", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + ToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware}, + }, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()}) + it := r.Run(ctx, []Message{schema.UserMessage("call the tool")}) + + var toolResultContents []string + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + if ev.Output.MessageOutput.Message != nil && + ev.Output.MessageOutput.Message.Role == schema.Tool { + toolResultContents = append(toolResultContents, ev.Output.MessageOutput.Message.Content) + } + if ev.Output.MessageOutput.IsStreaming && + ev.Output.MessageOutput.MessageStream != nil && + ev.Output.MessageOutput.Role == schema.Tool { + var msgs []*schema.Message + for { + msg, err := ev.Output.MessageOutput.MessageStream.Recv() + if err != nil { + break + } + msgs = append(msgs, msg) + } + if len(msgs) > 0 { + concated, err := schema.ConcatMessages(msgs) + if err == nil { + toolResultContents = append(toolResultContents, concated.Content) + } + } + } + } + } + + assert.NotEmpty(t, toolResultContents, "should have at least one tool result event") + for _, content := range toolResultContents { + assert.Equal(t, modifiedResult, content, + "tool result event should contain the middleware-modified result, not the original") + assert.NotEqual(t, originalResult, content, + "tool result event should NOT contain the original result") + } + }) +} + +type simpleToolForMiddlewareTest struct { + name string + result string +} + +func (s *simpleToolForMiddlewareTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: s.name, + Desc: "simple tool", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "input": { + Desc: "input", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (s *simpleToolForMiddlewareTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return s.result, nil +} + +func (s *simpleToolForMiddlewareTest) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) { + return schema.StreamReaderFromArray([]string{s.result}), nil +} + +func TestGetComposeOptions(t *testing.T) { + t.Run("WithChatModelOptions", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var capturedTemperature float32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + options := model.GetCommonOptions(&model.Options{}, opts...) + if options.Temperature != nil { + capturedTemperature = *options.Temperature + } + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + }) + assert.NoError(t, err) + + temp := float32(0.7) + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + WithChatModelOptions([]model.Option{model.WithTemperature(temp)})) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, temp, capturedTemperature, "Temperature should be passed through WithChatModelOptions") + }) + + t.Run("WithToolOptions", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var toolOptionsCaptured bool + testTool := &toolOptionCapturingTool{ + name: "test_tool", + onRun: func(opts []tool.Option) { + if len(opts) > 0 { + toolOptionsCaptured = true + } + }, + } + info, _ := testTool.Info(ctx) + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil).Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + WithToolOptions([]tool.Option{testToolOption("test_value")})) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, toolOptionsCaptured, "Tool options should be passed through WithToolOptions") + }) + +} + +type toolOptionCapturingTool struct { + name string + onRun func(opts []tool.Option) +} + +func (t *toolOptionCapturingTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: t.name + " description"}, nil +} + +func (t *toolOptionCapturingTool) InvokableRun(_ context.Context, _ string, opts ...tool.Option) (string, error) { + if t.onRun != nil { + t.onRun(opts) + } + return t.name + " result", nil +} + +type testToolOptions struct { + value string +} + +func testToolOption(value string) tool.Option { + return tool.WrapImplSpecificOptFn(func(o *testToolOptions) { + o.value = value + }) +} + +type errorTool struct { + infoErr error +} + +func (e *errorTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return nil, e.infoErr +} + +func (e *errorTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return "", nil +} + +func TestChatModelAgent_PrepareExecContextError(t *testing.T) { + t.Run("Run_WithToolInfoError_ReturnsError", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + expectedErr := errors.New("tool info error") + errTool := &errorTool{infoErr: expectedErr} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{errTool}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + event, ok := iter.Next() + assert.True(t, ok) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "tool info error") + + _, ok = iter.Next() + assert.False(t, ok) + }) + + t.Run("Resume_WithToolInfoError_ReturnsError", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + expectedErr := errors.New("tool info error for resume") + errTool := &errorTool{infoErr: expectedErr} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{errTool}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Resume(ctx, &ResumeInfo{ + InterruptState: []byte("dummy"), + EnableStreaming: false, + }) + + event, ok := iter.Next() + assert.True(t, ok) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "tool info error for resume") + + _, ok = iter.Next() + assert.False(t, ok) + }) +} + +// fakeEnhancedInvokableToolForTest implements tool.EnhancedInvokableTool for testing. +type fakeEnhancedInvokableToolForTest struct{} + +func (t *fakeEnhancedInvokableToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "test_enhanced_invokable_tool", + Desc: "test enhanced invokable tool", + }, nil +} + +func (t *fakeEnhancedInvokableToolForTest) InvokableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "enhanced invokable result"}, + }, + }, nil +} + +// fakeEnhancedStreamableToolForTest implements tool.EnhancedStreamableTool for testing. +type fakeEnhancedStreamableToolForTest struct{} + +func (t *fakeEnhancedStreamableToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "test_enhanced_streamable_tool", + Desc: "test enhanced streamable tool", + }, nil +} + +func (t *fakeEnhancedStreamableToolForTest) StreamableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr := schema.StreamReaderFromArray([]*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "enhanced streamable result"}}}, + }) + return sr, nil +} + +func TestPreprocessComposeCheckpoint_MigrateErrorIsReturned(t *testing.T) { + in := []byte("prefix\u0015" + stateGobNameV080 + "suffix") + _, err := preprocessComposeCheckpoint(in) + assert.Error(t, err) +} + +func TestNewChatModelAgent_FailoverConfigValidation(t *testing.T) { + ctx := context.Background() + cm := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil + }, + } + + t.Run("missing GetFailoverModel", func(t *testing.T) { + _, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: cm, + ModelFailoverConfig: &ModelFailoverConfig[*schema.Message]{ + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + }, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "ModelFailoverConfig.GetFailoverModel") + }) + + t.Run("missing ShouldFailover", func(t *testing.T) { + _, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "test", + Model: cm, + ModelFailoverConfig: &ModelFailoverConfig[*schema.Message]{ + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return cm, nil, nil + }, + }, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "ModelFailoverConfig.ShouldFailover") + }) +} + +// aliasCaptureTool captures the raw arguments JSON received by the tool. +type aliasCaptureTool struct { + name string + params map[string]*schema.ParameterInfo + receivedArgs string +} + +func (t *aliasCaptureTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: t.name + " tool", + ParamsOneOf: schema.NewParamsOneOfByParams(t.params), + }, nil +} + +func (t *aliasCaptureTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + t.receivedArgs = argumentsInJSON + return "ok", nil +} + +func TestToolAliasesPropagation(t *testing.T) { + t.Run("prepareExecContext_propagates_ToolAliases", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + + captureTool := &aliasCaptureTool{ + name: "grep", + params: map[string]*schema.ParameterInfo{ + "pattern": {Type: schema.String, Desc: "regex pattern"}, + "path": {Type: schema.String, Desc: "search path"}, + }, + } + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "grep", + Arguments: `{"grep_content": "TODO", "path": "/src"}`, + }, + }, + }, + }, nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "test", + Instruction: "test", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + ToolAliases: map[string]compose.ToolAliasConfig{ + "grep": { + ArgumentsAliases: map[string][]string{ + "pattern": {"grep_content"}, + }, + }, + }, + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("search for TODOs")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + require.NotEmpty(t, captureTool.receivedArgs, "tool should have been called") + var args map[string]any + err = json.Unmarshal([]byte(captureTool.receivedArgs), &args) + require.NoError(t, err) + assert.Equal(t, "TODO", args["pattern"], "alias 'grep_content' should be remapped to 'pattern'") + assert.NotContains(t, args, "grep_content", "alias key should not be present after remapping") + assert.Equal(t, "/src", args["path"]) + }) + + t.Run("applyBeforeAgent_preserves_ToolAliases_when_handler_modifies_tools", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + + captureTool := &aliasCaptureTool{ + name: "grep", + params: map[string]*schema.ParameterInfo{ + "pattern": {Type: schema.String, Desc: "regex pattern"}, + }, + } + + extraTool := &aliasCaptureTool{ + name: "extra_tool", + params: map[string]*schema.ParameterInfo{ + "input": {Type: schema.String}, + }, + } + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "grep", + Arguments: `{"grep_content": "FIXME"}`, + }, + }, + }, + }, nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + handler := &testToolsHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + tools: []tool.BaseTool{extraTool}, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "test", + Instruction: "test", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + ToolAliases: map[string]compose.ToolAliasConfig{ + "grep": { + ArgumentsAliases: map[string][]string{ + "pattern": {"grep_content"}, + }, + }, + }, + }, + }, + Handlers: []ChatModelAgentMiddleware{handler}, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("search for FIXMEs")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + require.NotEmpty(t, captureTool.receivedArgs, "tool should have been called") + var args map[string]any + err = json.Unmarshal([]byte(captureTool.receivedArgs), &args) + require.NoError(t, err) + assert.Equal(t, "FIXME", args["pattern"], "alias 'grep_content' should be remapped to 'pattern' even after handler rebuild") + assert.NotContains(t, args, "grep_content", "alias key should not be present after remapping") + }) + + t.Run("name_alias_propagated_through_prepareExecContext", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + + captureTool := &aliasCaptureTool{ + name: "grep", + params: map[string]*schema.ParameterInfo{ + "pattern": {Type: schema.String}, + }, + } + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "search_content", + Arguments: `{"pattern": "TODO"}`, + }, + }, + }, + }, nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "test", + Instruction: "test", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{captureTool}, + ToolAliases: map[string]compose.ToolAliasConfig{ + "grep": { + NameAliases: []string{"search_content"}, + }, + }, + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("search")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + require.NotEmpty(t, captureTool.receivedArgs, "tool should have been called via name alias 'search_content'") + var args map[string]any + err = json.Unmarshal([]byte(captureTool.receivedArgs), &args) + require.NoError(t, err) + assert.Equal(t, "TODO", args["pattern"]) + }) + + t.Run("handler_adds_tool_matching_preexisting_ToolAliases_with_no_initial_tools", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + + captureTool := &aliasCaptureTool{ + name: "grep", + params: map[string]*schema.ParameterInfo{ + "pattern": {Type: schema.String, Desc: "regex pattern"}, + }, + } + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "grep", + Arguments: `{"grep_content": "BUG"}`, + }, + }, + }, + }, nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + handler := &testToolsHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + tools: []tool.BaseTool{captureTool}, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "test", + Instruction: "test", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + ToolAliases: map[string]compose.ToolAliasConfig{ + "grep": { + ArgumentsAliases: map[string][]string{ + "pattern": {"grep_content"}, + }, + }, + }, + }, + }, + Handlers: []ChatModelAgentMiddleware{handler}, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("find bugs")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + require.NotEmpty(t, captureTool.receivedArgs, "tool added by handler should have been called") + var args map[string]any + err = json.Unmarshal([]byte(captureTool.receivedArgs), &args) + require.NoError(t, err) + assert.Equal(t, "BUG", args["pattern"], "alias 'grep_content' should be remapped to 'pattern' for handler-added tool") + assert.NotContains(t, args, "grep_content") + }) +} diff --git a/adk/config.go b/adk/config.go new file mode 100644 index 0000000..67ead86 --- /dev/null +++ b/adk/config.go @@ -0,0 +1,35 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import "github.com/cloudwego/eino/adk/internal" + +// Language represents the language setting for the ADK built-in prompts. +type Language = internal.Language + +const ( + // LanguageEnglish represents English language. + LanguageEnglish Language = internal.LanguageEnglish + // LanguageChinese represents Chinese language. + LanguageChinese Language = internal.LanguageChinese +) + +// SetLanguage sets the language for the ADK built-in prompts. +// The default language is English if not explicitly set. +func SetLanguage(lang Language) error { + return internal.SetLanguage(lang) +} diff --git a/adk/deterministic_transfer.go b/adk/deterministic_transfer.go new file mode 100644 index 0000000..ce5b200 --- /dev/null +++ b/adk/deterministic_transfer.go @@ -0,0 +1,301 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "runtime/debug" + "sync" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*deterministicTransferState]("_eino_adk_deterministic_transfer_state") +} + +type deterministicTransferState struct { + EventList []*agentEventWrapper +} + +// AgentWithDeterministicTransferTo wraps an agent to transfer to given agents deterministically. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func AgentWithDeterministicTransferTo(_ context.Context, config *DeterministicTransferConfig) Agent { + if ra, ok := config.Agent.(ResumableAgent); ok { + return &resumableAgentWithDeterministicTransferTo{ + agent: ra, + toAgentNames: config.ToAgentNames, + } + } + return &agentWithDeterministicTransferTo{ + agent: config.Agent, + toAgentNames: config.ToAgentNames, + } +} + +type agentWithDeterministicTransferTo struct { + agent Agent + toAgentNames []string +} + +func (a *agentWithDeterministicTransferTo) Description(ctx context.Context) string { + return a.agent.Description(ctx) +} + +func (a *agentWithDeterministicTransferTo) Name(ctx context.Context) string { + return a.agent.Name(ctx) +} + +func (a *agentWithDeterministicTransferTo) GetType() string { + if typer, ok := a.agent.(components.Typer); ok { + return typer.GetType() + } + return "DeterministicTransfer" +} + +func (a *agentWithDeterministicTransferTo) Run(ctx context.Context, + input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + + if fa, ok := a.agent.(*flowAgent); ok { + return runFlowAgentWithIsolatedSession(ctx, fa, input, a.toAgentNames, options...) + } + + aIter := a.agent.Run(ctx, input, options...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + go forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames) + + return iterator +} + +type resumableAgentWithDeterministicTransferTo struct { + agent ResumableAgent + toAgentNames []string +} + +func (a *resumableAgentWithDeterministicTransferTo) Description(ctx context.Context) string { + return a.agent.Description(ctx) +} + +func (a *resumableAgentWithDeterministicTransferTo) Name(ctx context.Context) string { + return a.agent.Name(ctx) +} + +func (a *resumableAgentWithDeterministicTransferTo) GetType() string { + if typer, ok := a.agent.(components.Typer); ok { + return typer.GetType() + } + return "DeterministicTransfer" +} + +func (a *resumableAgentWithDeterministicTransferTo) Run(ctx context.Context, + input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + + if fa, ok := a.agent.(*flowAgent); ok { + return runFlowAgentWithIsolatedSession(ctx, fa, input, a.toAgentNames, options...) + } + + aIter := a.agent.Run(ctx, input, options...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + go forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames) + + return iterator +} + +func (a *resumableAgentWithDeterministicTransferTo) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + if fa, ok := a.agent.(*flowAgent); ok { + return resumeFlowAgentWithIsolatedSession(ctx, fa, info, a.toAgentNames, opts...) + } + + aIter := a.agent.Resume(ctx, info, opts...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + go forwardEventsAndAppendTransfer(aIter, generator, a.toAgentNames) + + return iterator +} + +func forwardEventsAndAppendTransfer(iter *AsyncIterator[*AgentEvent], + generator *AsyncGenerator[*AgentEvent], toAgentNames []string) { + + defer func() { + if panicErr := recover(); panicErr != nil { + generator.Send(&AgentEvent{Err: safe.NewPanicErr(panicErr, debug.Stack())}) + } + generator.Close() + }() + + var lastEvent *AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + generator.Send(event) + lastEvent = event + } + + if lastEvent != nil && lastEvent.Action != nil && (lastEvent.Action.Interrupted != nil || lastEvent.Action.Exit) { + return + } + + sendTransferEvents(generator, toAgentNames) +} + +func runFlowAgentWithIsolatedSession(ctx context.Context, fa *flowAgent, input *AgentInput, + toAgentNames []string, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + + parentSession := getSession(ctx) + parentRunCtx := getRunCtx(ctx) + + isolatedSession := &runSession{ + Values: parentSession.Values, + valuesMtx: parentSession.valuesMtx, + } + if isolatedSession.valuesMtx == nil { + isolatedSession.valuesMtx = &sync.Mutex{} + } + if isolatedSession.Values == nil { + isolatedSession.Values = make(map[string]any) + } + + ctx = setRunCtx(ctx, &runContext{ + Session: isolatedSession, + RootInput: parentRunCtx.RootInput, + RunPath: parentRunCtx.RunPath, + }) + + iter := fa.Run(ctx, input, options...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + go handleFlowAgentEvents(ctx, iter, generator, isolatedSession, parentSession, toAgentNames) + + return iterator +} + +func resumeFlowAgentWithIsolatedSession(ctx context.Context, fa *flowAgent, info *ResumeInfo, + toAgentNames []string, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + + state, ok := info.InterruptState.(*deterministicTransferState) + if !ok || state == nil { + return genErrorIter(errors.New("invalid interrupt state for flowAgent resume in deterministic transfer")) + } + + parentSession := getSession(ctx) + parentRunCtx := getRunCtx(ctx) + + isolatedSession := &runSession{ + Values: parentSession.Values, + valuesMtx: parentSession.valuesMtx, + Events: state.EventList, + } + if isolatedSession.valuesMtx == nil { + isolatedSession.valuesMtx = &sync.Mutex{} + } + if isolatedSession.Values == nil { + isolatedSession.Values = make(map[string]any) + } + + ctx = setRunCtx(ctx, &runContext{ + Session: isolatedSession, + RootInput: parentRunCtx.RootInput, + RunPath: parentRunCtx.RunPath, + }) + + iter := fa.Resume(ctx, info, opts...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + go handleFlowAgentEvents(ctx, iter, generator, isolatedSession, parentSession, toAgentNames) + + return iterator +} + +func handleFlowAgentEvents(ctx context.Context, iter *AsyncIterator[*AgentEvent], + generator *AsyncGenerator[*AgentEvent], isolatedSession, parentSession *runSession, toAgentNames []string) { + + defer func() { + if panicErr := recover(); panicErr != nil { + generator.Send(&AgentEvent{Err: safe.NewPanicErr(panicErr, debug.Stack())}) + } + generator.Close() + }() + + var lastEvent *AgentEvent + + for { + event, ok := iter.Next() + if !ok { + break + } + + if parentSession != nil && (event.Action == nil || event.Action.Interrupted == nil) { + copied := copyTypedAgentEvent(event) + setAutomaticClose(copied) + setAutomaticClose(event) + parentSession.addEvent(copied) + } + + if event.Action != nil && event.Action.internalInterrupted != nil { + lastEvent = event + continue + } + + generator.Send(event) + lastEvent = event + } + + if lastEvent != nil && lastEvent.Action != nil { + if lastEvent.Action.internalInterrupted != nil { + events := isolatedSession.getEvents() + state := &deterministicTransferState{EventList: events} + compositeEvent := CompositeInterrupt(ctx, "deterministic transfer wrapper interrupted", + state, lastEvent.Action.internalInterrupted) + generator.Send(compositeEvent) + return + } + + if lastEvent.Action.Exit { + return + } + } + + sendTransferEvents(generator, toAgentNames) +} + +func sendTransferEvents(generator *AsyncGenerator[*AgentEvent], toAgentNames []string) { + for _, toAgentName := range toAgentNames { + aMsg, tMsg := GenTransferMessages(context.Background(), toAgentName) + + aEvent := EventFromMessage(aMsg, nil, schema.Assistant, "") + generator.Send(aEvent) + + tEvent := EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName) + tEvent.Action = &AgentAction{ + TransferToAgent: &TransferToAgentAction{ + DestAgentName: toAgentName, + }, + } + generator.Send(tEvent) + } +} diff --git a/adk/deterministic_transfer_test.go b/adk/deterministic_transfer_test.go new file mode 100644 index 0000000..da135da --- /dev/null +++ b/adk/deterministic_transfer_test.go @@ -0,0 +1,731 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +type dtTestStore struct { + data map[string][]byte +} + +func newDTTestStore() *dtTestStore { + return &dtTestStore{data: make(map[string][]byte)} +} + +func (s *dtTestStore) Set(_ context.Context, key string, value []byte) error { + s.data[key] = value + return nil +} + +func (s *dtTestStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.data[key] + return v, ok, nil +} + +type dtTestAgent struct { + name string + runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] + resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] +} + +func (a *dtTestAgent) Name(_ context.Context) string { + return a.name +} + +func (a *dtTestAgent) Description(_ context.Context) string { + return a.name + " description" +} + +func (a *dtTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return a.runFn(ctx, input, options...) +} + +func (a *dtTestAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + if a.resumeFn != nil { + return a.resumeFn(ctx, info, opts...) + } + return a.runFn(ctx, &AgentInput{}, opts...) +} + +func TestDeterministicTransferFlowAgentInterruptResume(t *testing.T) { + ctx := context.Background() + store := newDTTestStore() + + interruptData := "interrupt_data" + var runCount int + + innerAgent := &dtTestAgent{ + name: "inner", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + runCount++ + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(EventFromMessage(schema.AssistantMessage("before interrupt", nil), nil, schema.Assistant, "")) + intEvent := Interrupt(ctx, interruptData) + intEvent.Action.Interrupted.Data = interruptData + gen.Send(intEvent) + }() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + runCount++ + + assert.True(t, info.WasInterrupted, "innerAgent resumeFn: should be interrupted") + assert.True(t, info.IsResumeTarget, "innerAgent resumeFn: should be resume target") + + runCtx := getRunCtx(ctx) + assert.NotNil(t, runCtx, "innerAgent resumeFn: runCtx should not be nil") + assert.NotNil(t, runCtx.Session, "innerAgent resumeFn: runCtx.Session should not be nil") + + var agentEvents []*AgentEvent + for _, ev := range runCtx.Session.Events { + if ev.AgentEvent != nil { + agentEvents = append(agentEvents, ev.AgentEvent) + } + } + + assert.Len(t, agentEvents, 1, "innerAgent resumeFn: should have exactly 1 agent event") + if len(agentEvents) == 1 { + ev := agentEvents[0] + assert.Equal(t, "inner", ev.AgentName, "innerAgent resumeFn: event should be from inner agent") + assert.Equal(t, "before interrupt", ev.Output.MessageOutput.Message.Content, "innerAgent resumeFn: event content should be 'before interrupt'") + assert.Len(t, ev.RunPath, 2, "innerAgent resumeFn: RunPath should have 2 steps (outer agent, inner agent)") + if len(ev.RunPath) == 2 { + assert.Equal(t, "outer", ev.RunPath[0].agentName, "innerAgent resumeFn: RunPath[0] should be outer agent") + assert.Equal(t, "inner", ev.RunPath[1].agentName, "innerAgent resumeFn: RunPath[1] should be inner agent") + } + } + + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(EventFromMessage(schema.AssistantMessage("after resume", nil), nil, schema.Assistant, "")) + }() + return iter + }, + } + + innerFlowAgent := toFlowAgent(ctx, innerAgent) + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: innerFlowAgent, + ToAgentNames: []string{"next_agent"}, + }) + + outerAgent := &dtTestAgent{ + name: "outer", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return wrapped.Run(ctx, input, options...) + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.WasInterrupted, "outerAgent resumeFn: should be interrupted") + + runCtx := getRunCtx(ctx) + assert.NotNil(t, runCtx, "outerAgent resumeFn: runCtx should not be nil") + assert.NotNil(t, runCtx.Session, "outerAgent resumeFn: runCtx.Session should not be nil") + + var agentEvents []*AgentEvent + for _, ev := range runCtx.Session.Events { + if ev.AgentEvent != nil { + agentEvents = append(agentEvents, ev.AgentEvent) + } + } + + assert.Len(t, agentEvents, 1, "outerAgent resumeFn: should have exactly 1 agent event") + if len(agentEvents) == 1 { + ev := agentEvents[0] + assert.Equal(t, "inner", ev.AgentName, "outerAgent resumeFn: event should be from inner agent (preserved original)") + assert.Equal(t, "before interrupt", ev.Output.MessageOutput.Message.Content, "outerAgent resumeFn: event content should be 'before interrupt'") + assert.Len(t, ev.RunPath, 2, "outerAgent resumeFn: RunPath should have 2 steps") + if len(ev.RunPath) == 2 { + assert.Equal(t, "outer", ev.RunPath[0].agentName, "outerAgent resumeFn: RunPath[0] should be outer agent") + assert.Equal(t, "inner", ev.RunPath[1].agentName, "outerAgent resumeFn: RunPath[1] should be inner agent") + } + } + + ra := wrapped.(ResumableAgent) + return ra.Resume(ctx, info, opts...) + }, + } + + outerFlowAgent := toFlowAgent(ctx, outerAgent) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: outerFlowAgent, + EnableStreaming: true, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1")) + + var events []*AgentEvent + var interrupted bool + var interruptEvent *AgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + if ev.Action != nil && ev.Action.Interrupted != nil { + interrupted = true + interruptEvent = ev + } + } + + assert.Equal(t, 1, runCount, "run should have been called once") + assert.True(t, interrupted, "should have interrupted") + assert.Greater(t, len(events), 0, "should have events") + if interruptEvent == nil { + t.Fatal("should have interrupt event") + } + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, "should have interrupt contexts") + + _, exists, err := store.Get(ctx, "cp1") + assert.NoError(t, err) + assert.True(t, exists, "checkpoint should have been saved") + + var hasDeterministicTransferContext bool + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + t.Logf("InterruptContext: ID=%s, Info=%v, IsRootCause=%v, Addr=%v", intCtx.ID, intCtx.Info, intCtx.IsRootCause, intCtx.Address) + if intCtx.Info == "deterministic transfer wrapper interrupted" { + hasDeterministicTransferContext = true + } + for parent := intCtx.Parent; parent != nil; parent = parent.Parent { + t.Logf(" Parent: ID=%s, Info=%v, Addr=%v", parent.ID, parent.Info, parent.Address) + if parent.Info == "deterministic transfer wrapper interrupted" { + hasDeterministicTransferContext = true + } + } + } + assert.True(t, hasDeterministicTransferContext, "should have deterministic transfer interrupt context") + + var rootCauseID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + rootCauseID = intCtx.ID + break + } + } + assert.NotEmpty(t, rootCauseID, "should have root cause interrupt ID") + + resumeIter, err := runner.ResumeWithParams(ctx, "cp1", &ResumeParams{ + Targets: map[string]any{rootCauseID: nil}, + }) + assert.NoError(t, err) + + var resumeEvents []*AgentEvent + var resumeErr error + var hasTransfer bool + for { + ev, ok := resumeIter.Next() + if !ok { + break + } + + if ev.Err != nil { + resumeErr = ev.Err + } + if ev.Action != nil && ev.Action.TransferToAgent != nil { + hasTransfer = true + } + resumeEvents = append(resumeEvents, ev) + } + + assert.Equal(t, 2, runCount, "inner agent should be called twice (once for initial, once for resume)") + assert.NotEmpty(t, resumeEvents, "should have resume events") + assert.True(t, hasTransfer, "should have transfer action after resume") + assert.Error(t, resumeErr, "transfer should fail because next_agent doesn't exist") + assert.Contains(t, resumeErr.Error(), "next_agent", "error should mention the missing agent") +} + +func TestDeterministicTransferRunPathPreserved(t *testing.T) { + ctx := context.Background() + store := newDTTestStore() + + var collectedRunPaths [][]RunStep + + innerAgent := &dtTestAgent{ + name: "inner", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("from inner", nil), nil, schema.Assistant, "") + gen.Send(ev) + }() + return iter + }, + } + + innerFlowAgent := toFlowAgent(ctx, innerAgent) + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: innerFlowAgent, + ToAgentNames: []string{}, + }) + + outerAgent := &dtTestAgent{ + name: "outer", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + innerIter := wrapped.Run(ctx, input, options...) + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + for { + ev, ok := innerIter.Next() + if !ok { + break + } + collectedRunPaths = append(collectedRunPaths, ev.RunPath) + gen.Send(ev) + } + }() + return iter + }, + } + + outerFlowAgent := toFlowAgent(ctx, outerAgent) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: outerFlowAgent, + EnableStreaming: true, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1")) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.NotEmpty(t, collectedRunPaths, "should have collected RunPaths") + for _, rp := range collectedRunPaths { + assert.Len(t, rp, 2, "RunPath should have 2 steps (outer agent, inner agent)") + if len(rp) == 2 { + assert.Equal(t, "outer", rp[0].agentName, "RunPath[0] should be outer agent") + assert.Equal(t, "inner", rp[1].agentName, "RunPath[1] should be inner agent") + } + } +} + +func TestDeterministicTransferExitSkipsTransfer(t *testing.T) { + ctx := context.Background() + store := newDTTestStore() + + innerAgent := &dtTestAgent{ + name: "inner", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("inner exits", nil), nil, schema.Assistant, "") + ev.Action = &AgentAction{Exit: true} + gen.Send(ev) + }() + return iter + }, + } + + innerFlowAgent := toFlowAgent(ctx, innerAgent) + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: innerFlowAgent, + ToAgentNames: []string{"next_agent"}, + }) + + var outerSawExit bool + var transferGenerated bool + + outerAgent := &dtTestAgent{ + name: "outer", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + innerIter := wrapped.Run(ctx, input, options...) + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + for { + ev, ok := innerIter.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Exit { + outerSawExit = true + } + if ev.Action != nil && ev.Action.TransferToAgent != nil { + transferGenerated = true + } + gen.Send(ev) + } + }() + return iter + }, + } + + outerFlowAgent := toFlowAgent(ctx, outerAgent) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: outerFlowAgent, + EnableStreaming: true, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1")) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, outerSawExit, "outer should see exit event from inner") + assert.False(t, transferGenerated, "transfer should not be generated when inner exits") +} + +type nonFlowTestAgent struct { + name string + runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] + resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] +} + +func (a *nonFlowTestAgent) Name(_ context.Context) string { + return a.name +} + +func (a *nonFlowTestAgent) Description(_ context.Context) string { + return a.name + " description" +} + +func (a *nonFlowTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return a.runFn(ctx, input, options...) +} + +func (a *nonFlowTestAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + if a.resumeFn != nil { + return a.resumeFn(ctx, info, opts...) + } + return a.runFn(ctx, &AgentInput{}, opts...) +} + +type nonResumableTestAgent struct { + name string + runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] +} + +func (a *nonResumableTestAgent) Name(_ context.Context) string { + return a.name +} + +func (a *nonResumableTestAgent) Description(_ context.Context) string { + return a.name + " description" +} + +func (a *nonResumableTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return a.runFn(ctx, input, options...) +} + +func TestDeterministicTransferNonFlowAgent_ExitSkipsTransfer(t *testing.T) { + ctx := context.Background() + + agent := &nonFlowTestAgent{ + name: "test_agent", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("exiting", nil), nil, schema.Assistant, "") + ev.Action = &AgentAction{Exit: true} + gen.Send(ev) + }() + return iter + }, + } + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: agent, + ToAgentNames: []string{"next_agent"}, + }) + + iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + var events []*AgentEvent + var sawExit bool + var sawTransfer bool + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + if ev.Action != nil && ev.Action.Exit { + sawExit = true + } + if ev.Action != nil && ev.Action.TransferToAgent != nil { + sawTransfer = true + } + } + + assert.True(t, sawExit, "should see exit event") + assert.False(t, sawTransfer, "should NOT see transfer when exit is last event") + assert.Len(t, events, 1, "should have exactly 1 event (exit)") +} + +func TestDeterministicTransferNonFlowAgent_AppendsTransfer(t *testing.T) { + ctx := context.Background() + + agent := &nonFlowTestAgent{ + name: "test_agent", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("normal output", nil), nil, schema.Assistant, "") + gen.Send(ev) + }() + return iter + }, + } + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: agent, + ToAgentNames: []string{"next_agent"}, + }) + + iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + var events []*AgentEvent + var sawTransfer bool + var transferTarget string + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + if ev.Action != nil && ev.Action.TransferToAgent != nil { + sawTransfer = true + transferTarget = ev.Action.TransferToAgent.DestAgentName + } + } + + assert.True(t, sawTransfer, "should see transfer event after normal completion") + assert.Equal(t, "next_agent", transferTarget, "transfer target should be next_agent") + assert.Greater(t, len(events), 1, "should have more than 1 event (output + transfer messages)") +} + +func TestDeterministicTransferNonFlowAgent_InterruptSkipsTransfer(t *testing.T) { + ctx := context.Background() + + agent := &nonFlowTestAgent{ + name: "test_agent", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := &AgentEvent{ + Action: &AgentAction{ + Interrupted: &InterruptInfo{Data: "test interrupt"}, + }, + } + gen.Send(ev) + }() + return iter + }, + } + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: agent, + ToAgentNames: []string{"next_agent"}, + }) + + iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + var events []*AgentEvent + var sawInterrupt bool + var sawTransfer bool + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + if ev.Action != nil && ev.Action.Interrupted != nil { + sawInterrupt = true + } + if ev.Action != nil && ev.Action.TransferToAgent != nil { + sawTransfer = true + } + } + + assert.True(t, sawInterrupt, "should see interrupt event") + assert.False(t, sawTransfer, "should NOT see transfer when interrupted") + assert.Len(t, events, 1, "should have exactly 1 event (interrupt)") +} + +func TestDeterministicTransferNonFlowAgent_Resume(t *testing.T) { + ctx := context.Background() + + var resumeCalled bool + agent := &nonFlowTestAgent{ + name: "test_agent", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("from run", nil), nil, schema.Assistant, "") + gen.Send(ev) + }() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + resumeCalled = true + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("from resume", nil), nil, schema.Assistant, "") + gen.Send(ev) + }() + return iter + }, + } + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: agent, + ToAgentNames: []string{"next_agent"}, + }) + + ra, ok := wrapped.(ResumableAgent) + assert.True(t, ok, "wrapped agent should be ResumableAgent") + + iter := ra.Resume(ctx, &ResumeInfo{WasInterrupted: true}) + + var sawTransfer bool + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.TransferToAgent != nil { + sawTransfer = true + } + } + + assert.True(t, resumeCalled, "resume should have been called on inner agent") + assert.True(t, sawTransfer, "should see transfer event after resume completion") +} + +func TestDeterministicTransferFlowAgent_ResumeWithInvalidState(t *testing.T) { + ctx := context.Background() + + innerAgent := &dtTestAgent{ + name: "inner", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(EventFromMessage(schema.AssistantMessage("test", nil), nil, schema.Assistant, "")) + }() + return iter + }, + } + + innerFlowAgent := toFlowAgent(ctx, innerAgent) + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: innerFlowAgent, + ToAgentNames: []string{"next_agent"}, + }) + + ra, ok := wrapped.(ResumableAgent) + assert.True(t, ok, "wrapped flowAgent should be ResumableAgent") + + iter := ra.Resume(ctx, &ResumeInfo{ + WasInterrupted: true, + InterruptState: nil, + }) + + var gotError bool + var errorMsg string + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + gotError = true + errorMsg = ev.Err.Error() + } + } + + assert.True(t, gotError, "should get error for invalid state") + assert.Contains(t, errorMsg, "invalid interrupt state", "error should mention invalid state") +} + +func TestDeterministicTransferNonResumableAgent(t *testing.T) { + ctx := context.Background() + + agent := &nonResumableTestAgent{ + name: "non_resumable", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + ev := EventFromMessage(schema.AssistantMessage("output", nil), nil, schema.Assistant, "") + gen.Send(ev) + }() + return iter + }, + } + + wrapped := AgentWithDeterministicTransferTo(ctx, &DeterministicTransferConfig{ + Agent: agent, + ToAgentNames: []string{"next_agent"}, + }) + + _, isResumable := wrapped.(ResumableAgent) + assert.False(t, isResumable, "wrapped non-resumable agent should NOT be ResumableAgent") + + assert.Equal(t, "non_resumable", wrapped.Name(ctx), "Name should delegate to inner agent") + assert.Equal(t, "non_resumable description", wrapped.Description(ctx), "Description should delegate to inner agent") + + iter := wrapped.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + var sawTransfer bool + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.TransferToAgent != nil { + sawTransfer = true + } + } + + assert.True(t, sawTransfer, "should see transfer event") +} diff --git a/adk/failover_chatmodel.go b/adk/failover_chatmodel.go new file mode 100644 index 0000000..f12d890 --- /dev/null +++ b/adk/failover_chatmodel.go @@ -0,0 +1,490 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "io" + "log" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type failoverCurrentModelKey struct{} + +func typedSetFailoverCurrentModel[M MessageType](ctx context.Context, currentModel model.BaseModel[M]) context.Context { + return context.WithValue(ctx, failoverCurrentModelKey{}, currentModel) +} + +func typedGetFailoverCurrentModel[M MessageType](ctx context.Context) (model.BaseModel[M], bool) { + m, ok := ctx.Value(failoverCurrentModelKey{}).(model.BaseModel[M]) + return m, ok +} + +type failoverHasMoreAttemptsKey struct{} + +// withFailoverHasMoreAttempts sets a flag in context indicating whether additional failover +// attempts remain after the current one. This is read by buildErrWrapper to decide whether +// stream errors should be wrapped as WillRetryError. +func withFailoverHasMoreAttempts(ctx context.Context, hasMore bool) context.Context { + return context.WithValue(ctx, failoverHasMoreAttemptsKey{}, hasMore) +} + +// getFailoverHasMoreAttempts returns true if the current failover attempt has more attempts +// after it, false otherwise (including when no failover context is present). +func getFailoverHasMoreAttempts(ctx context.Context) bool { + v, _ := ctx.Value(failoverHasMoreAttemptsKey{}).(bool) + return v +} + +type typedFailoverProxyModel[M MessageType] struct { +} + +func (m *typedFailoverProxyModel[M]) prepareTarget(ctx context.Context) (model.BaseModel[M], error) { + target, ok := typedGetFailoverCurrentModel[M](ctx) + if !ok { + return nil, errors.New("failover current model not found in context") + } + + if !components.IsCallbacksEnabled(target) { + target = typedCallbackInjectionModelWrapper[M]{}.wrapModel(target) + } + + return target, nil +} + +func (m *typedFailoverProxyModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + target, err := m.prepareTarget(ctx) + if err != nil { + var zero M + return zero, err + } + + return target.Generate(ctx, input, opts...) +} + +func (m *typedFailoverProxyModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + target, err := m.prepareTarget(ctx) + if err != nil { + return nil, err + } + + return target.Stream(ctx, input, opts...) +} + +func (m *typedFailoverProxyModel[M]) IsCallbacksEnabled() bool { + return true +} + +func (m *typedFailoverProxyModel[M]) GetType() string { + return "FailoverProxyModel" +} + +type failoverProxyModel = typedFailoverProxyModel[*schema.Message] + +// FailoverContext contains context information during failover process. +type FailoverContext[M MessageType] struct { + // FailoverAttempt is the current failover attempt number, starting from 1. + FailoverAttempt uint + + // InputMessages is the original input messages before any transformation. + InputMessages []M + + // LastOutputMessage is the output message from the last failed attempt. + // May be nil if no output was produced. For streaming, this may be a partial message + // already received before the stream error. + LastOutputMessage M + + // LastErr is the error from the last failed attempt that triggered this failover. + // + // Note: When ModelRetryConfig is also configured, LastErr will be a *RetryExhaustedError + // (if retries were exhausted) rather than the original model error. The original error + // can be retrieved via RetryExhaustedError.LastErr. + LastErr error +} + +// ModelFailoverConfig configures failover behavior for ChatModel. +// When configured, each ChatModel call first tries the last successful model (initially the configured Model), +// and if that fails, calls GetFailoverModel to select alternate models. +type ModelFailoverConfig[M MessageType] struct { + // MaxRetries specifies the maximum number of failover attempts. + // + // When failover is triggered, GetFailoverModel will be called up to MaxRetries times + // (FailoverAttempt starts from 1). If GetFailoverModel returns an error, failover + // stops immediately and that error is returned. + // + // A value of 0 means no failover (GetFailoverModel will not be called). + // A value of 1 means GetFailoverModel may be called once. + // + // Note: if lastSuccessModel is set (from a previous successful call), it will be tried + // first before calling GetFailoverModel. + MaxRetries uint + + // ShouldFailover determines whether to fail over to the next model when an error occurs. + // It receives the output message (may be nil/zero if no output is available) and the error (non-nil on failure). + // For streaming errors, outputMessage can carry a partial message accumulated before the error. + // + // Note: When ModelRetryConfig is also configured, outputErr will be a *RetryExhaustedError + // (if retries were exhausted) rather than the original model error. Use errors.As to extract + // the RetryExhaustedError and access RetryExhaustedError.LastErr for the original error. + // + // Note: When the context itself is cancelled (ctx.Err() != nil), failover will stop immediately + // regardless of this function. However, if the model returns context.Canceled or context.DeadlineExceeded + // as an error while the context is still active, this function will still be called. + // Should not be nil when ModelFailoverConfig is set. + // Return true to fail over to the next model, false to stop and return the current result/error. + ShouldFailover func(ctx context.Context, outputMessage M, outputErr error) bool + + // GetFailoverModel is called when a model call fails and ShouldFailover returns true. + // It selects the next model to use for the failover attempt and optionally transforms input messages. + // It receives the failover context containing attempt number (starting from 1), original input, and last result. + // Return values: + // - failoverModel: The model to use for this failover attempt. + // - failoverModelInputMessages: The transformed input messages for the failover model. If nil, will use original input. + // - failoverErr: If non-nil, failover stops and this error is returned. + // Should not be nil when ModelFailoverConfig is set via ChatModelAgentConfig. + GetFailoverModel func(ctx context.Context, failoverCtx *FailoverContext[M]) ( + failoverModel model.BaseModel[M], failoverModelInputMessages []M, failoverErr error) +} + +func typedGetFailoverLastSuccessModel[M MessageType](ctx context.Context) model.BaseModel[M] { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil { + return nil + } + return execCtx.failoverLastSuccessModel +} + +func typedSetFailoverLastSuccessModel[M MessageType](ctx context.Context, m model.BaseModel[M]) { + if execCtx := getTypedChatModelAgentExecCtx[M](ctx); execCtx != nil { + execCtx.failoverLastSuccessModel = m + } +} + +type failoverModelWrapper[M MessageType] struct { + config *ModelFailoverConfig[M] + inner model.BaseModel[M] +} + +func newFailoverModelWrapper[M MessageType](inner model.BaseModel[M], config *ModelFailoverConfig[M]) *failoverModelWrapper[M] { + return &failoverModelWrapper[M]{ + config: config, + inner: inner, + } +} + +func (f *failoverModelWrapper[M]) needFailover(ctx context.Context, outputMessage M, outputErr error) bool { + if ctx.Err() != nil { + return false + } + + if _, ok := compose.ExtractInterruptInfo(outputErr); ok { + return false + } + + // ErrStreamCanceled means the caller voluntarily abandoned the stream; + // never retry or fail over in this case. + if errors.Is(outputErr, ErrStreamCanceled) { + return false + } + + // ShouldFailover is validated at agent construction; nil here indicates a programmer error. + return f.config.ShouldFailover(ctx, outputMessage, outputErr) +} + +func (f *failoverModelWrapper[M]) getFailoverModel(ctx context.Context, failoverCtx *FailoverContext[M]) (model.BaseModel[M], []M, error) { + currentModel, msgs, err := f.config.GetFailoverModel(ctx, failoverCtx) + if err != nil { + return nil, nil, err + } + if currentModel == nil { + return nil, nil, nil + } + return currentModel, msgs, nil +} + +func (f *failoverModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + // Defensive: GetFailoverModel is validated non-nil at agent construction. + if f.config.GetFailoverModel == nil { + return f.inner.Generate(ctx, input, opts...) + } + + var lastOutputMessage M + var lastErr error + + // Try lastSuccessModel first if available. + if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil { + if err := ctx.Err(); err != nil { + var zero M + return zero, err + } + + modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess) + modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0) + result, err := f.inner.Generate(modelCtx, input, opts...) + if err == nil { + return result, nil + } + + lastOutputMessage = result + lastErr = err + + if !f.needFailover(ctx, result, err) { + return result, err + } + + log.Printf("failover ChatModel.Generate lastSuccessModel failed: %v", err) + } + + for attempt := uint(1); attempt <= f.config.MaxRetries; attempt++ { + if err := ctx.Err(); err != nil { + var zero M + return zero, err + } + + failoverCtx := &FailoverContext[M]{ + FailoverAttempt: attempt, + InputMessages: input, + LastOutputMessage: lastOutputMessage, + LastErr: lastErr, + } + + currentModel, currentInput, err := f.getFailoverModel(ctx, failoverCtx) + if err != nil { + var zero M + return zero, err + } + if currentModel == nil { + var zero M + return zero, fmt.Errorf("failover GetFailoverModel returned nil model at attempt %d", attempt) + } + + if currentInput == nil { + currentInput = input + } + + modelCtx := typedSetFailoverCurrentModel(ctx, currentModel) + modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries) + result, err := f.inner.Generate(modelCtx, currentInput, opts...) + lastOutputMessage = result + lastErr = err + + if err == nil { + typedSetFailoverLastSuccessModel(ctx, currentModel) + return result, nil + } + + if !f.needFailover(ctx, result, err) { + return result, err + } + + if attempt < f.config.MaxRetries { + log.Printf("failover ChatModel.Generate attempt %d failed: %v", attempt, err) + } + } + + return lastOutputMessage, lastErr +} + +func (f *failoverModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) ( + *schema.StreamReader[M], error) { + // Defensive: GetFailoverModel is validated non-nil at agent construction. + if f.config.GetFailoverModel == nil { + return f.inner.Stream(ctx, input, opts...) + } + + var lastOutputMessage M + var lastErr error + + // Try lastSuccessModel first if available. + if lastSuccess := typedGetFailoverLastSuccessModel[M](ctx); lastSuccess != nil { + if err := ctx.Err(); err != nil { + return nil, err + } + + modelCtx := typedSetFailoverCurrentModel(ctx, lastSuccess) + modelCtx = withFailoverHasMoreAttempts(modelCtx, f.config.MaxRetries > 0) + stream, err := f.inner.Stream(modelCtx, input, opts...) + if err != nil { + lastErr = err + var zero M + if !f.needFailover(ctx, zero, err) { + return nil, err + } + log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", err) + } else { + copies := stream.Copy(2) + checkCopy := copies[0] + returnCopy := copies[1] + + outMsg, streamErr := typedConsumeStream(checkCopy) + if streamErr != nil { + lastOutputMessage = outMsg + lastErr = streamErr + returnCopy.Close() + + if !f.needFailover(ctx, outMsg, streamErr) { + return nil, streamErr + } + log.Printf("failover ChatModel.Stream lastSuccessModel failed: %v", streamErr) + } else { + return returnCopy, nil + } + } + } + + for attempt := uint(1); attempt <= f.config.MaxRetries; attempt++ { + if err := ctx.Err(); err != nil { + return nil, err + } + + failoverCtx := &FailoverContext[M]{ + FailoverAttempt: attempt, + InputMessages: input, + LastOutputMessage: lastOutputMessage, + LastErr: lastErr, + } + + currentModel, currentInput, err := f.getFailoverModel(ctx, failoverCtx) + if err != nil { + return nil, err + } + if currentModel == nil { + return nil, fmt.Errorf("failover GetFailoverModel returned nil model at attempt %d", attempt) + } + + if currentInput == nil { + currentInput = input + } + + modelCtx := typedSetFailoverCurrentModel(ctx, currentModel) + modelCtx = withFailoverHasMoreAttempts(modelCtx, attempt < f.config.MaxRetries) + stream, err := f.inner.Stream(modelCtx, currentInput, opts...) + if err != nil { + lastErr = err + var zero M + lastOutputMessage = zero + + if !f.needFailover(ctx, zero, err) { + return nil, err + } + + if attempt < f.config.MaxRetries { + log.Printf("failover ChatModel.Stream attempt %d failed: %v", attempt, err) + } + continue + } + + // The stream returned by f.inner.Stream is already Copy'd by the inner eventSender layer: one + // copy is forwarded to the client in real time via events. Therefore consuming a copy here does + // NOT block client-side streaming. + // + // We Copy the stream into two readers: + // - checkCopy: consumed synchronously to surface mid-stream errors and decide whether to fail over. + // - returnCopy: returned to the caller (stateModelWrapper), which also consumes synchronously to + // build state (AfterModelRewriteState), so waiting here adds no extra latency. + // + // If checkCopy errors and failover is allowed, we close returnCopy and retry with the next model. + // Otherwise we return returnCopy. + // + // NOTE on duplicate events during failover: when a retry happens, events from the failed attempt + // may already have been emitted to the client, and the retry will emit a new stream. Client-side + // handlers are expected to handle multiple rounds (e.g., reset on retry or deduplicate by attempt + // metadata). + copies := stream.Copy(2) + checkCopy := copies[0] + returnCopy := copies[1] + + outMsg, streamErr := typedConsumeStream(checkCopy) + if streamErr != nil { + lastOutputMessage = outMsg + lastErr = streamErr + returnCopy.Close() + + if !f.needFailover(ctx, outMsg, streamErr) { + return nil, streamErr + } + + if attempt < f.config.MaxRetries { + log.Printf("failover ChatModel.Stream attempt %d failed: %v", attempt, streamErr) + } + continue + } + + typedSetFailoverLastSuccessModel(ctx, currentModel) + return returnCopy, nil + } + + return nil, lastErr +} + +func typedConsumeStream[M MessageType](stream *schema.StreamReader[M]) (M, error) { + var zero M + defer stream.Close() + + switch s := any(stream).(type) { + case *schema.StreamReader[*schema.Message]: + chunks := make([]*schema.Message, 0) + for { + chunk, err := s.Recv() + if err == io.EOF { + break + } + if err != nil { + msg, _ := schema.ConcatMessages(chunks) + if msg != nil { + return any(msg).(M), err + } + return zero, err + } + chunks = append(chunks, chunk) + } + msg, _ := schema.ConcatMessages(chunks) + if msg != nil { + return any(msg).(M), nil + } + return zero, nil + case *schema.StreamReader[*schema.AgenticMessage]: + chunks := make([]*schema.AgenticMessage, 0) + for { + chunk, err := s.Recv() + if err == io.EOF { + break + } + if err != nil { + msg, _ := schema.ConcatAgenticMessages(chunks) + if msg != nil { + return any(msg).(M), err + } + return zero, err + } + chunks = append(chunks, chunk) + } + msg, _ := schema.ConcatAgenticMessages(chunks) + if msg != nil { + return any(msg).(M), nil + } + return zero, nil + default: + panic("unreachable: unknown MessageType") + } +} diff --git a/adk/failover_chatmodel_test.go b/adk/failover_chatmodel_test.go new file mode 100644 index 0000000..8b8ca57 --- /dev/null +++ b/adk/failover_chatmodel_test.go @@ -0,0 +1,742 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "io" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type fakeChatModel struct { + generate func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) + stream func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) + callbacksEnabled bool +} + +func (m *fakeChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return m.generate(ctx, input, opts...) +} + +func (m *fakeChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return m.stream(ctx, input, opts...) +} + +func (m *fakeChatModel) IsCallbacksEnabled() bool { + return m.callbacksEnabled +} + +func drainMessageStream(sr *schema.StreamReader[*schema.Message]) ([]*schema.Message, error) { + defer sr.Close() + var out []*schema.Message + for { + chunk, err := sr.Recv() + if err == io.EOF { + return out, nil + } + if err != nil { + return out, err + } + out = append(out, chunk) + } +} + +func streamWithMidError(chunks []*schema.Message, err error) *schema.StreamReader[*schema.Message] { + sr, sw := schema.Pipe[*schema.Message](2) + go func() { + defer sw.Close() + for _, c := range chunks { + sw.Send(c, nil) + } + sw.Send(nil, err) + }() + return sr +} + +func streamWithMidErrorControlled(chunks []*schema.Message, err error, firstSent chan struct{}, release chan struct{}) *schema.StreamReader[*schema.Message] { + sr, sw := schema.Pipe[*schema.Message](2) + go func() { + defer sw.Close() + for i, c := range chunks { + sw.Send(c, nil) + if i == 0 && firstSent != nil { + close(firstSent) + if release != nil { + <-release + } + } + } + sw.Send(nil, err) + }() + return sr +} + +func TestFailoverCurrentModelContext(t *testing.T) { + t.Run("set and get", func(t *testing.T) { + ctx := context.Background() + m := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil + }, + } + ctx = typedSetFailoverCurrentModel[*schema.Message](ctx, m) + got, ok := typedGetFailoverCurrentModel[*schema.Message](ctx) + require.True(t, ok) + require.Same(t, m, got) + }) + + t.Run("wrong type", func(t *testing.T) { + ctx := context.WithValue(context.Background(), failoverCurrentModelKey{}, "bad") + _, ok := typedGetFailoverCurrentModel[*schema.Message](ctx) + require.False(t, ok) + }) + + t.Run("missing", func(t *testing.T) { + _, ok := typedGetFailoverCurrentModel[*schema.Message](context.Background()) + require.False(t, ok) + }) +} + +func TestFailoverProxyModel(t *testing.T) { + t.Run("generate missing context", func(t *testing.T) { + p := &failoverProxyModel{} + _, err := p.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + }) + + t.Run("stream missing context", func(t *testing.T) { + p := &failoverProxyModel{} + _, err := p.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + }) + + t.Run("generate routes to current model", func(t *testing.T) { + var called int32 + target := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&called, 1) + return schema.AssistantMessage("routed", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("routed", nil)}), nil + }, + } + ctx := typedSetFailoverCurrentModel[*schema.Message](context.Background(), target) + p := &failoverProxyModel{} + msg, err := p.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "routed", msg.Content) + require.Equal(t, int32(1), atomic.LoadInt32(&called)) + }) +} + +func TestFailoverModelWrapper_Generate(t *testing.T) { + t.Run("delegates when GetFailoverModel nil", func(t *testing.T) { + var called int32 + inner := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&called, 1) + return schema.AssistantMessage("inner", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("inner", nil)}), nil + }, + } + w := newFailoverModelWrapper[*schema.Message](inner, &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 2, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: nil, + }) + msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "inner", msg.Content) + require.Equal(t, int32(1), atomic.LoadInt32(&called)) + }) + + t.Run("failover to second model", func(t *testing.T) { + wantErr := errors.New("first failed") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, wantErr + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + return errors.Is(err, wantErr) + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok", msg.Content) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("canceled error delegates to ShouldFailover", func(t *testing.T) { + var shouldCalls int32 + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, context.Canceled + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 5, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + // User decides to stop on canceled error + return !errors.Is(err, context.Canceled) + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m1, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.ErrorIs(t, err, context.Canceled) + // ShouldFailover is called once and returns false, stopping failover + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("stops when GetFailoverModel returns error", func(t *testing.T) { + wantErr := errors.New("get model failed") + var called int32 + inner := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&called, 1) + return schema.AssistantMessage("unused", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, wantErr + }, + } + + w := newFailoverModelWrapper[*schema.Message](inner, cfg) + _, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.ErrorIs(t, err, wantErr) + require.Equal(t, int32(0), atomic.LoadInt32(&called)) + }) + + t.Run("stops when GetFailoverModel returns nil model", func(t *testing.T) { + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, msg) + require.Error(t, err) + require.ErrorContains(t, err, "GetFailoverModel returned nil model") + }) +} + +func TestFailoverModelWrapper_Stream(t *testing.T) { + t.Run("returns stream when first attempt succeeds", func(t *testing.T) { + var shouldCalls int32 + in := schema.UserMessage("hi") + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + require.Len(t, input, 1) + require.Same(t, in, input[0]) + return schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("a", nil), + schema.AssistantMessage("b", nil), + }), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 0, + ShouldFailover: func(context.Context, *schema.Message, error) bool { + atomic.AddInt32(&shouldCalls, 1) + return false + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m1, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := w.Stream(ctx, []*schema.Message{in}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 2) + require.Equal(t, "a", msgs[0].Content) + require.Equal(t, "b", msgs[1].Content) + require.Equal(t, int32(0), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("failover when Stream returns error immediately", func(t *testing.T) { + wantErr := errors.New("stream init failed") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return nil, wantErr + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + return errors.Is(err, wantErr) + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "ok", msgs[0].Content) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("failover when stream errors mid-way", func(t *testing.T) { + streamErr := errors.New("mid error") + var shouldCalls int32 + var seenOutput atomic.Value + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p1", nil), + schema.AssistantMessage("p2", nil), + }, streamErr), nil + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final", nil)}), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + if errors.Is(err, streamErr) && out != nil { + seenOutput.Store(out.Content) + } + return errors.Is(err, streamErr) + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "final", msgs[0].Content) + require.Equal(t, "p1p2", seenOutput.Load()) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("stop when ShouldFailover returns false for mid-way error", func(t *testing.T) { + streamErr := errors.New("mid error") + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return streamWithMidError([]*schema.Message{schema.AssistantMessage("p", nil)}, streamErr), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(context.Context, *schema.Message, error) bool { + return false + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m1, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, sr) + require.ErrorIs(t, err, streamErr) + }) + + t.Run("canceled mid-way error delegates to ShouldFailover", func(t *testing.T) { + var shouldCalls int32 + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return streamWithMidError([]*schema.Message{schema.AssistantMessage("p", nil)}, context.Canceled), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + // User decides to stop on canceled error + return !errors.Is(err, context.Canceled) + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m1, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, sr) + require.ErrorIs(t, err, context.Canceled) + // ShouldFailover is called once and returns false, stopping failover + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("stop when Stream returns error immediately and ShouldFailover returns false", func(t *testing.T) { + wantErr := errors.New("stream init failed") + var shouldCalls int32 + var m1Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return nil, wantErr + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + require.ErrorIs(t, err, wantErr) + return false + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m1, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, sr) + require.ErrorIs(t, err, wantErr) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) + }) + + t.Run("stops when GetFailoverModel returns nil model", func(t *testing.T) { + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, sr) + require.Error(t, err) + require.ErrorContains(t, err, "GetFailoverModel returned nil model") + }) + + t.Run("stops when GetFailoverModel returns error", func(t *testing.T) { + wantErr := errors.New("get model failed") + var called int32 + inner := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&called, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return true }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, wantErr + }, + } + + w := newFailoverModelWrapper[*schema.Message](inner, cfg) + sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")}) + require.Nil(t, sr) + require.ErrorIs(t, err, wantErr) + require.Equal(t, int32(0), atomic.LoadInt32(&called)) + }) + + t.Run("stops when ctx canceled during mid-way error handling", func(t *testing.T) { + midErr := errors.New("mid error") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + firstSent := make(chan struct{}) + release := make(chan struct{}) + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidErrorControlled( + []*schema.Message{schema.AssistantMessage("p", nil)}, + midErr, + firstSent, + release, + ), nil + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil + }, + } + + cfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(context.Context, *schema.Message, error) bool { + atomic.AddInt32(&shouldCalls, 1) + return true + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg) + baseCtx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + ctx, cancel := context.WithCancel(baseCtx) + type result struct { + sr *schema.StreamReader[*schema.Message] + err error + } + ch := make(chan result, 1) + go func() { + sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + ch <- result{sr: sr, err: err} + }() + + <-firstSent + cancel() + close(release) + + res := <-ch + if res.sr != nil { + res.sr.Close() + } + require.Nil(t, res.sr) + require.ErrorIs(t, res.err, midErr) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&shouldCalls)) + }) +} + +func TestTypedConsumeStream_EmptyAgenticStream(t *testing.T) { + sr, sw := schema.Pipe[*schema.AgenticMessage](1) + sw.Close() + + msg, err := typedConsumeStream(sr) + assert.Nil(t, err, "empty stream should not return error") + assert.NotNil(t, msg, "empty stream should return non-nil message from ConcatAgenticMessages") +} + +func TestTypedConsumeStream_AgenticMidStreamError(t *testing.T) { + midErr := errors.New("mid-stream failure") + sr := streamWithMidErrorAgentic( + []*schema.AgenticMessage{agenticChunk("chunk1"), agenticChunk("chunk2")}, + midErr, + ) + + msg, err := typedConsumeStream(sr) + assert.ErrorIs(t, err, midErr, "should return the mid-stream error") + assert.NotNil(t, msg, "should return concatenated partial message from received chunks") +} + +func streamWithMidErrorAgentic(chunks []*schema.AgenticMessage, err error) *schema.StreamReader[*schema.AgenticMessage] { + sr, sw := schema.Pipe[*schema.AgenticMessage](len(chunks) + 1) + go func() { + defer sw.Close() + for _, c := range chunks { + sw.Send(c, nil) + } + sw.Send(nil, err) + }() + return sr +} + +func agenticChunk(text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: text}), + }, + } +} diff --git a/adk/filesystem/backend.go b/adk/filesystem/backend.go new file mode 100644 index 0000000..62ebee8 --- /dev/null +++ b/adk/filesystem/backend.go @@ -0,0 +1,304 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package filesystem provides file system operations. +package filesystem + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// FileInfo represents basic file metadata information. +type FileInfo struct { + // Path is the path of the file or directory, which can be a filename, relative path, or absolute path. + Path string + + // IsDir indicates whether the entry is a directory. + // true for directories, false for regular files. + IsDir bool + + // Size is the file size in bytes. + // For directories, this value may be 0 or platform-dependent. + Size int64 + + // ModifiedAt is the last modification time in ISO 8601 format. + // Example: "2025-01-15T10:30:00Z" + ModifiedAt string +} + +// GrepMatch represents a single pattern match result. +type GrepMatch struct { + Content string + + // Path is the file path where the match was found. + Path string + + // Line is the 1-based line number of the match. + Line int +} + +// LsInfoRequest contains parameters for listing file information. +type LsInfoRequest struct { + // Path specifies the directory path to list. + Path string +} + +// ReadRequest contains parameters for reading file content. +type ReadRequest struct { + // FilePath is the path to the file to be read. + FilePath string + + // Offset specifies the starting line number (1-based) for reading. + // Line 1 is the first line of the file. + // Use this when the file is too large to read at once. + // Defaults to 1 (start from the first line). + // Values < 1 will be treated as 1. + Offset int + + // Limit specifies the maximum number of lines to read. + // When Limit is 0 (default), the entire file content is returned. + Limit int +} + +// MultiModalReadRequest extends ReadRequest with parameters only applicable +// to MultiModalReader implementations (e.g. PDF page ranges). +type MultiModalReadRequest struct { + ReadRequest + + // Pages specifies the page range for PDF files (e.g. "1-5", "3", "10-20"). + Pages string +} + +// GrepRequest contains parameters for searching file content. +type GrepRequest struct { + // ===== Search Parameters ===== + + // Pattern is the search pattern, supports full regular expression syntax. + // Uses ripgrep syntax (not grep). Examples: + // - "log.*Error" matches lines with "log" followed by "Error" + // - "function\\s+\\w+" matches "function" followed by whitespace and word characters + // - Literal braces need escaping: "interface\\{\\}" matches "interface{}" + Pattern string + + // Path is an optional directory path to limit the search scope. + Path string + + // ===== File Filtering ===== + + // Glob is an optional pattern to filter the files to be searched. + // It filters by file path, not content. If empty, no files are filtered. + // Supports standard glob wildcards: + // - `*` matches any characters except path separators. + // - `**` matches any directories recursively. + // - `?` matches a single character. + // - `[abc]` matches one character from the set. + Glob string + + // FileType is the file type filter, e.g., "js", "py", "rust". + // More efficient than Glob for standard file types. + FileType string + + // ===== Search Options ===== + + // CaseInsensitive enables case insensitive search. + CaseInsensitive bool + + // EnableMultiline enables multiline mode where patterns can span lines. + // Default: false (patterns match within single lines only). + EnableMultiline bool + + // ===== Context Display (Content mode only) ===== + + // AfterLines shows N lines after each match. + // Only applicable when OutputMode is "content". + // Values <= 0 are treated as unset. + AfterLines int + + // BeforeLines shows N lines before each match. + // Only applicable when OutputMode is "content". + // Values <= 0 are treated as unset. + BeforeLines int +} + +// GlobInfoRequest contains parameters for glob pattern matching. +type GlobInfoRequest struct { + // Pattern is the glob expression used to match file paths. + // It supports standard glob syntax: + // - `*` matches any characters except path separators. + // - `**` matches any directories recursively. + // - `?` matches a single character. + // - `[abc]` matches one character from the set. + Pattern string + + // Path is the base directory from which to start the search. + Path string +} + +// WriteRequest contains parameters for writing file content. +type WriteRequest struct { + // FilePath is the path of the file to write. + FilePath string + + // Content is the data to be written to the file. + Content string +} + +// EditRequest contains parameters for editing file content. +type EditRequest struct { + // FilePath is the path of the file to edit. + FilePath string + + // OldString is the exact string to be replaced. It must be non-empty and will be matched literally, including whitespace. + OldString string + + // NewString is the string that will replace OldString. + // It must be different from OldString. + // An empty string can be used to effectively delete OldString. + NewString string + + // ReplaceAll controls the replacement behavior. + // If true, all occurrences of OldString are replaced. + // If false, the operation fails unless OldString appears exactly once in the file. + ReplaceAll bool +} + +// FileContentPartType defines the type of a multimodal file content part. +type FileContentPartType string + +const ( + // FileContentPartTypeImage represents an image part (e.g. PNG, JPG). + FileContentPartTypeImage FileContentPartType = "image" + // FileContentPartTypePDF represents a file part (e.g. PDF). + FileContentPartTypePDF FileContentPartType = "pdf" +) + +// FileContentPart represents a multimodal part of file content. +// Data holds raw bytes; encoding (e.g. base64) is handled by the consumer. +type FileContentPart struct { + // Type is the kind of content this part represents. + // Required. + Type FileContentPartType + + // MIMEType is the MIME type of the content (e.g. "image/png", "application/pdf"). + // Required. + MIMEType string + + // Data is the raw binary content. + // Required. + Data []byte +} + +// FileContent holds the result of a Read operation. +type FileContent struct { + // Content holds the plain text content of the file. + Content string +} + +// MultiFileContent holds the result of a MultiModalRead operation. +// +// FileContent and Parts are mutually exclusive (one-of): +// - Set FileContent for plain text results (same as a normal Read). +// - Set Parts for multimodal results (images, PDFs, etc.). +// +// When Parts is non-empty, FileContent is ignored. +type MultiFileContent struct { + *FileContent + + // Parts holds multimodal output parts (e.g. image, PDF). + Parts []FileContentPart +} + +// MultiModalReader is an optional extension interface for Backend. +// Backends that implement this interface support multimodal file reading, +// returning structured parts (images, PDFs) instead of plain text. +// +// For large file handling, there are two approaches to control output size: +// - Implement size control within MultiModalRead (e.g. reject files exceeding a threshold, +// downsample images, or limit PDF page counts at the backend level). +// - Use ToolMiddleware's EnhancedInvokable to customize result transformation, +// or use the built-in reduction middleware with configurable policies. +type MultiModalReader interface { + MultiModalRead(ctx context.Context, req *MultiModalReadRequest) (*MultiFileContent, error) +} + +// Backend is a pluggable, unified file backend protocol interface. +// +// All methods use struct-based parameters to allow future extensibility +// without breaking backward compatibility. +type Backend interface { + // LsInfo lists file information under the given path. + // + // Returns: + // - []FileInfo: List of matching file information + // - error: Error if the operation fails + LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) + + // Read reads file content with support for line-based offset and limit. + // + // Returns: + // - string: The file content read + // - error: Error if file does not exist or read fails + Read(ctx context.Context, req *ReadRequest) (*FileContent, error) + + // GrepRaw searches for content matching the specified pattern in files. + // + // Returns: + // - []GrepMatch: List of all matching results + // - error: Error if the search fails + GrepRaw(ctx context.Context, req *GrepRequest) ([]GrepMatch, error) + + // GlobInfo returns file information matching the glob pattern. + // + // Returns: + // - []FileInfo: List of matching file information + // - error: Error if the pattern is invalid or operation fails + GlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error) + + // Write creates or updates file content. + // + // Returns: + // - error: Error if the write operation fails + Write(ctx context.Context, req *WriteRequest) error + + // Edit replaces string occurrences in a file. + // + // Returns: + // - error: Error if file does not exist, OldString is empty, or OldString is not found + Edit(ctx context.Context, req *EditRequest) error +} + +// ExecuteRequest contains parameters for executing a command. +type ExecuteRequest struct { + Command string // The command to execute + RunInBackendGround bool +} + +// ExecuteResponse contains the response result of command execution. +type ExecuteResponse struct { + Output string // Command output content + ExitCode *int // Command exit code + Truncated bool // Whether the output was truncated +} + +type Shell interface { + Execute(ctx context.Context, input *ExecuteRequest) (result *ExecuteResponse, err error) +} + +type StreamingShell interface { + ExecuteStreaming(ctx context.Context, input *ExecuteRequest) (result *schema.StreamReader[*ExecuteResponse], err error) +} diff --git a/adk/filesystem/backend_inmemory.go b/adk/filesystem/backend_inmemory.go new file mode 100644 index 0000000..1a81181 --- /dev/null +++ b/adk/filesystem/backend_inmemory.go @@ -0,0 +1,774 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "fmt" + "path/filepath" + "regexp" + "strings" + "sync" + "time" + + "github.com/bmatcuk/doublestar/v4" +) + +type fileEntry struct { + content string + modifiedAt time.Time +} + +// InMemoryBackend is an in-memory implementation of the Backend interface. +// It stores files in a map and is safe for concurrent use. +type InMemoryBackend struct { + mu sync.RWMutex + files map[string]*fileEntry +} + +// NewInMemoryBackend creates a new in-memory backend. +func NewInMemoryBackend() *InMemoryBackend { + return &InMemoryBackend{ + files: make(map[string]*fileEntry), + } +} + +// LsInfo lists file information under the given path. +func (b *InMemoryBackend) LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + // Normalize path + path := normalizePath(req.Path) + + var result []FileInfo + seen := make(map[string]bool) + dirInfo := make(map[string]*FileInfo) + + for filePath, entry := range b.files { + normalizedFilePath := normalizePath(filePath) + + // Check if file is under the given path + if path == "/" || strings.HasPrefix(normalizedFilePath, path+"/") || normalizedFilePath == path { + // For directory listing, we want to show immediate children + relativePath := strings.TrimPrefix(normalizedFilePath, path) + relativePath = strings.TrimPrefix(relativePath, "/") + + if relativePath == "" { + // The path itself is a file + if !seen[normalizedFilePath] { + result = append(result, FileInfo{ + Path: filepath.Base(normalizedFilePath), + IsDir: false, + Size: int64(len(entry.content)), + ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano), + }) + seen[normalizedFilePath] = true + } + continue + } + + // Get the first segment (immediate child) + parts := strings.SplitN(relativePath, "/", 2) + if len(parts) > 0 { + childPath := path + if path != "/" { + childPath += "/" + } + childPath += parts[0] + + isDir := len(parts) > 1 + if !seen[childPath] { + if isDir { + dirInfo[childPath] = &FileInfo{ + Path: parts[0], + IsDir: true, + Size: 0, + ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano), + } + } else { + result = append(result, FileInfo{ + Path: parts[0], + IsDir: false, + Size: int64(len(entry.content)), + ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano), + }) + } + seen[childPath] = true + } else if isDir { + if info, ok := dirInfo[childPath]; ok { + if entry.modifiedAt.After(mustParseTime(info.ModifiedAt)) { + info.ModifiedAt = entry.modifiedAt.Format(time.RFC3339Nano) + } + } + } + } + } + } + + for _, info := range dirInfo { + result = append(result, *info) + } + + return result, nil +} + +func mustParseTime(s string) time.Time { + t, _ := time.Parse(time.RFC3339Nano, s) + return t +} + +// Read reads file content with offset and limit. +func (b *InMemoryBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + filePath := normalizePath(req.FilePath) + + entry, exists := b.files[filePath] + if !exists { + return nil, fmt.Errorf("file not found: %s", filePath) + } + + // Convert 1-based offset to 0-based index; values < 1 default to line 1 + offset := req.Offset - 1 + if offset < 0 { + offset = 0 + } + limit := req.Limit + + content := entry.content + + // Fast path: no offset and no limit — return as-is + if offset == 0 && limit <= 0 { + return &FileContent{Content: content}, nil + } + + // Fast path: no offset, content fits within limit — return as-is + if offset == 0 { + lineCount := strings.Count(content, "\n") + 1 + if lineCount <= limit { + return &FileContent{Content: content}, nil + } + } + + // Skip `offset` lines by scanning for newlines directly + start := 0 + for i := 0; i < offset; i++ { + idx := strings.IndexByte(content[start:], '\n') + if idx == -1 { + // offset exceeds total lines + return &FileContent{}, nil + } + start += idx + 1 + } + + // No limit: return everything from start + if limit <= 0 { + return &FileContent{Content: content[start:]}, nil + } + + // Find the end position after `limit` lines + end := start + for i := 0; i < limit; i++ { + idx := strings.IndexByte(content[end:], '\n') + if idx == -1 { + // Reached the end of content + return &FileContent{Content: content[start:]}, nil + } + end += idx + 1 + } + + // Trim the trailing newline from the last included line + return &FileContent{Content: content[start : end-1]}, nil +} + +// GrepRaw returns matches for the given pattern. +func (b *InMemoryBackend) GrepRaw(ctx context.Context, req *GrepRequest) ([]GrepMatch, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + if req.Pattern == "" { + return nil, fmt.Errorf("pattern cannot be empty") + } + + re, err := b.compilePattern(req) + if err != nil { + return nil, err + } + + searchPath := "/" + if req.Path != "" { + searchPath = normalizePath(req.Path) + } + + filteredFiles, err := b.filterFiles(searchPath, req) + if err != nil { + return nil, err + } + + if len(filteredFiles) == 0 { + return []GrepMatch{}, nil + } + + if len(filteredFiles) == 1 { + collector := newGrepCollector() + entry := b.files[filteredFiles[0]] + collector.processFile(filteredFiles[0], entry.content, re, req) + return collector.buildResults(b, req) + } + + matches, err := b.grepFilesInParallel(filteredFiles, re, req) + if err != nil { + return nil, err + } + + if req.BeforeLines > 0 || req.AfterLines > 0 { + matches = b.applyContext(matches, req) + } + + return matches, nil +} + +func (b *InMemoryBackend) grepFilesInParallel(filteredFiles []string, re *regexp.Regexp, req *GrepRequest) ([]GrepMatch, error) { + numWorkers := len(filteredFiles) + if numWorkers > 10 { + numWorkers = 10 + } + + type fileTask struct { + path string + content string + } + + tasks := make(chan fileTask, len(filteredFiles)) + results := make(chan []GrepMatch, len(filteredFiles)) + errChan := make(chan error, numWorkers) + + var wg sync.WaitGroup + for i := 0; i < numWorkers; i++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + defer func() { + if r := recover(); r != nil { + errChan <- fmt.Errorf("worker %d panic: %v", workerID, r) + } + }() + + collector := newGrepCollector() + for task := range tasks { + fileMatches := collector.findMatches(task.path, task.content, re, req) + if len(fileMatches) > 0 { + results <- fileMatches + } + } + }(i) + } + + for _, filePath := range filteredFiles { + entry := b.files[filePath] + tasks <- fileTask{ + path: filePath, + content: entry.content, + } + } + close(tasks) + + go func() { + wg.Wait() + close(results) + close(errChan) + }() + + var allMatches []GrepMatch + var errs []error + + for { + select { + case matches, ok := <-results: + if !ok { + results = nil + } else { + allMatches = append(allMatches, matches...) + } + case err, ok := <-errChan: + if !ok { + errChan = nil + } else if err != nil { + errs = append(errs, err) + } + } + + if results == nil && errChan == nil { + break + } + } + + if len(errs) > 0 { + return nil, fmt.Errorf("grep failed with %d error(s): %v", len(errs), errs[0]) + } + + return allMatches, nil +} + +func (b *InMemoryBackend) compilePattern(req *GrepRequest) (*regexp.Regexp, error) { + pattern := req.Pattern + if req.CaseInsensitive { + pattern = "(?i)" + pattern + } + re, err := regexp.Compile(pattern) + if err != nil { + return nil, fmt.Errorf("invalid regex pattern: %w", err) + } + return re, nil +} + +func (b *InMemoryBackend) filterFiles(searchPath string, req *GrepRequest) ([]string, error) { + var candidateFiles []string + + for filePath := range b.files { + normalizedFilePath := normalizePath(filePath) + + if searchPath != "/" && !strings.HasPrefix(normalizedFilePath, searchPath+"/") && normalizedFilePath != searchPath { + continue + } + + candidateFiles = append(candidateFiles, normalizedFilePath) + } + + if req.Glob != "" { + filtered, err := b.filterByGlob(candidateFiles, searchPath, req.Glob) + if err != nil { + return nil, err + } + candidateFiles = filtered + } + + if req.FileType != "" { + candidateFiles = b.filterByFileType(candidateFiles, req.FileType) + } + + return candidateFiles, nil +} + +func (b *InMemoryBackend) filterByGlob(files []string, searchPath string, globPattern string) ([]string, error) { + var result []string + + for _, filePath := range files { + var matchPath string + if strings.Contains(globPattern, "/") || strings.Contains(globPattern, "**") { + if searchPath == "/" { + matchPath = strings.TrimPrefix(filePath, "/") + } else { + matchPath = strings.TrimPrefix(filePath, searchPath+"/") + } + } else { + matchPath = filepath.Base(filePath) + } + + matched, err := doublestar.Match(globPattern, matchPath) + if err != nil { + return nil, fmt.Errorf("invalid glob pattern: %w", err) + } + if matched { + result = append(result, filePath) + } + } + + return result, nil +} + +func (b *InMemoryBackend) filterByFileType(files []string, fileType string) []string { + var result []string + + for _, filePath := range files { + ext := strings.TrimPrefix(filepath.Ext(filePath), ".") + if matchFileType(ext, fileType) { + result = append(result, filePath) + } + } + + return result +} + +// matchFileType checks if the file extension matches the given file type. +func matchFileType(ext, fileType string) bool { + typeMap := map[string][]string{ + "ada": {"adb", "ads"}, + "agda": {"agda", "lagda"}, + "aidl": {"aidl"}, + "amake": {"bp", "mk"}, + "asciidoc": {"adoc", "asc", "asciidoc"}, + "asm": {"S", "asm", "s"}, + "asp": {"ascx", "asp", "aspx"}, + "ats": {"ats", "dats", "hats", "sats"}, + "avro": {"avdl", "avpr", "avsc"}, + "awk": {"awk"}, + "bat": {"bat"}, + "bazel": {"BUILD", "bazel", "bzl"}, + "bitbake": {"bb", "bbappend", "bbclass", "conf", "inc"}, + "c": {"c", "h", "H", "cats"}, + "cabal": {"cabal"}, + "cbor": {"cbor"}, + "ceylon": {"ceylon"}, + "clojure": {"clj", "cljc", "cljs", "cljx"}, + "cmake": {"cmake"}, + "coffeescript": {"coffee"}, + "config": {"cfg", "conf", "config", "ini"}, + "coq": {"v"}, + "cpp": {"C", "cc", "cpp", "cxx", "c++", "h", "hh", "hpp", "hxx", "h++", "inl"}, + "crystal": {"cr", "ecr"}, + "cs": {"cs"}, + "csharp": {"cs"}, + "cshtml": {"cshtml"}, + "css": {"css", "scss", "sass", "less"}, + "csv": {"csv"}, + "cuda": {"cu", "cuh"}, + "cython": {"pxd", "pxi", "pyx"}, + "d": {"d"}, + "dart": {"dart"}, + "devicetree": {"dts", "dtsi"}, + "dhall": {"dhall"}, + "diff": {"diff", "patch"}, + "docker": {"dockerfile"}, + "go": {"go"}, + "groovy": {"gradle", "groovy"}, + "haskell": {"c2hs", "cpphs", "hs", "hsc", "lhs"}, + "html": {"ejs", "htm", "html"}, + "java": {"java", "jsp", "jspx", "properties"}, + "js": {"cjs", "js", "jsx", "mjs", "vue"}, + "json": {"json", "sarif"}, + "jsonl": {"jsonl"}, + "julia": {"jl"}, + "jupyter": {"ipynb", "jpynb"}, + "kotlin": {"kt", "kts"}, + "less": {"less"}, + "lua": {"lua"}, + "make": {"mak", "mk"}, + "markdown": {"markdown", "md", "mdown", "mdwn", "mdx", "mkd", "mkdn"}, + "md": {"markdown", "md", "mdown", "mdwn", "mdx", "mkd", "mkdn"}, + "matlab": {"m"}, + "ocaml": {"ml", "mli", "mll", "mly"}, + "perl": {"PL", "perl", "pl", "plh", "plx", "pm", "t"}, + "php": {"php", "php3", "php4", "php5", "php7", "php8", "pht", "phtml"}, + "python": {"py", "pyi"}, + "py": {"py", "pyi"}, + "ruby": {"gemspec", "rb", "rbw"}, + "rust": {"rs"}, + "sass": {"sass", "scss"}, + "scala": {"sbt", "scala"}, + "sh": {"bash", "sh", "zsh"}, + "sql": {"psql", "sql"}, + "swift": {"swift"}, + "toml": {"toml"}, + "ts": {"cts", "mts", "ts", "tsx"}, + "typescript": {"cts", "mts", "ts", "tsx"}, + "txt": {"txt"}, + "vue": {"vue"}, + "xml": {"dtd", "xml", "xsd", "xsl", "xslt"}, + "yaml": {"yaml", "yml"}, + "zig": {"zig"}, + } + + if exts, ok := typeMap[fileType]; ok { + for _, e := range exts { + if ext == e { + return true + } + } + } + return ext == fileType +} + +// applyContext adds context lines around matches. +func (b *InMemoryBackend) applyContext(matches []GrepMatch, req *GrepRequest) []GrepMatch { + if len(matches) == 0 { + return matches + } + + beforeLines := 0 + afterLines := 0 + + if req.BeforeLines > 0 { + beforeLines = req.BeforeLines + } + if req.AfterLines > 0 { + afterLines = req.AfterLines + } + + if beforeLines <= 0 && afterLines <= 0 { + return matches + } + + // Group matches by file path for efficient processing + matchesByFile := make(map[string][]GrepMatch) + fileOrder := make([]string, 0) + seenFiles := make(map[string]bool) + + for _, match := range matches { + if !seenFiles[match.Path] { + fileOrder = append(fileOrder, match.Path) + seenFiles[match.Path] = true + } + matchesByFile[match.Path] = append(matchesByFile[match.Path], match) + } + + var result []GrepMatch + + // Process each file once + for _, filePath := range fileOrder { + fileMatches := matchesByFile[filePath] + + // Get file content once per file + b.mu.RLock() + entry, exists := b.files[filePath] + b.mu.RUnlock() + + if !exists { + // If file doesn't exist, keep original matches + result = append(result, fileMatches...) + continue + } + + lines := strings.Split(entry.content, "\n") + processedLines := make(map[int]bool) + + // Process all matches for this file + for _, match := range fileMatches { + startLine := match.Line - beforeLines + if startLine < 1 { + startLine = 1 + } + + endLine := match.Line + afterLines + if endLine > len(lines) { + endLine = len(lines) + } + + for lineNum := startLine; lineNum <= endLine; lineNum++ { + if !processedLines[lineNum] { + processedLines[lineNum] = true + result = append(result, GrepMatch{ + Path: filePath, + Line: lineNum, + Content: lines[lineNum-1], + }) + } + } + } + } + + return result +} + +// GlobInfo returns file info entries matching the glob pattern. +func (b *InMemoryBackend) GlobInfo(ctx context.Context, req *GlobInfoRequest) ([]FileInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + basePath := normalizePath(req.Path) + isAbsolutePattern := strings.HasPrefix(req.Pattern, "/") + + var result []FileInfo + + for filePath, entry := range b.files { + normalizedFilePath := normalizePath(filePath) + + var matchPath string + var resultPath string + + if isAbsolutePattern { + matchPath = normalizedFilePath + resultPath = normalizedFilePath + } else { + if basePath != "/" && !strings.HasPrefix(normalizedFilePath, basePath+"/") && normalizedFilePath != basePath { + continue + } + + if basePath == "/" { + matchPath = strings.TrimPrefix(normalizedFilePath, "/") + } else { + matchPath = strings.TrimPrefix(normalizedFilePath, basePath+"/") + } + resultPath = matchPath + } + + matched, err := doublestar.Match(req.Pattern, matchPath) + if err != nil { + return nil, fmt.Errorf("invalid glob pattern: %w", err) + } + + if matched { + result = append(result, FileInfo{ + Path: resultPath, + IsDir: false, + Size: int64(len(entry.content)), + ModifiedAt: entry.modifiedAt.Format(time.RFC3339Nano), + }) + } + } + + return result, nil +} + +// Write creates or overwrites file content. +func (b *InMemoryBackend) Write(ctx context.Context, req *WriteRequest) error { + b.mu.Lock() + defer b.mu.Unlock() + + filePath := normalizePath(req.FilePath) + b.files[filePath] = &fileEntry{ + content: req.Content, + modifiedAt: time.Now(), + } + + return nil +} + +// Edit replaces string occurrences in a file. +func (b *InMemoryBackend) Edit(ctx context.Context, req *EditRequest) error { + b.mu.Lock() + defer b.mu.Unlock() + + filePath := normalizePath(req.FilePath) + + entry, exists := b.files[filePath] + if !exists { + return fmt.Errorf("file not found: %s", filePath) + } + + if req.OldString == "" { + return fmt.Errorf("oldString must be non-empty") + } + + content := entry.content + if !strings.Contains(content, req.OldString) { + return fmt.Errorf("oldString not found in file: %s", filePath) + } + + if !req.ReplaceAll { + firstIndex := strings.Index(content, req.OldString) + if firstIndex != -1 { + // Check if there's another occurrence after the first one + if strings.Contains(content[firstIndex+len(req.OldString):], req.OldString) { + return fmt.Errorf("multiple occurrences of oldString found in file %s, but ReplaceAll is false", filePath) + } + } + } + + var newContent string + if req.ReplaceAll { + newContent = strings.ReplaceAll(content, req.OldString, req.NewString) + } else { + newContent = strings.Replace(content, req.OldString, req.NewString, 1) + } + + b.files[filePath] = &fileEntry{ + content: newContent, + modifiedAt: time.Now(), + } + + return nil +} + +// normalizePath normalizes a file path by ensuring it starts with "/" and removing trailing slashes. +func normalizePath(path string) string { + if path == "" { + return "/" + } + + // Ensure path starts with "/" + if !strings.HasPrefix(path, "/") { + path = "/" + path + } + + return filepath.Clean(path) +} + +type grepCollector struct { + allMatches []GrepMatch +} + +func newGrepCollector() *grepCollector { + return &grepCollector{ + allMatches: []GrepMatch{}, + } +} + +func (c *grepCollector) processFile(filePath, content string, re *regexp.Regexp, req *GrepRequest) { + fileMatches := c.findMatches(filePath, content, re, req) + if len(fileMatches) > 0 { + c.allMatches = append(c.allMatches, fileMatches...) + } +} + +func (c *grepCollector) findMatches(filePath, content string, re *regexp.Regexp, req *GrepRequest) []GrepMatch { + if req.EnableMultiline { + return c.findMultilineMatches(filePath, content, re) + } + return c.findSingleLineMatches(filePath, content, re) +} + +func (c *grepCollector) findMultilineMatches(filePath, content string, re *regexp.Regexp) []GrepMatch { + var fileMatches []GrepMatch + matches := re.FindAllStringIndex(content, -1) + lines := strings.Split(content, "\n") + + for _, match := range matches { + matchStart := match[0] + matchEnd := match[1] + startLineNum := 1 + strings.Count(content[:matchStart], "\n") + endLineNum := 1 + strings.Count(content[:matchEnd], "\n") + + for lineNum := startLineNum; lineNum <= endLineNum && lineNum <= len(lines); lineNum++ { + fileMatches = append(fileMatches, GrepMatch{ + Path: filePath, + Line: lineNum, + Content: lines[lineNum-1], + }) + } + } + return fileMatches +} + +func (c *grepCollector) findSingleLineMatches(filePath, content string, re *regexp.Regexp) []GrepMatch { + var fileMatches []GrepMatch + lines := strings.Split(content, "\n") + for lineNum, line := range lines { + if re.MatchString(line) { + fileMatches = append(fileMatches, GrepMatch{ + Path: filePath, + Line: lineNum + 1, + Content: line, + }) + } + } + return fileMatches +} + +func (c *grepCollector) buildResults(b *InMemoryBackend, req *GrepRequest) ([]GrepMatch, error) { + return c.buildContentResult(b, req), nil +} + +func (c *grepCollector) buildContentResult(b *InMemoryBackend, req *GrepRequest) []GrepMatch { + results := c.allMatches + if req.BeforeLines > 0 || req.AfterLines > 0 { + results = b.applyContext(c.allMatches, req) + } + return results +} diff --git a/adk/filesystem/backend_inmemory_test.go b/adk/filesystem/backend_inmemory_test.go new file mode 100644 index 0000000..fae9bf0 --- /dev/null +++ b/adk/filesystem/backend_inmemory_test.go @@ -0,0 +1,2441 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "fmt" + "path/filepath" + "strconv" + "strings" + "testing" + "time" +) + +func TestInMemoryBackend_WriteAndRead(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + // Test Write + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "line1\nline2\nline3\nline4\nline5", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + // Test Read - full content + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Limit: 100, + }) + if err != nil { + t.Fatalf("Read failed: %v", err) + } + expected := "line1\nline2\nline3\nline4\nline5" + if content.Content != expected { + t.Errorf("Read content mismatch. Expected: %q, Got: %q", expected, content.Content) + } + + // Test Read - with offset and limit + content, err = backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Offset: 1, + Limit: 2, + }) + if err != nil { + t.Fatalf("Read with offset failed: %v", err) + } + expected = "line1\nline2" + if content.Content != expected { + t.Errorf("Read with offset content mismatch. Expected: %q, Got: %q", expected, content.Content) + } + + // Test Read - non-existent file + _, err = backend.Read(ctx, &ReadRequest{ + FilePath: "/nonexistent.txt", + Limit: 10, + }) + if err == nil { + t.Error("Expected error for non-existent file, got nil") + } +} + +func TestInMemoryBackend_LsInfo(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + // Create some files + backend.Write(ctx, &WriteRequest{ + FilePath: "/file1.txt", + Content: "content1", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/file2.txt", + Content: "content2", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file3.txt", + Content: "content3", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/subdir/file4.txt", + Content: "content4", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir2/file5.txt", + Content: "content5", + }) + + // Test LsInfo - root + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos) != 4 { // file1.txt, file2.txt, dir1, dir2 + t.Errorf("Expected 4 items in root, got %d", len(infos)) + } + + // Test LsInfo - specific directory + infos, err = backend.LsInfo(ctx, &LsInfoRequest{Path: "/dir1"}) + if err != nil { + t.Fatalf("LsInfo for /dir1 failed: %v", err) + } + if len(infos) != 2 { // file3.txt, subdir + t.Errorf("Expected 2 items in /dir1, got %d", len(infos)) + } +} + +func TestInMemoryBackend_Edit(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + // Create a file + backend.Write(ctx, &WriteRequest{ + FilePath: "/edit.txt", + Content: "hello world\nhello again\nhello world", + }) + + // Test Edit - report error if old string occurs + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/edit.txt", + OldString: "hello", + NewString: "hi", + ReplaceAll: false, + }) + if err == nil { + t.Fatal("should have failed") + } + + // Test Edit - replace all occurrences + backend.Write(ctx, &WriteRequest{ + FilePath: "/edit2.txt", + Content: "hello world\nhello again\nhello world", + }) + err = backend.Edit(ctx, &EditRequest{ + FilePath: "/edit2.txt", + OldString: "hello", + NewString: "hi", + ReplaceAll: true, + }) + if err != nil { + t.Fatalf("Edit (replace all) failed: %v", err) + } + + content, _ := backend.Read(ctx, &ReadRequest{ + FilePath: "/edit2.txt", + Limit: 100, + }) + expected := "hi world\nhi again\nhi world" + if content.Content != expected { + t.Errorf("Edit (replace all) content mismatch. Expected: %q, Got: %q", expected, content.Content) + } + + // Test Edit - non-existent file + err = backend.Edit(ctx, &EditRequest{ + FilePath: "/nonexistent.txt", + OldString: "old", + NewString: "new", + ReplaceAll: false, + }) + if err == nil { + t.Error("Expected error for non-existent file, got nil") + } + + // Test Edit - empty oldString + err = backend.Edit(ctx, &EditRequest{ + FilePath: "/edit.txt", + OldString: "", + NewString: "new", + ReplaceAll: false, + }) + if err == nil { + t.Error("Expected error for empty oldString, got nil") + } +} + +func TestInMemoryBackend_LsInfo_PathIsFilename(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/file1.txt", + Content: "content1", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/file2.txt", + Content: "content2", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file3.txt", + Content: "content3", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/subdir/file4.txt", + Content: "content4", + }) + + t.Run("RootDirectory", func(t *testing.T) { + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + for _, info := range infos { + if strings.Contains(info.Path, "/") { + t.Errorf("Path should be filename only, got: %s", info.Path) + } + if info.IsDir { + if info.Path != "dir1" { + t.Errorf("Expected directory name 'dir1', got: %s", info.Path) + } + } else { + if info.Path != "file1.txt" && info.Path != "file2.txt" { + t.Errorf("Expected filename 'file1.txt' or 'file2.txt', got: %s", info.Path) + } + } + } + }) + + t.Run("Subdirectory", func(t *testing.T) { + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/dir1"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + for _, info := range infos { + if strings.Contains(info.Path, "/") { + t.Errorf("Path should be filename only, got: %s", info.Path) + } + if info.IsDir { + if info.Path != "subdir" { + t.Errorf("Expected directory name 'subdir', got: %s", info.Path) + } + } else { + if info.Path != "file3.txt" { + t.Errorf("Expected filename 'file3.txt', got: %s", info.Path) + } + } + } + }) + + t.Run("NestedSubdirectory", func(t *testing.T) { + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/dir1/subdir"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + if len(infos) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos)) + } + + info := infos[0] + if info.Path != "file4.txt" { + t.Errorf("Expected filename 'file4.txt', got: %s", info.Path) + } + if strings.Contains(info.Path, "/") { + t.Errorf("Path should be filename only, got: %s", info.Path) + } + }) +} + +func TestInMemoryBackend_GlobInfo(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + // Create some files + backend.Write(ctx, &WriteRequest{ + FilePath: "/file1.txt", + Content: "content1", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/file2.py", + Content: "content2", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file3.txt", + Content: "content3", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file4.py", + Content: "content4", + }) + + // Test GlobInfo - match .txt files in root only + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "*.txt", + Path: "/", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + if len(infos) != 1 { // only file1.txt in root + t.Errorf("Expected 1 .txt file in root, got %d", len(infos)) + } + if infos[0].Path != "file1.txt" { + t.Errorf("Expected relative path 'file1.txt', got %s", infos[0].Path) + } + + // Test GlobInfo - match all .py files in dir1 + infos, err = backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "*.py", + Path: "/dir1", + }) + if err != nil { + t.Fatalf("GlobInfo for /dir1 failed: %v", err) + } + if len(infos) != 1 { // file4.py + t.Errorf("Expected 1 .py file in /dir1, got %d", len(infos)) + } + if infos[0].Path != "file4.py" { + t.Errorf("Expected relative path 'file4.py', got %s", infos[0].Path) + } +} + +func TestInMemoryBackend_GlobInfo_RelativePath(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/Users/bytedance/Desktop/github/eino/file1.go", + Content: "content1", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/Users/bytedance/Desktop/github/openai-go/paginationmanual_test.go", + Content: "content2", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/Users/bytedance/Desktop/github/openai-go/paginationauto_test.go", + Content: "content3", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/Users/bytedance/Desktop/other/test.go", + Content: "content4", + }) + + t.Run("GlobFromRootWithPattern", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "**/*.go", + Path: "/Users/bytedance/Desktop/github", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 3 { + t.Fatalf("Expected 3 .go files, got %d", len(infos)) + } + + expectedPaths := map[string]bool{ + "eino/file1.go": false, + "openai-go/paginationmanual_test.go": false, + "openai-go/paginationauto_test.go": false, + } + + for _, info := range infos { + if _, exists := expectedPaths[info.Path]; exists { + expectedPaths[info.Path] = true + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + + for path, found := range expectedPaths { + if !found { + t.Errorf("Expected path not found: %s", path) + } + } + }) + + t.Run("GlobFromSubdirectory", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "*.go", + Path: "/Users/bytedance/Desktop/github/openai-go", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 2 { + t.Fatalf("Expected 2 .go files, got %d", len(infos)) + } + + expectedPaths := map[string]bool{ + "paginationmanual_test.go": false, + "paginationauto_test.go": false, + } + + for _, info := range infos { + if _, exists := expectedPaths[info.Path]; exists { + expectedPaths[info.Path] = true + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + + for path, found := range expectedPaths { + if !found { + t.Errorf("Expected path not found: %s", path) + } + } + }) + + t.Run("GlobFromRootWithAbsolutePattern", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "/Users/bytedance/Desktop/github/**/*.go", + Path: "/", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + expected := map[string]bool{ + "/Users/bytedance/Desktop/github/eino/file1.go": false, + "/Users/bytedance/Desktop/github/openai-go/paginationmanual_test.go": false, + "/Users/bytedance/Desktop/github/openai-go/paginationauto_test.go": false, + } + for _, info := range infos { + if _, ok := expected[info.Path]; ok { + expected[info.Path] = true + } + } + for path, found := range expected { + if !found { + t.Errorf("Expected absolute path not found: %s", path) + } + } + }) + + t.Run("GlobRecursiveWithRelativePattern", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "**/*.go", + Path: "/Users/bytedance/Desktop/github", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 3 { + t.Fatalf("Expected 3 .go files with ** pattern, got %d", len(infos)) + } + + expected := map[string]bool{ + "eino/file1.go": false, + "openai-go/paginationmanual_test.go": false, + "openai-go/paginationauto_test.go": false, + } + + for _, info := range infos { + if _, ok := expected[info.Path]; ok { + expected[info.Path] = true + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + + for path, found := range expected { + if !found { + t.Errorf("Expected relative path not found: %s", path) + } + } + }) +} + +func TestInMemoryBackend_GlobInfo_RecursivePattern(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/src/main.go", + Content: "main", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/src/utils/helper.go", + Content: "helper", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/src/utils/deep/nested.go", + Content: "nested", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/test/test.go", + Content: "test", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/README.md", + Content: "readme", + }) + + t.Run("DoubleStarMatchesAllSubdirectories", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "**/*.go", + Path: "/project", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 4 { + t.Fatalf("Expected 4 .go files, got %d", len(infos)) + } + + expected := map[string]bool{ + "src/main.go": false, + "src/utils/helper.go": false, + "src/utils/deep/nested.go": false, + "test/test.go": false, + } + + for _, info := range infos { + if _, ok := expected[info.Path]; ok { + expected[info.Path] = true + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + + for path, found := range expected { + if !found { + t.Errorf("Expected path not found: %s", path) + } + } + }) + + t.Run("DoubleStarInMiddleOfPattern", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "src/**/*.go", + Path: "/project", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 3 { + t.Fatalf("Expected 3 .go files under src/, got %d", len(infos)) + } + + expected := map[string]bool{ + "src/main.go": false, + "src/utils/helper.go": false, + "src/utils/deep/nested.go": false, + } + + for _, info := range infos { + if _, ok := expected[info.Path]; ok { + expected[info.Path] = true + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + + for path, found := range expected { + if !found { + t.Errorf("Expected path not found: %s", path) + } + } + }) + + t.Run("DoubleStarAtEnd", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "src/**", + Path: "/project", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 3 { + t.Fatalf("Expected 3 files under src/, got %d", len(infos)) + } + + expected := map[string]bool{ + "src/main.go": false, + "src/utils/helper.go": false, + "src/utils/deep/nested.go": false, + } + + for _, info := range infos { + if _, ok := expected[info.Path]; ok { + expected[info.Path] = true + } + } + + for path, found := range expected { + if !found { + t.Errorf("Expected path not found: %s", path) + } + } + }) + + t.Run("AbsolutePatternWithDoubleStarRecursive", func(t *testing.T) { + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "/project/**/*.go", + Path: "/", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 4 { + t.Fatalf("Expected 4 .go files, got %d", len(infos)) + } + + expected := map[string]bool{ + "/project/src/main.go": false, + "/project/src/utils/helper.go": false, + "/project/src/utils/deep/nested.go": false, + "/project/test/test.go": false, + } + + for _, info := range infos { + if _, ok := expected[info.Path]; ok { + expected[info.Path] = true + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + + for path, found := range expected { + if !found { + t.Errorf("Expected absolute path not found: %s", path) + } + } + }) +} + +func TestInMemoryBackend_Concurrent(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + // Test concurrent writes and reads + done := make(chan bool) + for i := 0; i < 10; i++ { + go func(n int) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/concurrent.txt", + Content: "content", + }) + backend.Read(ctx, &ReadRequest{ + FilePath: "/concurrent.txt", + Limit: 10, + }) + done <- true + }(i) + } + + for i := 0; i < 10; i++ { + <-done + } +} + +func TestInMemoryBackend_LsInfo_FileInfoMetadata(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + t.Run("FileMetadata", func(t *testing.T) { + content := "hello world" + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: content, + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + if len(infos) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos)) + } + + info := infos[0] + if info.Path != "test.txt" { + t.Errorf("Expected path test.txt, got %s", info.Path) + } + if info.IsDir { + t.Error("Expected IsDir to be false for file") + } + if info.Size != int64(len(content)) { + t.Errorf("Expected size %d, got %d", len(content), info.Size) + } + if info.ModifiedAt == "" { + t.Error("Expected ModifiedAt to be non-empty") + } + _, err = time.Parse(time.RFC3339Nano, info.ModifiedAt) + if err != nil { + t.Errorf("ModifiedAt is not valid RFC3339 format: %v", err) + } + }) + + t.Run("DirectoryMetadata", func(t *testing.T) { + backend := NewInMemoryBackend() + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file1.txt", + Content: "content1", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + if len(infos) != 1 { + t.Fatalf("Expected 1 directory, got %d", len(infos)) + } + + info := infos[0] + if info.Path != "dir1" { + t.Errorf("Expected path dir1, got %s", info.Path) + } + if !info.IsDir { + t.Error("Expected IsDir to be true for directory") + } + if info.Size != 0 { + t.Errorf("Expected size 0 for directory, got %d", info.Size) + } + if info.ModifiedAt == "" { + t.Error("Expected ModifiedAt to be non-empty for directory") + } + }) + + t.Run("MixedFilesAndDirectories", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{ + FilePath: "/file1.txt", + Content: "content1", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file2.txt", + Content: "content2", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/subdir/file3.txt", + Content: "content3", + }) + + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + if len(infos) != 2 { + t.Fatalf("Expected 2 items (file1.txt, dir1), got %d", len(infos)) + } + + fileCount := 0 + dirCount := 0 + for _, info := range infos { + if info.IsDir { + dirCount++ + if info.Path != "dir1" { + t.Errorf("Expected directory path dir1, got %s", info.Path) + } + } else { + fileCount++ + if info.Path != "file1.txt" { + t.Errorf("Expected file path file1.txt, got %s", info.Path) + } + if info.Size != int64(len("content1")) { + t.Errorf("Expected file size %d, got %d", len("content1"), info.Size) + } + } + } + + if fileCount != 1 { + t.Errorf("Expected 1 file, got %d", fileCount) + } + if dirCount != 1 { + t.Errorf("Expected 1 directory, got %d", dirCount) + } + }) + + t.Run("SubdirectoryListing", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file1.txt", + Content: "short", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/subdir/file2.txt", + Content: "longer content here", + }) + + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/dir1"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + if len(infos) != 2 { + t.Fatalf("Expected 2 items (file1.txt, subdir), got %d", len(infos)) + } + + for _, info := range infos { + if info.Path == "file1.txt" { + if info.IsDir { + t.Error("Expected file1.txt to be a file") + } + if info.Size != int64(len("short")) { + t.Errorf("Expected size %d, got %d", len("short"), info.Size) + } + } else if info.Path == "subdir" { + if !info.IsDir { + t.Error("Expected subdir to be a directory") + } + } else { + t.Errorf("Unexpected path: %s", info.Path) + } + } + }) + + t.Run("DirectoryModifiedAtUsesLatestFile", func(t *testing.T) { + backend := NewInMemoryBackend() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file1.txt", + Content: "content1", + }) + time.Sleep(10 * time.Millisecond) + + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir1/file2.txt", + Content: "content2", + }) + + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + + if len(infos) != 1 { + t.Fatalf("Expected 1 directory, got %d", len(infos)) + } + + dirInfo := infos[0] + if !dirInfo.IsDir { + t.Fatal("Expected directory") + } + + dirModTime, _ := time.Parse(time.RFC3339Nano, dirInfo.ModifiedAt) + + subInfos, _ := backend.LsInfo(ctx, &LsInfoRequest{Path: "/dir1"}) + var latestFileTime time.Time + for _, info := range subInfos { + fileTime, _ := time.Parse(time.RFC3339Nano, info.ModifiedAt) + if fileTime.After(latestFileTime) { + latestFileTime = fileTime + } + } + + if !dirModTime.Equal(latestFileTime) && dirModTime.Before(latestFileTime) { + t.Logf("Directory mod time: %v, Latest file time: %v", dirModTime, latestFileTime) + } + }) +} + +func TestInMemoryBackend_GlobInfo_FileInfoMetadata(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + t.Run("BasicMetadata", func(t *testing.T) { + content := "test content" + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: content, + }) + + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "*.txt", + Path: "/", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos)) + } + + info := infos[0] + if info.Path != "test.txt" { + t.Errorf("Expected path test.txt, got %s", info.Path) + } + if info.IsDir { + t.Error("Expected IsDir to be false") + } + if info.Size != int64(len(content)) { + t.Errorf("Expected size %d, got %d", len(content), info.Size) + } + if info.ModifiedAt == "" { + t.Error("Expected ModifiedAt to be non-empty") + } + }) + + t.Run("MultipleFilesMetadata", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{ + FilePath: "/file1.txt", + Content: "short", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/file2.txt", + Content: "much longer content", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/file3.py", + Content: "python", + }) + + infos, err := backend.GlobInfo(ctx, &GlobInfoRequest{ + Pattern: "*.txt", + Path: "/", + }) + if err != nil { + t.Fatalf("GlobInfo failed: %v", err) + } + + if len(infos) != 2 { + t.Fatalf("Expected 2 .txt files, got %d", len(infos)) + } + + for _, info := range infos { + if info.IsDir { + t.Errorf("Expected IsDir to be false for %s", info.Path) + } + if info.Size <= 0 { + t.Errorf("Expected positive size for %s, got %d", info.Path, info.Size) + } + if info.ModifiedAt == "" { + t.Errorf("Expected ModifiedAt to be non-empty for %s", info.Path) + } + } + }) +} + +func TestInMemoryBackend_WriteAndEdit_ModifiedAt(t *testing.T) { + ctx := context.Background() + + t.Run("WriteUpdatesModifiedAt", func(t *testing.T) { + backend := NewInMemoryBackend() + beforeWrite := time.Now() + time.Sleep(1 * time.Millisecond) + + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "initial content", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + time.Sleep(1 * time.Millisecond) + afterWrite := time.Now() + + infos, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos)) + } + + modTime, err := time.Parse(time.RFC3339Nano, infos[0].ModifiedAt) + if err != nil { + t.Fatalf("Failed to parse ModifiedAt: %v", err) + } + + if modTime.Before(beforeWrite) || modTime.After(afterWrite) { + t.Errorf("ModifiedAt %v should be between %v and %v", modTime, beforeWrite, afterWrite) + } + }) + + t.Run("EditUpdatesModifiedAt", func(t *testing.T) { + backend := NewInMemoryBackend() + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/edit.txt", + Content: "hello world", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + infos1, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos1) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos1)) + } + modTime1, err := time.Parse(time.RFC3339Nano, infos1[0].ModifiedAt) + if err != nil { + t.Fatalf("Failed to parse ModifiedAt: %v", err) + } + + time.Sleep(10 * time.Millisecond) + + err = backend.Edit(ctx, &EditRequest{ + FilePath: "/edit.txt", + OldString: "hello", + NewString: "hi", + ReplaceAll: true, + }) + if err != nil { + t.Fatalf("Edit failed: %v", err) + } + + infos2, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos2) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos2)) + } + modTime2, err := time.Parse(time.RFC3339Nano, infos2[0].ModifiedAt) + if err != nil { + t.Fatalf("Failed to parse ModifiedAt: %v", err) + } + + if !modTime2.After(modTime1) { + t.Errorf("ModifiedAt should be updated after edit. Before: %v, After: %v", modTime1, modTime2) + } + }) + + t.Run("OverwriteUpdatesModifiedAt", func(t *testing.T) { + backend := NewInMemoryBackend() + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/overwrite.txt", + Content: "original", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + infos1, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos1) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos1)) + } + modTime1, err := time.Parse(time.RFC3339Nano, infos1[0].ModifiedAt) + if err != nil { + t.Fatalf("Failed to parse ModifiedAt: %v", err) + } + + time.Sleep(10 * time.Millisecond) + + err = backend.Write(ctx, &WriteRequest{ + FilePath: "/overwrite.txt", + Content: "new content", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + infos2, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos2) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos2)) + } + modTime2, err := time.Parse(time.RFC3339Nano, infos2[0].ModifiedAt) + if err != nil { + t.Fatalf("Failed to parse ModifiedAt: %v", err) + } + + if !modTime2.After(modTime1) { + t.Errorf("ModifiedAt should be updated after overwrite. Before: %v, After: %v", modTime1, modTime2) + } + }) + + t.Run("SizeUpdatesAfterEdit", func(t *testing.T) { + backend := NewInMemoryBackend() + err := backend.Write(ctx, &WriteRequest{ + FilePath: "/size.txt", + Content: "hello", + }) + if err != nil { + t.Fatalf("Write failed: %v", err) + } + + infos1, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos1) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos1)) + } + size1 := infos1[0].Size + + err = backend.Edit(ctx, &EditRequest{ + FilePath: "/size.txt", + OldString: "hello", + NewString: "hello world", + ReplaceAll: true, + }) + if err != nil { + t.Fatalf("Edit failed: %v", err) + } + + infos2, err := backend.LsInfo(ctx, &LsInfoRequest{Path: "/"}) + if err != nil { + t.Fatalf("LsInfo failed: %v", err) + } + if len(infos2) != 1 { + t.Fatalf("Expected 1 file, got %d", len(infos2)) + } + size2 := infos2[0].Size + + if size2 <= size1 { + t.Errorf("Size should increase after edit. Before: %d, After: %d", size1, size2) + } + if size2 != int64(len("hello world")) { + t.Errorf("Expected size %d, got %d", len("hello world"), size2) + } + }) +} + +func TestInMemoryBackend_Read_EdgeCases(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "line1\nline2\nline3", + }) + + t.Run("negative offset should be treated as zero", func(t *testing.T) { + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Offset: -5, + Limit: 2, + }) + if err != nil { + t.Fatalf("Read failed: %v", err) + } + expected := "line1\nline2" + if content.Content != expected { + t.Errorf("Expected: %q, Got: %q", expected, content.Content) + } + }) + + t.Run("offset exceeds file length", func(t *testing.T) { + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Offset: 100, + Limit: 10, + }) + if err != nil { + t.Fatalf("Read failed: %v", err) + } + if content.Content != "" { + t.Errorf("Expected empty content, got: %q", content.Content) + } + }) + + t.Run("zero or negative limit should use default 200", func(t *testing.T) { + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Offset: 0, + Limit: 0, + }) + if err != nil { + t.Fatalf("Read failed: %v", err) + } + lines := strings.Split(content.Content, "\n") + if len(lines) != 3 { + t.Errorf("Expected 3 lines, got %d", len(lines)) + } + }) + + t.Run("limit exceeds remaining lines", func(t *testing.T) { + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Offset: 1, + Limit: 100, + }) + if err != nil { + t.Fatalf("Read failed: %v", err) + } + lines := strings.Split(content.Content, "\n") + if len(lines) != 3 { + t.Errorf("Expected 3 lines, got %d", len(lines)) + } + }) +} + +func TestInMemoryBackend_Edit_EdgeCases(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + t.Run("edit non-existent file", func(t *testing.T) { + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/nonexistent.txt", + OldString: "old", + NewString: "new", + }) + if err == nil { + t.Error("Expected error for non-existent file") + } + if !strings.Contains(err.Error(), "not found") { + t.Errorf("Expected 'not found' error, got: %v", err) + } + }) + + t.Run("empty oldString", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "content", + }) + + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/test.txt", + OldString: "", + NewString: "new", + }) + if err == nil { + t.Error("Expected error for empty oldString") + } + if !strings.Contains(err.Error(), "non-empty") { + t.Errorf("Expected 'non-empty' error, got: %v", err) + } + }) + + t.Run("oldString not found", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "hello world", + }) + + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/test.txt", + OldString: "notfound", + NewString: "new", + }) + if err == nil { + t.Error("Expected error when oldString not found") + } + if !strings.Contains(err.Error(), "not found in file") { + t.Errorf("Expected 'not found in file' error, got: %v", err) + } + }) + + t.Run("multiple occurrences with ReplaceAll false", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "foo bar foo baz", + }) + + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/test.txt", + OldString: "foo", + NewString: "FOO", + ReplaceAll: false, + }) + if err == nil { + t.Error("Expected error for multiple occurrences with ReplaceAll=false") + } + if !strings.Contains(err.Error(), "multiple occurrences") { + t.Errorf("Expected 'multiple occurrences' error, got: %v", err) + } + }) + + t.Run("single occurrence with ReplaceAll false", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "foo bar baz", + }) + + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/test.txt", + OldString: "foo", + NewString: "FOO", + ReplaceAll: false, + }) + if err != nil { + t.Fatalf("Edit failed: %v", err) + } + + content, _ := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Limit: 100, + }) + if !strings.Contains(content.Content, "FOO") { + t.Error("Expected content to contain 'FOO'") + } + }) + + t.Run("ReplaceAll replaces all occurrences", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.txt", + Content: "foo bar foo baz foo", + }) + + err := backend.Edit(ctx, &EditRequest{ + FilePath: "/test.txt", + OldString: "foo", + NewString: "FOO", + ReplaceAll: true, + }) + if err != nil { + t.Fatalf("Edit failed: %v", err) + } + + content, _ := backend.Read(ctx, &ReadRequest{ + FilePath: "/test.txt", + Limit: 100, + }) + if strings.Contains(content.Content, "foo") { + t.Error("Expected all 'foo' to be replaced") + } + fooCount := strings.Count(content.Content, "FOO") + if fooCount != 3 { + t.Errorf("Expected 3 occurrences of 'FOO', got %d", fooCount) + } + }) +} + +func TestInMemoryBackend_NormalizePath(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + t.Run("paths are normalized on write", func(t *testing.T) { + testCases := []struct { + inputPath string + normalizedPath string + }{ + {"test.txt", "/test.txt"}, + {"/test.txt", "/test.txt"}, + {"//test.txt", "/test.txt"}, + {"/dir//file.txt", "/dir/file.txt"}, + {"/dir/../file.txt", "/file.txt"}, + } + + for _, tc := range testCases { + backend.Write(ctx, &WriteRequest{ + FilePath: tc.inputPath, + Content: "content", + }) + + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: tc.normalizedPath, + Limit: 10, + }) + if err != nil { + t.Errorf("Failed to read normalized path %s (from %s): %v", tc.normalizedPath, tc.inputPath, err) + } + if !strings.Contains(content.Content, "content") { + t.Errorf("Content not found for normalized path %s (from %s)", tc.normalizedPath, tc.inputPath) + } + } + }) +} + +func TestInMemoryBackend_MatchFileType(t *testing.T) { + testCases := []struct { + ext string + fileType string + expected bool + }{ + {"go", "go", true}, + {"py", "python", true}, + {"py", "py", true}, + {"js", "js", true}, + {"ts", "typescript", true}, + {"ts", "ts", true}, + {"cpp", "cpp", true}, + {"c", "c", true}, + {"h", "c", true}, + {"md", "markdown", true}, + {"txt", "txt", true}, + {"go", "python", false}, + {"js", "typescript", false}, + {"unknown", "go", false}, + } + + for _, tc := range testCases { + t.Run(fmt.Sprintf("%s matches %s", tc.ext, tc.fileType), func(t *testing.T) { + result := matchFileType(tc.ext, tc.fileType) + if result != tc.expected { + t.Errorf("matchFileType(%q, %q) = %v, expected %v", tc.ext, tc.fileType, result, tc.expected) + } + }) + } +} + +func TestInMemoryBackend_GrepRaw(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.go", + Content: "package main\nfunc main() {\n\tlog.Error(\"error\")\n\tfmt.Println(\"hello\")\n}", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/test.py", + Content: "def hello():\n print('error')\n print('world')", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/dir/file.go", + Content: "package test\nfunc TestError() {\n\tlog.Error(\"test error\")\n}", + }) + + t.Run("basic pattern search", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "error", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 3 { + t.Errorf("Expected 2 matches, got %d", len(matches)) + } + }) + + t.Run("empty pattern error", func(t *testing.T) { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "", + }) + if err == nil { + t.Error("Expected error for empty pattern") + } + if !strings.Contains(err.Error(), "cannot be empty") { + t.Errorf("Expected 'cannot be empty' error, got: %v", err) + } + }) + + t.Run("invalid regex pattern", func(t *testing.T) { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "[invalid", + }) + if err == nil { + t.Error("Expected error for invalid regex") + } + if !strings.Contains(err.Error(), "invalid regex") { + t.Errorf("Expected 'invalid regex' error, got: %v", err) + } + }) + + t.Run("case sensitive search", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "Error", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 3 { + t.Errorf("Expected 2 matches, got %d", len(matches)) + } + }) + + t.Run("case insensitive search", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "ERROR", + CaseInsensitive: true, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 3 { + t.Errorf("Expected at least 2 matches, got %d", len(matches)) + } + }) + + t.Run("filter by file type", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "error", + FileType: "go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasSuffix(match.Path, ".go") { + t.Errorf("Expected only .go files, got: %s", match.Path) + } + } + }) + + t.Run("filter by glob pattern", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "Error", + Glob: "*.go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasSuffix(match.Path, ".go") { + t.Errorf("Expected only .go files, got: %s", match.Path) + } + } + }) + + t.Run("invalid glob pattern", func(t *testing.T) { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "error", + Glob: "[invalid", + }) + if err == nil { + t.Error("Expected error for invalid glob pattern") + } + if !strings.Contains(err.Error(), "invalid glob") { + t.Errorf("Expected 'invalid glob' error, got: %v", err) + } + }) + + t.Run("search in specific path", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "Error", + Path: "/dir", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasPrefix(match.Path, "/dir") { + t.Errorf("Expected matches only from /dir, got: %s", match.Path) + } + } + }) + + t.Run("search with non-existent path", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "error", + Path: "/nonexistent", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 0 { + t.Errorf("Expected 0 matches for non-existent path, got %d", len(matches)) + } + }) + + t.Run("regex pattern matching", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "log\\..*Error", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 1 { + t.Errorf("Expected at least 1 match, got %d", len(matches)) + } + }) + + t.Run("no matches found", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "nonexistent_pattern_xyz", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 0 { + t.Errorf("Expected 0 matches, got %d", len(matches)) + } + }) + + t.Run("match line numbers", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "log\\.Error", + FileType: "go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if match.Line <= 0 { + t.Errorf("Expected positive line number, got %d", match.Line) + } + } + }) + + t.Run("match content is returned", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "package main", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 1 { + t.Fatal("Expected at least 1 match") + } + found := false + for _, match := range matches { + if strings.Contains(match.Content, "package main") { + found = true + break + } + } + if !found { + t.Error("Expected match content to contain 'package main'") + } + }) +} + +func TestInMemoryBackend_GrepRaw_WithContext(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/context.txt", + Content: "line1\nline2\ntarget line\nline4\nline5\nline6", + }) + + t.Run("with before context", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "target", + BeforeLines: 2, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 3 { + t.Errorf("Expected at least 3 matches (2 before + target), got %d", len(matches)) + } + }) + + t.Run("with after context", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "target", + AfterLines: 2, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 3 { + t.Errorf("Expected at least 3 matches (target + 2 after), got %d", len(matches)) + } + }) + + t.Run("with both before and after context", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "target", + BeforeLines: 1, + AfterLines: 1, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 3 { + t.Errorf("Expected at least 3 matches (1 before + target + 1 after), got %d", len(matches)) + } + }) + + t.Run("context at file boundaries", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/boundary.txt", + Content: "first line target\nsecond line", + }) + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "target", + Path: "/boundary.txt", + BeforeLines: 5, + AfterLines: 5, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) == 0 { + t.Error("Expected at least 1 match") + } + }) + + t.Run("zero context lines", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "target", + BeforeLines: 0, + AfterLines: 0, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 1 { + t.Error("Expected at least 1 match") + } + }) + + t.Run("negative context lines treated as zero", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "target", + BeforeLines: -5, + AfterLines: -5, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) < 1 { + t.Error("Expected at least 1 match") + } + }) +} + +func TestInMemoryBackend_GrepRaw_Multiline(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/multiline.txt", + Content: "start\nmiddle line\nend", + }) + + t.Run("single line mode (default)", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "start.*end", + EnableMultiline: false, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 0 { + t.Errorf("Expected 0 matches in single-line mode, got %d", len(matches)) + } + }) + + t.Run("multiline mode enabled", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "start[\\s\\S]*end", + EnableMultiline: true, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) == 0 { + t.Error("Expected matches in multiline mode") + } + }) + + t.Run("multiline with multiple matches", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/multiline2.txt", + Content: "block1 start\nblock1 middle\nblock1 end\n\nblock2 start\nblock2 end", + }) + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "start[\\s\\S]*?end", + Path: "/multiline2.txt", + EnableMultiline: true, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) == 0 { + t.Error("Expected matches in multiline mode") + } + }) + + t.Run("multiline with multiple matches v2", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/multiline3.txt", + Content: ` +const a = 1; +function calculateTotal( + items, + discount +) { + return items.reduce((sum, item) => sum + item.price, 0); +} + +const b = 2; + +/* + * This is a comment + * spanning multiple lines + */ + +class UserService { + constructor(db) { + this.db = db; + } +} +`, + }) + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "function calculateTotal\\([^\\)]*\\)", + Path: "/multiline3.txt", + EnableMultiline: true, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) == 0 { + t.Error("Expected matches in multiline mode") + } + + foundLastLine := false + for _, match := range matches { + if match.Line == 6 && strings.Contains(match.Content, ") {") { + foundLastLine = true + break + } + } + if !foundLastLine { + t.Error("Expected to find line 5 with ') {' in content") + for _, match := range matches { + t.Logf("Line %d: %s", match.Line, match.Content) + } + } + }) + +} + +func TestInMemoryBackend_GrepRaw_EmptyFiles(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + t.Run("search in empty file", func(t *testing.T) { + backend.Write(ctx, &WriteRequest{ + FilePath: "/empty.txt", + Content: "", + }) + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "anything", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 0 { + t.Errorf("Expected 0 matches in empty file, got %d", len(matches)) + } + }) + + t.Run("search with no files", func(t *testing.T) { + emptyBackend := NewInMemoryBackend() + matches, err := emptyBackend.GrepRaw(ctx, &GrepRequest{ + Pattern: "anything", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 0 { + t.Errorf("Expected 0 matches with no files, got %d", len(matches)) + } + }) +} + +func TestInMemoryBackend_GrepRaw_SpecialCharacters(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/special.txt", + Content: "interface{}\nmap[string]int\nfunc() error\n$variable\n*pointer", + }) + + t.Run("match curly braces", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "interface\\{\\}", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 1 { + t.Errorf("Expected 1 match, got %d", len(matches)) + } + }) + + t.Run("match square brackets", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "map\\[.*\\]", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 1 { + t.Errorf("Expected 1 match, got %d", len(matches)) + } + }) + + t.Run("match parentheses", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "func\\(\\)", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 1 { + t.Errorf("Expected 1 match, got %d", len(matches)) + } + }) + + t.Run("match dollar sign", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "\\$variable", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 1 { + t.Errorf("Expected 1 match, got %d", len(matches)) + } + }) + + t.Run("match asterisk", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "\\*pointer", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 1 { + t.Errorf("Expected 1 match, got %d", len(matches)) + } + }) +} + +func TestInMemoryBackend_GrepRaw_Concurrent(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + for i := 0; i < 10; i++ { + backend.Write(ctx, &WriteRequest{ + FilePath: fmt.Sprintf("/file%d.txt", i), + Content: fmt.Sprintf("content%d with error message", i), + }) + } + + t.Run("concurrent grep operations", func(t *testing.T) { + done := make(chan bool) + for i := 0; i < 10; i++ { + go func() { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "error", + }) + if err != nil { + t.Errorf("Concurrent GrepRaw failed: %v", err) + } + done <- true + }() + } + + for i := 0; i < 10; i++ { + <-done + } + }) + + t.Run("parallel file processing", func(t *testing.T) { + backend := NewInMemoryBackend() + for i := 0; i < 100; i++ { + backend.Write(ctx, &WriteRequest{ + FilePath: fmt.Sprintf("/large/file%d.go", i), + Content: fmt.Sprintf("package main\nimport \"log\"\nfunc test%d() {\n\tlog.Error(\"error %d\")\n}", i, i), + }) + } + + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "log\\.Error", + FileType: "go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 100 { + t.Errorf("Expected 100 matches, got %d", len(matches)) + } + }) + + t.Run("single file no parallelism", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{ + FilePath: "/single.txt", + Content: "error line 1\nerror line 2\nerror line 3", + }) + + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "error", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 3 { + t.Errorf("Expected 3 matches, got %d", len(matches)) + } + }) + + t.Run("empty files list", func(t *testing.T) { + backend := NewInMemoryBackend() + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "anything", + Path: "/nonexistent", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) != 0 { + t.Errorf("Expected 0 matches, got %d", len(matches)) + } + }) + + t.Run("concurrent operations are safe", func(t *testing.T) { + backend := NewInMemoryBackend() + for i := 0; i < 20; i++ { + backend.Write(ctx, &WriteRequest{ + FilePath: fmt.Sprintf("/concurrent/file%d.txt", i), + Content: fmt.Sprintf("line1\nline2\npattern%d\nline4", i), + }) + } + + done := make(chan error, 5) + for i := 0; i < 5; i++ { + go func(id int) { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "pattern\\d+", + }) + done <- err + }(i) + } + + for i := 0; i < 5; i++ { + if err := <-done; err != nil { + t.Errorf("Concurrent operation %d failed: %v", i, err) + } + } + }) +} + +func BenchmarkInMemoryBackend_GrepRaw(b *testing.B) { + backend := NewInMemoryBackend() + ctx := context.Background() + + for i := 0; i < 100; i++ { + content := fmt.Sprintf(`package main + +import ( + "fmt" + "log" +) + +func process%d() error { + log.Error("processing error %d") + fmt.Println("hello world") + return nil +} + +func calculate%d(x, y int) int { + return x + y +} +`, i, i, i) + backend.Write(ctx, &WriteRequest{ + FilePath: fmt.Sprintf("/project/src/file%d.go", i), + Content: content, + }) + } + + b.Run("parallel_grep", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "log\\.Error", + FileType: "go", + }) + if err != nil { + b.Fatalf("GrepRaw failed: %v", err) + } + } + }) + + b.Run("with_glob_filter", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "Error", + Glob: "**/*.go", + }) + if err != nil { + b.Fatalf("GrepRaw failed: %v", err) + } + } + }) + + b.Run("case_insensitive", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "ERROR", + CaseInsensitive: true, + }) + if err != nil { + b.Fatalf("GrepRaw failed: %v", err) + } + } + }) +} + +func TestInMemoryBackend_GrepRaw_ComplexScenarios(t *testing.T) { + backend := NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/src/main.go", + Content: "package main\nimport \"log\"\nfunc main() {\n\tlog.Error(\"error\")\n}", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/src/utils/helper.go", + Content: "package utils\nfunc Helper() error {\n\treturn nil\n}", + }) + backend.Write(ctx, &WriteRequest{ + FilePath: "/project/test/main_test.go", + Content: "package main\nimport \"testing\"\nfunc TestMain(t *testing.T) {\n}", + }) + + t.Run("combine path and file type filters", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "package", + Path: "/project/src", + FileType: "go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasPrefix(match.Path, "/project/src") { + t.Errorf("Expected path to start with /project/src, got: %s", match.Path) + } + if !strings.HasSuffix(match.Path, ".go") { + t.Errorf("Expected .go file, got: %s", match.Path) + } + } + }) + + t.Run("complex regex with case insensitive", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "func\\s+\\w+", + CaseInsensitive: true, + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + if len(matches) == 0 { + t.Error("Expected at least 1 match for function declarations") + } + }) + + t.Run("glob with directory structure", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "package", + Glob: "*_test.go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasSuffix(match.Path, "_test.go") { + t.Errorf("Expected test file, got: %s", match.Path) + } + } + }) + + t.Run("glob with recursive pattern", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "package", + Glob: "**/*.go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasSuffix(match.Path, ".go") { + t.Errorf("Expected .go file, got: %s", match.Path) + } + } + if len(matches) == 0 { + t.Error("Expected at least 1 match for **/*.go pattern") + } + }) + + t.Run("glob with path prefix", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "package", + Glob: "src/**/*.go", + Path: "/project", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if !strings.HasPrefix(match.Path, "/project/src") { + t.Errorf("Expected path to start with /project/src, got: %s", match.Path) + } + if !strings.HasSuffix(match.Path, ".go") { + t.Errorf("Expected .go file, got: %s", match.Path) + } + } + }) + + t.Run("glob simple filename pattern", func(t *testing.T) { + matches, err := backend.GrepRaw(ctx, &GrepRequest{ + Pattern: "package", + Glob: "main.go", + }) + if err != nil { + t.Fatalf("GrepRaw failed: %v", err) + } + for _, match := range matches { + if filepath.Base(match.Path) != "main.go" { + t.Errorf("Expected filename 'main.go', got: %s", match.Path) + } + } + }) +} + +func TestInMemoryBackend_Read_Scenarios(t *testing.T) { + ctx := context.Background() + + t.Run("empty file returns empty content", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/empty.txt", Content: ""}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/empty.txt"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "" { + t.Errorf("expected empty content, got %q", content.Content) + } + }) + + t.Run("single-line file without newline", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/single.txt", Content: "hello"}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/single.txt"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "hello" { + t.Errorf("expected %q, got %q", "hello", content.Content) + } + }) + + t.Run("offset 0 and offset 1 both start from first line", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "a\nb\nc"}) + + c0, _ := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: 0, Limit: 1}) + c1, _ := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: 1, Limit: 1}) + if c0.Content != c1.Content { + t.Errorf("Offset=0 (%q) and Offset=1 (%q) should return the same first line", c0.Content, c1.Content) + } + if c0.Content != "a" { + t.Errorf("expected first line %q, got %q", "a", c0.Content) + } + }) + + t.Run("file with trailing newline preserves trailing empty line", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/trail.txt", Content: "line1\nline2\n"}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/trail.txt"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "line1\nline2\n" { + t.Errorf("expected %q, got %q", "line1\nline2\n", content.Content) + } + lines := strings.Split(content.Content, "\n") + if len(lines) != 3 { // ["line1", "line2", ""] + t.Errorf("expected 3 elements from split, got %d", len(lines)) + } + }) + + t.Run("offset exactly at last line", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "a\nb\nc"}) + + // Offset=3 (1-based) → last line "c" + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: 3, Limit: 10}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "c" { + t.Errorf("expected %q, got %q", "c", content.Content) + } + }) + + t.Run("offset one beyond last line returns empty", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "a\nb\nc"}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: 4, Limit: 10}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "" { + t.Errorf("expected empty content, got %q", content.Content) + } + }) + + t.Run("limit=1 reads exactly one line", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "a\nb\nc"}) + + for i, expected := range []string{"a", "b", "c"} { + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: i + 1, Limit: 1}) + if err != nil { + t.Fatalf("line %d: unexpected error: %v", i+1, err) + } + if content.Content != expected { + t.Errorf("line %d: expected %q, got %q", i+1, expected, content.Content) + } + } + }) + + t.Run("sliding window reads consecutive ranges correctly", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "l1\nl2\nl3\nl4\nl5"}) + + tests := []struct { + offset int + limit int + expected string + }{ + {1, 2, "l1\nl2"}, + {2, 2, "l2\nl3"}, + {3, 2, "l3\nl4"}, + {4, 2, "l4\nl5"}, + {5, 2, "l5"}, + } + for _, tt := range tests { + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: tt.offset, Limit: tt.limit}) + if err != nil { + t.Fatalf("offset=%d limit=%d: unexpected error: %v", tt.offset, tt.limit, err) + } + if content.Content != tt.expected { + t.Errorf("offset=%d limit=%d: expected %q, got %q", tt.offset, tt.limit, tt.expected, content.Content) + } + } + }) + + t.Run("file with only newlines", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/newlines.txt", Content: "\n\n\n"}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/newlines.txt", Offset: 2, Limit: 1}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Line 2 is an empty string between two newlines + if content.Content != "" { + t.Errorf("expected empty line content, got %q", content.Content) + } + }) +} + +func TestInMemoryBackend_Read_NoLimit(t *testing.T) { + ctx := context.Background() + + t.Run("limit=0 reads all lines", func(t *testing.T) { + backend := NewInMemoryBackend() + fullContent := "line1\nline2\nline3\nline4\nline5" + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: fullContent}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Limit: 0}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != fullContent { + t.Errorf("expected %q, got %q", fullContent, content.Content) + } + }) + + t.Run("negative limit reads all lines", func(t *testing.T) { + backend := NewInMemoryBackend() + fullContent := "a\nb\nc" + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: fullContent}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Limit: -1}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != fullContent { + t.Errorf("expected %q, got %q", fullContent, content.Content) + } + }) + + t.Run("limit=0 with offset reads from offset to end", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "a\nb\nc\nd\ne"}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: 3, Limit: 0}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "c\nd\ne" { + t.Errorf("expected %q, got %q", "c\nd\ne", content.Content) + } + }) + + t.Run("limit=0 with offset beyond content returns empty", func(t *testing.T) { + backend := NewInMemoryBackend() + backend.Write(ctx, &WriteRequest{FilePath: "/f.txt", Content: "a\nb"}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/f.txt", Offset: 10, Limit: 0}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != "" { + t.Errorf("expected empty content, got %q", content.Content) + } + }) + + t.Run("limit=0 reads file with more than 2000 lines", func(t *testing.T) { + backend := NewInMemoryBackend() + var b strings.Builder + totalLines := 2500 + for i := 1; i <= totalLines; i++ { + if i > 1 { + b.WriteString("\n") + } + b.WriteString("line" + strconv.Itoa(i)) + } + fullContent := b.String() + backend.Write(ctx, &WriteRequest{FilePath: "/big.txt", Content: fullContent}) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: "/big.txt", Limit: 0}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if content.Content != fullContent { + t.Errorf("expected all %d lines, got %d lines", + totalLines, strings.Count(content.Content, "\n")+1) + } + }) +} diff --git a/adk/flow.go b/adk/flow.go new file mode 100644 index 0000000..e84fb84 --- /dev/null +++ b/adk/flow.go @@ -0,0 +1,809 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "log" + "runtime/debug" + "strings" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/compose" + icb "github.com/cloudwego/eino/internal/callbacks" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +type HistoryEntry struct { + IsUserInput bool + AgentName string + Message Message +} + +type HistoryRewriter func(ctx context.Context, entries []*HistoryEntry) ([]Message, error) + +type flowAgent struct { + Agent + + subAgents []*flowAgent + parentAgent *flowAgent + + disallowTransferToParent bool + historyRewriter HistoryRewriter + + checkPointStore compose.CheckPointStore +} + +func (a *flowAgent) deepCopy() *flowAgent { + ret := &flowAgent{ + Agent: a.Agent, + subAgents: make([]*flowAgent, 0, len(a.subAgents)), + parentAgent: a.parentAgent, + disallowTransferToParent: a.disallowTransferToParent, + historyRewriter: a.historyRewriter, + checkPointStore: a.checkPointStore, + } + + for _, sa := range a.subAgents { + ret.subAgents = append(ret.subAgents, sa.deepCopy()) + } + return ret +} + +// SetSubAgents sets sub-agents for the given agent and returns the updated agent. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func SetSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (ResumableAgent, error) { + return setSubAgents(ctx, agent, subAgents) +} + +type AgentOption func(options *flowAgent) + +// WithDisallowTransferToParent prevents a sub-agent from transferring to its parent. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func WithDisallowTransferToParent() AgentOption { + return func(fa *flowAgent) { + fa.disallowTransferToParent = true + } +} + +// WithHistoryRewriter sets a rewriter to transform conversation history +// during agent transfers. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func WithHistoryRewriter(h HistoryRewriter) AgentOption { + return func(fa *flowAgent) { + fa.historyRewriter = h + } +} + +func toFlowAgent(ctx context.Context, agent Agent, opts ...AgentOption) *flowAgent { + var fa *flowAgent + var ok bool + if fa, ok = agent.(*flowAgent); !ok { + fa = &flowAgent{Agent: agent} + } else { + fa = fa.deepCopy() + } + for _, opt := range opts { + opt(fa) + } + + if fa.historyRewriter == nil { + fa.historyRewriter = buildDefaultHistoryRewriter(agent.Name(ctx)) + } + + return fa +} + +// AgentWithOptions wraps an agent with flow-specific options and returns it. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func AgentWithOptions(ctx context.Context, agent Agent, opts ...AgentOption) Agent { + return toFlowAgent(ctx, agent, opts...) +} + +func setSubAgents(ctx context.Context, agent Agent, subAgents []Agent) (*flowAgent, error) { + fa := toFlowAgent(ctx, agent) + + if len(fa.subAgents) > 0 { + return nil, errors.New("agent's sub-agents has already been set") + } + + if onAgent, ok_ := fa.Agent.(OnSubAgents); ok_ { + err := onAgent.OnSetSubAgents(ctx, subAgents) + if err != nil { + return nil, err + } + } + + for _, s := range subAgents { + fsa := toFlowAgent(ctx, s) + + if fsa.parentAgent != nil { + return nil, errors.New("agent has already been set as a sub-agent of another agent") + } + + fsa.parentAgent = fa + if onAgent, ok__ := fsa.Agent.(OnSubAgents); ok__ { + err := onAgent.OnSetAsSubAgent(ctx, agent) + if err != nil { + return nil, err + } + + if fsa.disallowTransferToParent { + err = onAgent.OnDisallowTransferToParent(ctx) + if err != nil { + return nil, err + } + } + } + + fa.subAgents = append(fa.subAgents, fsa) + } + + return fa, nil +} + +func (a *flowAgent) getAgent(ctx context.Context, name string) *flowAgent { + for _, subAgent := range a.subAgents { + if subAgent.Name(ctx) == name { + return subAgent + } + } + + if a.parentAgent != nil && a.parentAgent.Name(ctx) == name { + return a.parentAgent + } + + return nil +} + +func rewriteMessage(msg Message, agentName string) Message { + var sb strings.Builder + sb.WriteString("For context:") + if msg.Role == schema.Assistant { + if msg.Content != "" { + sb.WriteString(fmt.Sprintf(" [%s] said: %s.", agentName, msg.Content)) + } + if len(msg.ToolCalls) > 0 { + for i := range msg.ToolCalls { + f := msg.ToolCalls[i].Function + sb.WriteString(fmt.Sprintf(" [%s] called tool: `%s` with arguments: %s.", + agentName, f.Name, f.Arguments)) + } + } + } else if msg.Role == schema.Tool && msg.Content != "" { + sb.WriteString(fmt.Sprintf(" [%s] `%s` tool returned result: %s.", + agentName, msg.ToolName, msg.Content)) + } + + rewritten := schema.UserMessage(sb.String()) + if msg.MultiContent != nil { + rewritten.MultiContent = append([]schema.ChatMessagePart{}, msg.MultiContent...) //nolint:staticcheck // backward compat with deprecated MultiContent field + } + if msg.UserInputMultiContent != nil { + rewritten.UserInputMultiContent = append([]schema.MessageInputPart{}, msg.UserInputMultiContent...) + } + + // Convert AssistantGenMultiContent to UserInputMultiContent, since the role changes to User. + // Reasoning parts have no user input equivalent and are dropped. + for _, part := range msg.AssistantGenMultiContent { + switch part.Type { + case schema.ChatMessagePartTypeText: + rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{ + Type: part.Type, + Text: part.Text, + Extra: part.Extra, + }) + case schema.ChatMessagePartTypeImageURL: + if part.Image != nil { + rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{ + Type: part.Type, + Image: &schema.MessageInputImage{MessagePartCommon: part.Image.MessagePartCommon}, + Extra: part.Extra, + }) + } + case schema.ChatMessagePartTypeAudioURL: + if part.Audio != nil { + rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{ + Type: part.Type, + Audio: &schema.MessageInputAudio{MessagePartCommon: part.Audio.MessagePartCommon}, + Extra: part.Extra, + }) + } + case schema.ChatMessagePartTypeVideoURL: + if part.Video != nil { + rewritten.UserInputMultiContent = append(rewritten.UserInputMultiContent, schema.MessageInputPart{ + Type: part.Type, + Video: &schema.MessageInputVideo{MessagePartCommon: part.Video.MessagePartCommon}, + Extra: part.Extra, + }) + } + } + } + + return rewritten +} + +func genMsg(entry *HistoryEntry, agentName string) (Message, error) { + msg := entry.Message + if entry.AgentName != agentName { + msg = rewriteMessage(msg, entry.AgentName) + } + + return msg, nil +} + +func deepCopyAgentInput(ai *AgentInput) *AgentInput { + copied := &AgentInput{ + Messages: make([]Message, len(ai.Messages)), + EnableStreaming: ai.EnableStreaming, + } + + copy(copied.Messages, ai.Messages) + + return copied +} + +func (a *flowAgent) genAgentInput(ctx context.Context, runCtx *runContext, skipTransferMessages bool) (*AgentInput, error) { + input := deepCopyAgentInput(runCtx.RootInput) + + events := runCtx.Session.getEvents() + historyEntries := make([]*HistoryEntry, 0) + + for _, m := range input.Messages { + historyEntries = append(historyEntries, &HistoryEntry{ + IsUserInput: true, + Message: m, + }) + } + + for _, event := range events { + if skipTransferMessages && event.Action != nil && event.Action.TransferToAgent != nil { + // If skipTransferMessages is true and the event contain transfer action, the message in this event won't be appended to history entries. + if event.Output != nil && + event.Output.MessageOutput != nil && + event.Output.MessageOutput.Role == schema.Tool && + len(historyEntries) > 0 { + // If the skipped message's role is Tool, remove the previous history entry as it's also a transfer message(from ChatModelAgent and GenTransferMessages). + historyEntries = historyEntries[:len(historyEntries)-1] + } + continue + } + + msg, err := getMessageFromWrappedEvent(event) + if err != nil { + var retryErr *WillRetryError + if errors.As(err, &retryErr) { + log.Printf("failed to get message from event, but will retry: %v", err) + continue + } + return nil, err + } + + if msg == nil { + continue + } + + historyEntries = append(historyEntries, &HistoryEntry{ + AgentName: event.AgentName, + Message: msg, + }) + } + + messages, err := a.historyRewriter(ctx, historyEntries) + if err != nil { + return nil, err + } + input.Messages = messages + + return input, nil +} + +func buildDefaultHistoryRewriter(agentName string) HistoryRewriter { + return func(ctx context.Context, entries []*HistoryEntry) ([]Message, error) { + messages := make([]Message, 0, len(entries)) + var err error + for _, entry := range entries { + msg := entry.Message + if !entry.IsUserInput { + msg, err = genMsg(entry, agentName) + if err != nil { + return nil, fmt.Errorf("gen agent input failed: %w", err) + } + } + + if msg != nil { + messages = append(messages, msg) + } + } + + return messages, nil + } +} + +func (a *flowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + agentName := a.Name(ctx) + + var runCtx *runContext + ctx, runCtx = initRunCtx(ctx, agentName, input) + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName) + + o := getCommonOptions(nil, opts...) + cancelCtx := o.cancelCtx + + processedInput, err := a.genAgentInput(ctx, runCtx, o.skipTransferMessages) + if err != nil { + if cancelCtx != nil { + cancelCtx.markDone() + } + cbInput := &AgentCallbackInput{Input: input} + ctx = callbacks.OnStart(ctx, cbInput) + return wrapIterWithOnEnd(ctx, genErrorIter(err)) + } + + ctxForSubAgents := ctx + + agentType := getAgentType(a.Agent) + ctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...) + cbInput := &AgentCallbackInput{Input: processedInput} + ctx = callbacks.OnStart(ctx, cbInput) + + input = processedInput + + if wf, ok := a.Agent.(*workflowAgent); ok { + ctx = withCancelContext(ctx, cancelCtx) + filteredOpts := filterCancelOption(filterCallbackHandlersForNestedAgents(agentName, opts)) + iter := wf.Run(ctx, input, filteredOpts...) + iter = wrapIterWithCancelCtx(iter, cancelCtx) + return wrapIterWithOnEnd(ctx, iter) + } + + aIter := a.Agent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...) + + return wrapIterWithCancelCtx(iterator, cancelCtx) +} + +func (a *flowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + agentName := a.Name(ctx) + + ctx, info = buildResumeInfo(ctx, agentName, info) + + ctxForSubAgents := ctx + + o := getCommonOptions(nil, opts...) + cancelCtx := o.cancelCtx + + agentType := getAgentType(a.Agent) + ctx = initAgentCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...) + cbInput := &AgentCallbackInput{ResumeInfo: info} + ctx = callbacks.OnStart(ctx, cbInput) + + if info.WasInterrupted { + if ra, ok := a.Agent.(ResumableAgent); ok { + if _, ok := ra.(*workflowAgent); ok { + ctx = withCancelContext(ctx, cancelCtx) + filteredOpts := filterCancelOption(filterCallbackHandlersForNestedAgents(agentName, opts)) + aIter := ra.Resume(ctx, info, filteredOpts...) + aIter = wrapIterWithCancelCtx(aIter, cancelCtx) + return wrapIterWithOnEnd(ctx, aIter) + } + + aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...) + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...) + return wrapIterWithCancelCtx(iterator, cancelCtx) + } + + if cancelCtx != nil { + cancelCtx.markDone() + } + return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+ + "but is not a ResumableAgent", agentName))) + } + + nextAgentName, err := getNextResumeAgent(ctx, info) + if err != nil { + if cancelCtx != nil { + cancelCtx.markDone() + } + return wrapIterWithOnEnd(ctx, genErrorIter(err)) + } + + subAgent := a.getAgent(ctxForSubAgents, nextAgentName) + if subAgent == nil { + if len(a.subAgents) == 0 { + if ra, ok := a.Agent.(ResumableAgent); ok { + ctx = withCancelContext(ctx, cancelCtx) + innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...) + return wrapIterWithCancelCtx(wrapIterWithOnEnd(ctx, innerIter), cancelCtx) + } + return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf( + "failed to resume agent: agent '%s' (type %T) has no sub-agents and does not implement ResumableAgent interface. "+ + "To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.Agent))) + } + if cancelCtx != nil { + cancelCtx.markDone() + } + return wrapIterWithOnEnd(ctx, genErrorIter(fmt.Errorf("failed to resume agent: sub-agent '%s' not found in agent '%s'", nextAgentName, agentName))) + } + + ctxForSubAgents = withCancelContext(ctxForSubAgents, cancelCtx) + innerIter := subAgent.Resume(ctxForSubAgents, info, filterCancelOption(opts)...) + return wrapIterWithCancelCtx(wrapIterWithOnEnd(ctx, innerIter), cancelCtx) +} + +// DeterministicTransferConfig is the configuration for AgentWithDeterministicTransferTo. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +type DeterministicTransferConfig struct { + Agent Agent + ToAgentNames []string +} + +func (a *flowAgent) run( + ctx context.Context, + ctxForSubAgents context.Context, + runCtx *runContext, + aIter *AsyncIterator[*AgentEvent], + generator *AsyncGenerator[*AgentEvent], + opts ...AgentRunOption) { + + cbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]() + + cbOutput := &AgentCallbackOutput{Events: cbIter} + icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false) + + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&AgentEvent{Err: e}) + } + + cbGen.Close() + generator.Close() + }() + + var lastAction *AgentAction + for { + event, ok := aIter.Next() + if !ok { + break + } + + // RunPath ownership: the eino framework sets RunPath exactly once. + // If event.RunPath is already set (e.g., by agentTool), we don't modify it. + // If event.RunPath is nil/empty, we set it to the current runCtx.RunPath. + // This ensures RunPath is set exactly once and not duplicated. + if len(event.RunPath) == 0 { + event.AgentName = a.Name(ctx) + event.RunPath = runCtx.RunPath + } + // Recording policy: exact RunPath match (non-interrupt) indicates events belonging to this agent execution. + // This prevents parent recording of child/tool-internal emissions. + if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) { + // copy the event so that the copied event's stream is exclusive for any potential consumer + // copy before adding to session because once added to session it's stream could be consumed by genAgentInput at any time + // interrupt action are not added to session, because ALL information contained in it + // is either presented to end-user, or made available to agents through other means + copied := copyTypedAgentEvent(event) + setAutomaticClose(copied) + setAutomaticClose(event) + runCtx.Session.addEvent(copied) + } + // Action gating uses exact run-path match as well: + // only actions originating from this agent execution (not child/tool runs) + // should influence parent control flow (exit/transfer/interrupt). + if exactRunPathMatch(runCtx.RunPath, event.RunPath) { + lastAction = event.Action + } + copied := copyTypedAgentEvent(event) + setAutomaticClose(copied) + setAutomaticClose(event) + cbGen.Send(copied) + generator.Send(event) + } + + var destName string + if lastAction != nil { + if lastAction.Interrupted != nil { + return + } + if lastAction.Exit { + return + } + + if lastAction.TransferToAgent != nil { + destName = lastAction.TransferToAgent.DestAgentName + } + } + + // handle transferring to another agent + if destName != "" { + agentToRun := a.getAgent(ctxForSubAgents, destName) + if agentToRun == nil { + e := fmt.Errorf("transfer failed: agent '%s' not found when transferring from '%s'", + destName, a.Name(ctxForSubAgents)) + generator.Send(&AgentEvent{Err: e}) + return + } + + subAIter := agentToRun.Run(ctxForSubAgents, nil /*subagents get input from runCtx*/, opts...) + for { + subEvent, ok_ := subAIter.Next() + if !ok_ { + break + } + + setAutomaticClose(subEvent) + generator.Send(subEvent) + } + } +} + +func exactRunPathMatch(aPath, bPath []RunStep) bool { + if len(aPath) != len(bPath) { + return false + } + for i := range aPath { + if !aPath[i].Equals(bPath[i]) { + return false + } + } + return true +} + +func wrapIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*AgentEvent]) *AsyncIterator[*AgentEvent] { + cbIter, cbGen := NewAsyncIteratorPair[*AgentEvent]() + cbOutput := &AgentCallbackOutput{Events: cbIter} + icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyAgentCallbackOutput), callbacks.TimingOnEnd, false) + + outIter, outGen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer func() { + cbGen.Close() + outGen.Close() + }() + for { + event, ok := iter.Next() + if !ok { + break + } + copied := copyTypedAgentEvent(event) + cbGen.Send(copied) + outGen.Send(event) + } + }() + return outIter +} + +// --------------------------------------------------------------------------- +// Typed wrapper for the agentic path (TypedAgent[*schema.AgenticMessage]). +// +// typedFlowAgent is a minimal wrapper used exclusively by TypedRunner and +// AgentTool to execute a TypedAgent[*schema.AgenticMessage]. It handles +// callbacks, event recording, and run-path tracking. Transfer, sub-agent +// orchestration, and history rewriting are handled solely by the concrete +// flowAgent (the *schema.Message path). +// --------------------------------------------------------------------------- + +type typedFlowAgent[M MessageType] struct { + TypedAgent[M] + + checkPointStore compose.CheckPointStore +} + +func toTypedFlowAgent[M MessageType](agent TypedAgent[M]) *typedFlowAgent[M] { + if fa, ok := agent.(*typedFlowAgent[M]); ok { + return fa + } + return &typedFlowAgent[M]{TypedAgent: agent} +} + +func getTypedAgentType[M MessageType](agent TypedAgent[M]) string { + if msgAgent, ok := any(agent).(Agent); ok { + return getAgentType(msgAgent) + } + if typer, ok := any(agent).(interface{ GetType() string }); ok { + return typer.GetType() + } + return "" +} + +func (a *typedFlowAgent[M]) Run(ctx context.Context, input *TypedAgentInput[M], opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + agentName := a.Name(ctx) + + var runCtx *runContext + ctx, runCtx = initTypedRunCtx(ctx, agentName, input) + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, agentName) + + o := getCommonOptions(nil, opts...) + cancelCtx := o.cancelCtx + + ctxForSubAgents := ctx + + agentType := getTypedAgentType(a.TypedAgent) + ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...) + cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{Input: any(input).(*TypedAgentInput[*schema.AgenticMessage])} + ctx = callbacks.OnStart(ctx, cbInput) + + aIter := a.TypedAgent.Run(withCancelContext(ctx, cancelCtx), input, filterOptions(agentName, opts)...) + + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + + go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), runCtx, aIter, generator, filterCancelOption(opts)...) + + return wrapIterWithCancelCtx(iterator, cancelCtx) +} + +func (a *typedFlowAgent[M]) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + agentName := a.Name(ctx) + + ctx, info = buildResumeInfo(ctx, agentName, info) + + ctxForSubAgents := ctx + + o := getCommonOptions(nil, opts...) + cancelCtx := o.cancelCtx + + agentType := getTypedAgentType(a.TypedAgent) + ctx = initAgenticCallbacks(ctx, agentName, agentType, filterOptions(agentName, opts)...) + cbInput := &TypedAgentCallbackInput[*schema.AgenticMessage]{ResumeInfo: info} + ctx = callbacks.OnStart(ctx, cbInput) + + if info.WasInterrupted { + if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok { + aIter := ra.Resume(withCancelContext(ctx, cancelCtx), info, opts...) + + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go a.run(withCancelContext(ctx, cancelCtx), withCancelContext(ctxForSubAgents, cancelCtx), getRunCtx(ctxForSubAgents), aIter, generator, filterCancelOption(opts)...) + return wrapIterWithCancelCtx(iterator, cancelCtx) + } + + if cancelCtx != nil { + cancelCtx.markDone() + } + return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf("failed to resume agent: agent '%s' is an interrupt point "+ + "but is not a ResumableAgent", agentName)) + } + + _, err := getNextResumeAgent(ctx, info) + if err != nil { + if cancelCtx != nil { + cancelCtx.markDone() + } + return typedErrorIterWithOnEnd[M](ctx, err) + } + + if ra, ok := a.TypedAgent.(TypedResumableAgent[M]); ok { + ctx = withCancelContext(ctx, cancelCtx) + innerIter := ra.Resume(ctx, info, filterCancelOption(opts)...) + return wrapIterWithCancelCtx(typedWrapIterWithOnEnd(ctx, innerIter), cancelCtx) + } + return typedErrorIterWithOnEnd[M](ctx, fmt.Errorf( + "failed to resume agent: agent '%s' (type %T) does not implement ResumableAgent interface. "+ + "To support resume, your custom agent wrapper must implement the ResumableAgent interface", agentName, a.TypedAgent)) +} + +func (a *typedFlowAgent[M]) run( + ctx context.Context, + _ context.Context, + runCtx *runContext, + aIter *AsyncIterator[*TypedAgentEvent[M]], + generator *AsyncGenerator[*TypedAgentEvent[M]], + _ ...AgentRunOption) { + + agenticCbIter, agenticCbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: agenticCbIter} + icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false) + + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&TypedAgentEvent[M]{Err: e}) + } + + agenticCbGen.Close() + generator.Close() + }() + + for { + event, ok := aIter.Next() + if !ok { + break + } + + if len(event.RunPath) == 0 { + event.AgentName = a.Name(ctx) + event.RunPath = runCtx.RunPath + } + if (event.Action == nil || event.Action.Interrupted == nil) && exactRunPathMatch(runCtx.RunPath, event.RunPath) { + copied := copyTypedAgentEvent(event) + typedSetAutomaticClose(copied) + typedSetAutomaticClose(event) + addTypedEvent(runCtx.Session, copied) + } + + agenticCopied := copyTypedAgentEvent(event) + typedSetAutomaticClose(agenticCopied) + typedSetAutomaticClose(event) + agenticCbGen.Send(any(agenticCopied).(*TypedAgentEvent[*schema.AgenticMessage])) + generator.Send(event) + } +} + +func wrapAgenticIterWithOnEnd(ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + cbIter, cbGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + cbOutput := &TypedAgentCallbackOutput[*schema.AgenticMessage]{Events: cbIter} + icb.On(ctx, cbOutput, icb.BuildOnEndHandleWithCopy(copyTypedCallbackOutput[*schema.AgenticMessage]), callbacks.TimingOnEnd, false) + + outIter, outGen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer func() { + cbGen.Close() + outGen.Close() + }() + for { + event, ok := iter.Next() + if !ok { + break + } + copied := copyTypedAgentEvent(event) + cbGen.Send(copied) + outGen.Send(event) + } + }() + return outIter +} + +func genAgenticErrorIter(err error) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{Err: err}) + gen.Close() + return iter +} + +func typedWrapIterWithOnEnd[M MessageType](ctx context.Context, iter *AsyncIterator[*TypedAgentEvent[M]]) *AsyncIterator[*TypedAgentEvent[M]] { + agenticIter := any(iter).(*AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]]) + return any(wrapAgenticIterWithOnEnd(ctx, agenticIter)).(*AsyncIterator[*TypedAgentEvent[M]]) +} + +func typedErrorIterWithOnEnd[M MessageType](ctx context.Context, err error) *AsyncIterator[*TypedAgentEvent[M]] { + return typedWrapIterWithOnEnd(ctx, typedErrorIter[M](err)) +} diff --git a/adk/flow_test.go b/adk/flow_test.go new file mode 100644 index 0000000..1cb095e --- /dev/null +++ b/adk/flow_test.go @@ -0,0 +1,389 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/callbacks" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +func strPtr(s string) *string { return &s } + +func TestRewriteMessage(t *testing.T) { + imageCommon := schema.MessagePartCommon{URL: strPtr("http://img.example.com")} + audioCommon := schema.MessagePartCommon{URL: strPtr("http://audio.example.com")} + videoCommon := schema.MessagePartCommon{URL: strPtr("http://video.example.com")} + + msg := &schema.Message{ + Role: schema.Assistant, + Content: "hello", + MultiContent: []schema.ChatMessagePart{ //nolint:staticcheck // testing deprecated field backward compat + {Type: schema.ChatMessagePartTypeText, Text: "legacy"}, + }, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "pre-existing"}, + }, + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "gen-text", Extra: map[string]any{"k": "v"}}, + {Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageOutputImage{MessagePartCommon: imageCommon}}, + {Type: schema.ChatMessagePartTypeAudioURL, Audio: &schema.MessageOutputAudio{MessagePartCommon: audioCommon}}, + {Type: schema.ChatMessagePartTypeVideoURL, Video: &schema.MessageOutputVideo{MessagePartCommon: videoCommon}}, + {Type: schema.ChatMessagePartTypeReasoning, Reasoning: &schema.MessageOutputReasoning{Text: "secret thoughts"}}, + }, + } + + rewritten := rewriteMessage(msg, "OtherAgent") + + assert.Equal(t, schema.User, rewritten.Role) + + // MultiContent: copied, not shared + assert.Equal(t, msg.MultiContent, rewritten.MultiContent) + rewritten.MultiContent[0].Text = "mutated" + assert.Equal(t, "legacy", msg.MultiContent[0].Text) + + // UserInputMultiContent: pre-existing entry copied, AssistantGenMultiContent appended (reasoning dropped) + assert.Len(t, rewritten.UserInputMultiContent, 5) // 1 pre-existing + 4 converted (text/image/audio/video) + + // pre-existing entry is not shared + rewritten.UserInputMultiContent[0].Text = "mutated" + assert.Equal(t, "pre-existing", msg.UserInputMultiContent[0].Text) + + // text conversion + assert.Equal(t, schema.ChatMessagePartTypeText, rewritten.UserInputMultiContent[1].Type) + assert.Equal(t, "gen-text", rewritten.UserInputMultiContent[1].Text) + assert.Equal(t, map[string]any{"k": "v"}, rewritten.UserInputMultiContent[1].Extra) + + // image conversion + assert.Equal(t, schema.ChatMessagePartTypeImageURL, rewritten.UserInputMultiContent[2].Type) + assert.Equal(t, imageCommon, rewritten.UserInputMultiContent[2].Image.MessagePartCommon) + + // audio conversion + assert.Equal(t, schema.ChatMessagePartTypeAudioURL, rewritten.UserInputMultiContent[3].Type) + assert.Equal(t, audioCommon, rewritten.UserInputMultiContent[3].Audio.MessagePartCommon) + + // video conversion + assert.Equal(t, schema.ChatMessagePartTypeVideoURL, rewritten.UserInputMultiContent[4].Type) + assert.Equal(t, videoCommon, rewritten.UserInputMultiContent[4].Video.MessagePartCommon) + + // reasoning is dropped; AssistantGenMultiContent is not set on rewritten message + assert.Empty(t, rewritten.AssistantGenMultiContent) +} + +// TestTransferToAgent tests the TransferToAgent functionality +func TestTransferToAgent(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create mock models for parent and child agents + parentModel := mockModel.NewMockToolCallingChatModel(ctrl) + childModel := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the parent model + // First call: parent model generates a message with TransferToAgent tool call + parentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("I'll transfer this to the child agent", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "ChildAgent"}`, + }, + }, + }), nil). + Times(1) + + // Set up expectations for the child model + // Second call: child model generates a response + childModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello from child agent", nil), nil). + Times(1) + + // Both models should implement WithTools + parentModel.EXPECT().WithTools(gomock.Any()).Return(parentModel, nil).AnyTimes() + childModel.EXPECT().WithTools(gomock.Any()).Return(childModel, nil).AnyTimes() + + // Create parent agent + parentAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ParentAgent", + Description: "Parent agent that will transfer to child", + Instruction: "You are a parent agent.", + Model: parentModel, + }) + assert.NoError(t, err) + assert.NotNil(t, parentAgent) + + // Create child agent + childAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ChildAgent", + Description: "Child agent that handles specific tasks", + Instruction: "You are a child agent.", + Model: childModel, + }) + assert.NoError(t, err) + assert.NotNil(t, childAgent) + + // Set up parent-child relationship + flowAgent, err := SetSubAgents(ctx, parentAgent, []Agent{childAgent}) + assert.NoError(t, err) + assert.NotNil(t, flowAgent) + + assert.NotNil(t, parentAgent.subAgents) + assert.NotNil(t, childAgent.parentAgent) + + // Run the parent agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Please transfer this to the child agent"), + }, + } + ctx, _ = initRunCtx(ctx, flowAgent.Name(ctx), input) + iterator := flowAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // First event: parent model output with tool call + event1, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event1) + assert.Nil(t, event1.Err) + assert.NotNil(t, event1.Output) + assert.NotNil(t, event1.Output.MessageOutput) + assert.Equal(t, schema.Assistant, event1.Output.MessageOutput.Role) + + // Second event: tool output (TransferToAgent) + event2, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event2) + assert.Nil(t, event2.Err) + assert.NotNil(t, event2.Output) + assert.NotNil(t, event2.Output.MessageOutput) + assert.Equal(t, schema.Tool, event2.Output.MessageOutput.Role) + + // Verify the action is TransferToAgent + assert.NotNil(t, event2.Action) + assert.NotNil(t, event2.Action.TransferToAgent) + assert.Equal(t, "ChildAgent", event2.Action.TransferToAgent.DestAgentName) + + // Third event: child model output + event3, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event3) + assert.Nil(t, event3.Err) + assert.NotNil(t, event3.Output) + assert.NotNil(t, event3.Output.MessageOutput) + assert.Equal(t, schema.Assistant, event3.Output.MessageOutput.Role) + + // Verify the message content from child agent + msg := event3.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.Equal(t, "Hello from child agent", msg.Content) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestResumeNonResumableWrapperError(t *testing.T) { + ctx := context.Background() + + interruptingAgent := &nonResumableFlowTestAgent{ + name: "inner", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(Interrupt(ctx, "please confirm")) + }() + return iter + }, + } + + innerFlowAgent := toFlowAgent(ctx, interruptingAgent) + + wrapper := &nonResumableFlowTestAgent{ + name: "wrapper", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return innerFlowAgent.Run(ctx, input, options...) + }, + } + + store := &flowTestStore{m: map[string][]byte{}} + runner := NewRunner(ctx, RunnerConfig{ + Agent: wrapper, + EnableStreaming: true, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("test")}, WithCheckPointID("cp1")) + + var interruptID string + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Action != nil && ev.Action.Interrupted != nil { + for _, intCtx := range ev.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + interruptID = intCtx.ID + } + } + } + } + assert.NotEmpty(t, interruptID, "should have an interrupt ID") + + resumeIter, err := runner.ResumeWithParams(ctx, "cp1", &ResumeParams{ + Targets: map[string]any{interruptID: nil}, + }) + assert.NoError(t, err) + + var resumeErr error + for { + ev, ok := resumeIter.Next() + if !ok { + break + } + if ev.Err != nil { + resumeErr = ev.Err + } + } + + assert.Error(t, resumeErr) + assert.Contains(t, resumeErr.Error(), "does not implement ResumableAgent interface") + assert.Contains(t, resumeErr.Error(), "custom agent wrapper must implement the ResumableAgent interface") +} + +type nonResumableFlowTestAgent struct { + name string + runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] +} + +func (a *nonResumableFlowTestAgent) Name(_ context.Context) string { + return a.name +} + +func (a *nonResumableFlowTestAgent) Description(_ context.Context) string { + return a.name + " description" +} + +func (a *nonResumableFlowTestAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return a.runFn(ctx, input, options...) +} + +type flowTestStore struct { + m map[string][]byte +} + +func (s *flowTestStore) Set(_ context.Context, key string, value []byte) error { + s.m[key] = value + return nil +} + +func (s *flowTestStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.m[key] + return v, ok, nil +} + +func TestTransferToAgentWithDesignatedCallback(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + parentModel := mockModel.NewMockToolCallingChatModel(ctrl) + childModel := mockModel.NewMockToolCallingChatModel(ctrl) + + parentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("I'll transfer this to the child agent", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: TransferToAgentToolName, + Arguments: `{"agent_name": "ChildAgent"}`, + }, + }, + }), nil). + Times(1) + + childModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello from child agent", nil), nil). + Times(1) + + parentModel.EXPECT().WithTools(gomock.Any()).Return(parentModel, nil).AnyTimes() + childModel.EXPECT().WithTools(gomock.Any()).Return(childModel, nil).AnyTimes() + + parentAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ParentAgent", + Description: "Parent agent that will transfer to child", + Instruction: "You are a parent agent.", + Model: parentModel, + }) + assert.NoError(t, err) + + childAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ChildAgent", + Description: "Child agent that handles specific tasks", + Instruction: "You are a child agent.", + Model: childModel, + }) + assert.NoError(t, err) + + flowAgent, err := SetSubAgents(ctx, parentAgent, []Agent{childAgent}) + assert.NoError(t, err) + + var childCallbackCount int + var mu sync.Mutex + + handler := callbacks.NewHandlerBuilder().OnStartFn( + func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component == ComponentOfAgent && info.Name == "ChildAgent" { + mu.Lock() + childCallbackCount++ + mu.Unlock() + } + return ctx + }).Build() + + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Please transfer this to the child agent"), + }, + } + ctx, _ = initRunCtx(ctx, flowAgent.Name(ctx), input) + iterator := flowAgent.Run(ctx, input, WithCallbacks(handler).DesignateAgent("ChildAgent")) + + for { + _, ok := iterator.Next() + if !ok { + break + } + } + + assert.Equal(t, 1, childCallbackCount, "designated callback for ChildAgent should fire exactly once during transfer") +} diff --git a/adk/handler.go b/adk/handler.go new file mode 100644 index 0000000..f952441 --- /dev/null +++ b/adk/handler.go @@ -0,0 +1,451 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "encoding/gob" + "fmt" + "io" + "reflect" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// InvokableToolCallEndpoint is the function signature for invoking a tool synchronously. +// Middleware authors implement wrappers around this endpoint to add custom behavior. +type InvokableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) + +// StreamableToolCallEndpoint is the function signature for invoking a tool with streaming output. +// Middleware authors implement wrappers around this endpoint to add custom behavior. +type StreamableToolCallEndpoint func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) + +type EnhancedInvokableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) + +type EnhancedStreamableToolCallEndpoint func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) + +// ToolContext provides metadata about the tool being wrapped. +type ToolContext struct { + Name string + CallID string +} + +// ToolCallsContext contains metadata about the tool calls that just completed. +type ToolCallsContext struct { + // ToolCalls contains the tool call metadata from the model's response. + ToolCalls []ToolContext +} + +// TypedModelContext contains context information passed to WrapModel. +type TypedModelContext[M MessageType] struct { + // Tools contains the current tool list configured for the agent. + // This is populated at request time with the tools that will be sent to the model. + // + // Deprecated: Use TypedChatModelAgentState.ToolInfos in BeforeModelRewriteState instead. + // ModelContext.Tools remains populated for backward compatibility with existing WrapModel handlers, + // but new code should read and modify state.ToolInfos which is the source of truth for the model call. + Tools []*schema.ToolInfo + + // ModelRetryConfig contains the retry configuration for the model. + // This is populated at request time from the agent's ModelRetryConfig. + // Used by EventSenderModelWrapper to wrap stream errors appropriately. + ModelRetryConfig *TypedModelRetryConfig[M] + + // ModelFailoverConfig contains the failover configuration for the model. + // This is populated at request time from the agent's ModelFailoverConfig. + // Used by EventSenderModelWrapper to wrap stream errors so that failed failover + // attempts are skipped (not treated as fatal) by the flow event processor. + ModelFailoverConfig *ModelFailoverConfig[M] + + cancelContext *cancelContext +} + +// ModelContext is the default model context type using *schema.Message. +type ModelContext = TypedModelContext[*schema.Message] + +// ChatModelAgentContext contains runtime information passed to handlers before each ChatModelAgent run. +// Handlers can modify Instruction, Tools, and ReturnDirectly to customize agent behavior. +// +// This type is specific to ChatModelAgent. Other agent types may define their own context types. +type ChatModelAgentContext struct { + // Instruction is the current instruction for the Agent execution. + // It includes the instruction configured for the agent, additional instructions appended by framework + // and AgentMiddleware, and modifications applied by previous BeforeAgent handlers. + // The finalized instruction after all BeforeAgent handlers are then passed to GenModelInput, + // to be (optionally) formatted with SessionValues and converted to system message. + Instruction string + + // Tools are the raw tools (without any wrapper or tool middleware) currently configured for the Agent execution. + // They includes tools passed in AgentConfig, implicit tools added by framework such as transfer / exit tools, + // and other tools already added by middlewares. + Tools []tool.BaseTool + + // ReturnDirectly is the set of tool names currently configured to cause the Agent to return directly. + // This is based on the return directly map configured for the agent, plus any modifications + // by previous BeforeAgent handlers. + ReturnDirectly map[string]bool + + // ToolSearchTool is the tool info for the model's native tool search capability. + // When set by a BeforeAgent handler, the framework passes it to the model via model.WithToolSearchTool. + ToolSearchTool *schema.ToolInfo +} + +// TypedChatModelAgentMiddleware defines the interface for customizing TypedChatModelAgent behavior. +// +// IMPORTANT: This interface is specifically designed for TypedChatModelAgent and agents built +// on top of it (e.g., DeepAgent). +// +// Why TypedChatModelAgentMiddleware instead of AgentMiddleware? +// +// AgentMiddleware is a struct type, which has inherent limitations: +// - Struct types are closed: users cannot add new methods to extend functionality +// - The framework only recognizes AgentMiddleware's fixed fields, so even if users +// embed AgentMiddleware in a custom struct and add methods, the framework cannot +// call those methods (config.Middlewares is []AgentMiddleware, not a user type) +// - Callbacks in AgentMiddleware only return error, cannot return modified context +// +// TypedChatModelAgentMiddleware is an interface type, which is open for extension: +// - Users can implement custom handlers with arbitrary internal state and methods +// - Hook methods return (context.Context, ..., error) for direct context propagation +// - Wrapper methods (WrapToolCall, WrapModel) enable context propagation through the +// wrapped endpoint chain: wrappers can pass modified context to the next wrapper +// - Configuration is centralized in struct fields rather than scattered in closures +// +// TypedChatModelAgentMiddleware vs AgentMiddleware: +// - Use AgentMiddleware for simple, static additions (extra instruction/tools) +// - Use TypedChatModelAgentMiddleware for dynamic behavior, context modification, or call wrapping +// - AgentMiddleware is kept for backward compatibility with existing users +// - Both can be used together; see AgentMiddleware documentation for execution order +// +// Use *TypedBaseChatModelAgentMiddleware as an embedded struct to provide default no-op +// implementations for all methods. +type TypedChatModelAgentMiddleware[M MessageType] interface { + // BeforeAgent is called before each agent run, allowing modification of + // the agent's instruction and tools configuration. + BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) + + // AfterAgent is called after the agent run reaches a successful terminal state. + // Successful terminal states are: final answer (model response with no tool calls), + // and return-directly tool result. + // + // AfterAgent is NOT called when the agent terminates with an error (e.g., + // ErrExceedMaxIterations, context cancellation, model errors). + // + // The state parameter contains the final conversation state, including all messages + // from the completed run. + // + // AfterAgent handlers are called in the same order as BeforeAgent handlers + // (first registered = first called). Consistent with all other middleware hooks, + // if any handler returns an error, subsequent handlers are NOT called (fail-fast) + // and the error is sent to the event stream. + AfterAgent(ctx context.Context, state *TypedChatModelAgentState[M]) (context.Context, error) + + // BeforeModelRewriteState is called before each model invocation. + // The returned state is persisted to the agent's internal state and passed to the model. + // The returned context is propagated to the model call and subsequent handlers. + // + // The ChatModelAgentState struct provides access to: + // - Messages: the conversation history + // - ToolInfos: the tool list that will be sent to the model (modifiable) + // - DeferredToolInfos: tools for server-side search (modifiable, nil if unused) + // + // This is the recommended place to modify messages and tools before a model call. + // Changes here are persisted in state and reflected in subsequent iterations. + BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error) + + // AfterModelRewriteState is called after each model invocation. + // The input state includes the model's response as the last message. + // The returned state is persisted to the agent's internal state. + // + // The ChatModelAgentState struct provides access to: + // - Messages: the conversation history including the model's response + // - ToolInfos: the tool list that was sent to the model + // - DeferredToolInfos: tools for server-side search (nil if unused) + AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error) + + // WrapInvokableToolCall wraps a tool's synchronous execution with custom behavior. + // Return the input endpoint unchanged and nil error if no wrapping is needed. + // + // This method is only called for tools that implement InvokableTool. + // If a tool only implements StreamableTool, this method will not be called for that tool. + // + // This method is called at request time when the tool is about to be executed. + // The tCtx parameter provides metadata about the tool: + // - Name: The name of the tool being wrapped + // - CallID: The unique identifier for this specific tool call + WrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error) + + // WrapStreamableToolCall wraps a tool's streaming execution with custom behavior. + // Return the input endpoint unchanged and nil error if no wrapping is needed. + // + // This method is only called for tools that implement StreamableTool. + // If a tool only implements InvokableTool, this method will not be called for that tool. + // + // This method is called at request time when the tool is about to be executed. + // The tCtx parameter provides metadata about the tool: + // - Name: The name of the tool being wrapped + // - CallID: The unique identifier for this specific tool call + WrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error) + + // WrapEnhancedInvokableToolCall wraps an enhanced tool's synchronous execution with custom behavior. + // Return the input endpoint unchanged and nil error if no wrapping is needed. + // + // This method is only called for tools that implement EnhancedInvokableTool. + // If a tool only implements EnhancedStreamableTool, this method will not be called for that tool. + // + // This method is called at request time when the tool is about to be executed. + // The tCtx parameter provides metadata about the tool: + // - Name: The name of the tool being wrapped + // - CallID: The unique identifier for this specific tool call + WrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error) + + // WrapEnhancedStreamableToolCall wraps an enhanced tool's streaming execution with custom behavior. + // Return the input endpoint unchanged and nil error if no wrapping is needed. + // + // This method is only called for tools that implement EnhancedStreamableTool. + // If a tool only implements EnhancedInvokableTool, this method will not be called for that tool. + // + // This method is called at request time when the tool is about to be executed. + // The tCtx parameter provides metadata about the tool: + // - Name: The name of the tool being wrapped + // - CallID: The unique identifier for this specific tool call + WrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error) + + // WrapModel wraps a chat model with custom behavior around the actual model call. + // Return the input model unchanged and nil error if no wrapping is needed. + // + // This method is called at request time when the model is about to be invoked. + // Note: The parameter is model.BaseModel[M] (not ToolCallingChatModel) because wrappers + // only need to intercept Generate/Stream calls. Tool binding (WithTools) is handled + // separately by the framework and does not flow through user wrappers. + // + // Recommended use cases (behavior around the model call itself): + // - Model call retry logic + // - Model failover (switching to a backup model) + // - Sending events (e.g. streaming progress) + // - Processing or transforming the response stream + // - Changing call configurations (temperature, top_p, etc.) + // + // Discouraged use cases (use BeforeModelRewriteState instead): + // - Modifying input messages: changes here are NOT persisted in state, only + // affect a single model call, and break prompt cache across iterations. + // - Modifying the tool list: use state.ToolInfos / state.DeferredToolInfos in + // BeforeModelRewriteState, which is the source of truth for tool configuration. + // + // The mc parameter provides read-only context about the current model call: + // - Tools: The tool infos that will be sent to the model (Deprecated: read state.ToolInfos instead) + WrapModel(ctx context.Context, m model.BaseModel[M], mc *TypedModelContext[M]) (model.BaseModel[M], error) +} + +// ChatModelAgentMiddleware is the default middleware type using *schema.Message. +// See TypedChatModelAgentMiddleware for full documentation. +type ChatModelAgentMiddleware = TypedChatModelAgentMiddleware[*schema.Message] + +type TypedBaseChatModelAgentMiddleware[M MessageType] struct{} + +// BaseChatModelAgentMiddleware provides default no-op implementations for ChatModelAgentMiddleware. +// Embed *BaseChatModelAgentMiddleware in custom handlers to only override the methods you need. +// +// Example: +// +// type MyHandler struct { +// *adk.BaseChatModelAgentMiddleware +// // custom fields +// } +// +// func (h *MyHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, mc *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) { +// // custom logic +// return ctx, state, nil +// } +type BaseChatModelAgentMiddleware = TypedBaseChatModelAgentMiddleware[*schema.Message] + +func (b *TypedBaseChatModelAgentMiddleware[M]) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) (InvokableToolCallEndpoint, error) { + return endpoint, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) (StreamableToolCallEndpoint, error) { + return endpoint, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) (EnhancedInvokableToolCallEndpoint, error) { + return endpoint, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) (EnhancedStreamableToolCallEndpoint, error) { + return endpoint, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) WrapModel(_ context.Context, m model.BaseModel[M], _ *TypedModelContext[M]) (model.BaseModel[M], error) { + return m, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + return ctx, runCtx, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) AfterAgent(ctx context.Context, state *TypedChatModelAgentState[M]) (context.Context, error) { + return ctx, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error) { + return ctx, state, nil +} + +func (b *TypedBaseChatModelAgentMiddleware[M]) AfterModelRewriteState(ctx context.Context, state *TypedChatModelAgentState[M], mc *TypedModelContext[M]) (context.Context, *TypedChatModelAgentState[M], error) { + return ctx, state, nil +} + +func processTypedState(ctx context.Context, fn func(extra map[string]any) map[string]any) error { + runCtx := getRunCtx(ctx) + if runCtx != nil && runCtx.AgenticRootInput != nil { + return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.AgenticMessage]) error { + st.Extra = fn(st.Extra) + return nil + }) + } + return compose.ProcessState(ctx, func(_ context.Context, st *typedState[*schema.Message]) error { + st.Extra = fn(st.Extra) + return nil + }) +} + +// SetRunLocalValue sets a key-value pair that persists for the duration of the current agent Run() invocation. +// The value is scoped to this specific execution and is not shared across different Run() calls or agent instances. +// +// Values stored here are compatible with interrupt/resume cycles - they will be serialized and restored +// when the agent is resumed. For custom types, you must register them using schema.RegisterName[T]() +// in an init() function to ensure proper serialization. +// +// This function can only be called from within a ChatModelAgentMiddleware during agent execution. +// Returns an error if called outside of an agent execution context. +func SetRunLocalValue(ctx context.Context, key string, value any) error { + if err := checkGobEncodability(key, value); err != nil { + return err + } + + err := processTypedState(ctx, func(extra map[string]any) map[string]any { + if extra == nil { + extra = make(map[string]any) + } + extra[key] = value + return extra + }) + if err != nil { + return fmt.Errorf("SetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err) + } + + return nil +} + +// GetRunLocalValue retrieves a value that was set during the current agent Run() invocation. +// The value is scoped to this specific execution and is not shared across different Run() calls or agent instances. +// +// Values stored via SetRunLocalValue are compatible with interrupt/resume cycles - they will be serialized +// and restored when the agent is resumed. For custom types, you must register them using schema.RegisterName[T]() +// in an init() function to ensure proper serialization. +// +// This function can only be called from within a ChatModelAgentMiddleware during agent execution. +// Returns the value and true if found, or nil and false if not found or if called outside of an agent execution context. +func GetRunLocalValue(ctx context.Context, key string) (any, bool, error) { + var val any + var found bool + err := processTypedState(ctx, func(extra map[string]any) map[string]any { + if extra != nil { + val, found = extra[key] + } + return extra + }) + if err != nil { + return nil, false, fmt.Errorf("GetRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err) + } + return val, found, nil +} + +// DeleteRunLocalValue removes a value that was set during the current agent Run() invocation. +// +// This function can only be called from within a ChatModelAgentMiddleware during agent execution. +// Returns an error if called outside of an agent execution context. +func DeleteRunLocalValue(ctx context.Context, key string) error { + err := processTypedState(ctx, func(extra map[string]any) map[string]any { + if extra != nil { + delete(extra, key) + } + return extra + }) + if err != nil { + return fmt.Errorf("DeleteRunLocalValue failed: must be called within a ChatModelAgent Run() or Resume() execution context: %w", err) + } + return nil +} + +// TypedSendEvent sends a custom TypedAgentEvent to the event stream during agent execution. +// This allows TypedChatModelAgentMiddleware implementations to emit custom events that will be +// received by the caller iterating over the agent's event stream. +// +// Note: TypedSendEvent is a pure transport — it does NOT auto-assign message IDs. +// Framework-created messages (model output, tool results) receive IDs automatically +// via internal wrapper layers. If your middleware constructs its own messages, call +// EnsureMessageID before sending to assign an ID. +// +// This function can only be called from within a TypedChatModelAgentMiddleware during agent execution. +// Returns an error if called outside of an agent execution context. +func TypedSendEvent[M MessageType](ctx context.Context, event *TypedAgentEvent[M]) error { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil || execCtx.generator == nil { + return fmt.Errorf("TypedSendEvent failed: must be called within a ChatModelAgent Run() or Resume() execution context") + } + + execCtx.send(event) + return nil +} + +// SendEvent sends a custom AgentEvent to the event stream during agent execution. +// This allows ChatModelAgentMiddleware implementations to emit custom events that will be +// received by the caller iterating over the agent's event stream. +// +// This function can only be called from within a ChatModelAgentMiddleware during agent execution. +// Returns an error if called outside of an agent execution context. +func SendEvent(ctx context.Context, event *AgentEvent) error { + return TypedSendEvent(ctx, event) +} + +// checkGobEncodability probes whether the value can be gob-encoded as part of +// a map[string]any, which is exactly how State.Extra is serialized during +// checkpoint. This catches unregistered types early at Set time, rather than +// letting them fail at checkpoint/resume time with a confusing error. +func checkGobEncodability(key string, value any) error { + probe := map[string]any{key: value} + if err := gob.NewEncoder(io.Discard).Encode(probe); err != nil { + typeName := reflect.TypeOf(value).String() + return fmt.Errorf("SetRunLocalValue: the value (type %s) for key %q is not gob-serializable, "+ + "which means it will fail when the agent checkpoint is saved or resumed.\n\n"+ + "To fix this, register the type in an init() function in your package:\n\n"+ + " func init() {\n"+ + " schema.RegisterName[%s](\"a_unique_name_for_this_type\")\n"+ + " }\n\n"+ + "This is required because agent state (including values set via SetRunLocalValue) is "+ + "persisted using gob encoding for interrupt/resume support. All concrete types stored "+ + "in interface-typed fields (like map[string]any) must be registered with gob.\n\n"+ + "If this value does not need to survive interrupt/resume, store it on the context instead, "+ + "for example via context.WithValue, so you don't need gob registration.\n\n"+ + "Underlying error: %w", typeName, key, typeName, err) + } + return nil +} diff --git a/adk/handler_test.go b/adk/handler_test.go new file mode 100644 index 0000000..811cd2b --- /dev/null +++ b/adk/handler_test.go @@ -0,0 +1,2571 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "fmt" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type testInstructionHandler struct { + *BaseChatModelAgentMiddleware + text string +} + +func (h *testInstructionHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + if runCtx.Instruction == "" { + runCtx.Instruction = h.text + } else if h.text != "" { + runCtx.Instruction = runCtx.Instruction + "\n" + h.text + } + return ctx, runCtx, nil +} + +type testInstructionFuncHandler struct { + *BaseChatModelAgentMiddleware + fn func(ctx context.Context, instruction string) (context.Context, string, error) +} + +func (h *testInstructionFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + newCtx, newInstruction, err := h.fn(ctx, runCtx.Instruction) + if err != nil { + return ctx, runCtx, err + } + runCtx.Instruction = newInstruction + return newCtx, runCtx, nil +} + +type testToolsHandler struct { + *BaseChatModelAgentMiddleware + tools []tool.BaseTool +} + +func (h *testToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + runCtx.Tools = append(runCtx.Tools, h.tools...) + return ctx, runCtx, nil +} + +type testToolsFuncHandler struct { + *BaseChatModelAgentMiddleware + fn func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) +} + +func (h *testToolsFuncHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + newCtx, newTools, newReturnDirectly, err := h.fn(ctx, runCtx.Tools, runCtx.ReturnDirectly) + if err != nil { + return ctx, runCtx, err + } + runCtx.Tools = newTools + runCtx.ReturnDirectly = newReturnDirectly + return newCtx, runCtx, nil +} + +type testBeforeAgentHandler struct { + *BaseChatModelAgentMiddleware + fn func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) +} + +func (h *testBeforeAgentHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + return h.fn(ctx, runCtx) +} + +type testBeforeModelRewriteStateHandler struct { + *BaseChatModelAgentMiddleware + fn func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) +} + +func (h *testBeforeModelRewriteStateHandler) BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + return h.fn(ctx, state, mc) +} + +type testAfterModelRewriteStateHandler struct { + *BaseChatModelAgentMiddleware + fn func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) +} + +func (h *testAfterModelRewriteStateHandler) AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + return h.fn(ctx, state, mc) +} + +type testToolWrapperHandler struct { + *BaseChatModelAgentMiddleware + wrapInvokableFn func(context.Context, InvokableToolCallEndpoint, *ToolContext) InvokableToolCallEndpoint + wrapStreamableFn func(context.Context, StreamableToolCallEndpoint, *ToolContext) StreamableToolCallEndpoint +} + +func (h *testToolWrapperHandler) WrapInvokableToolCall(ctx context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error) { + if h.wrapInvokableFn != nil { + return h.wrapInvokableFn(ctx, endpoint, tCtx), nil + } + return endpoint, nil +} + +func (h *testToolWrapperHandler) WrapStreamableToolCall(ctx context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error) { + if h.wrapStreamableFn != nil { + return h.wrapStreamableFn(ctx, endpoint, tCtx), nil + } + return endpoint, nil +} + +type testModelWrapperHandler struct { + *BaseChatModelAgentMiddleware + fn func(context.Context, model.BaseChatModel, *ModelContext) model.BaseChatModel +} + +func (h *testModelWrapperHandler) WrapModel(ctx context.Context, m model.BaseChatModel, mc *ModelContext) (model.BaseChatModel, error) { + return h.fn(ctx, m, mc), nil +} + +func newTestInvokableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, InvokableToolCallEndpoint, *ToolContext) InvokableToolCallEndpoint { + return func(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) InvokableToolCallEndpoint { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + if beforeFn != nil { + beforeFn() + } + result, err := endpoint(ctx, argumentsInJSON, opts...) + if afterFn != nil { + afterFn() + } + return result, err + } + } +} + +func newResultModifyingInvokableToolCallWrapper(modifyFn func(string) string) func(context.Context, InvokableToolCallEndpoint, *ToolContext) InvokableToolCallEndpoint { + return func(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) InvokableToolCallEndpoint { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + result, err := endpoint(ctx, argumentsInJSON, opts...) + if err == nil && modifyFn != nil { + result = modifyFn(result) + } + return result, err + } + } +} + +func newTestStreamableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, StreamableToolCallEndpoint, *ToolContext) StreamableToolCallEndpoint { + return func(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) StreamableToolCallEndpoint { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + if beforeFn != nil { + beforeFn() + } + result, err := endpoint(ctx, argumentsInJSON, opts...) + if afterFn != nil { + afterFn() + } + return result, err + } + } +} + +func TestHandlerExecutionOrder(t *testing.T) { + t.Run("MultipleInstructionHandlersPipeline", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var capturedInstruction string + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + if len(msgs) > 0 && msgs[0].Role == schema.System { + capturedInstruction = msgs[0].Content + } + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "Base instruction.", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testInstructionHandler{text: "Handler 1 addition."}, + &testInstructionHandler{text: "Handler 2 addition."}, + &testInstructionFuncHandler{fn: func(ctx context.Context, instruction string) (context.Context, string, error) { + return ctx, instruction + "\nHandler 3 dynamic.", nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Contains(t, capturedInstruction, "Base instruction.") + assert.Contains(t, capturedInstruction, "Handler 1 addition.") + assert.Contains(t, capturedInstruction, "Handler 2 addition.") + assert.Contains(t, capturedInstruction, "Handler 3 dynamic.") + }) + + t.Run("MiddlewaresBeforeHandlers", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var capturedInstruction string + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + if len(msgs) > 0 && msgs[0].Role == schema.System { + capturedInstruction = msgs[0].Content + } + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "Base.", + Model: cm, + Middlewares: []AgentMiddleware{ + {AdditionalInstruction: "Middleware instruction."}, + }, + Handlers: []ChatModelAgentMiddleware{ + &testInstructionHandler{text: "Handler instruction."}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + middlewareIdx := len(capturedInstruction) - len("Middleware instruction.") - len("\nHandler instruction.") + handlerIdx := len(capturedInstruction) - len("Handler instruction.") + assert.True(t, middlewareIdx < handlerIdx, "Middleware should be applied before Handler") + }) +} + +func TestToolsHandlerCombinations(t *testing.T) { + t.Run("MultipleToolsHandlersAppend", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &fakeToolForTest{tarCount: 1} + tool2 := &fakeToolForTest{tarCount: 2} + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + var capturedToolCount int + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + options := model.GetCommonOptions(&model.Options{}, opts...) + capturedToolCount = len(options.Tools) + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolsHandler{tools: []tool.BaseTool{tool2}}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, 2, capturedToolCount) + }) + + t.Run("ToolsFuncCanRemoveTools", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "tool1"} + tool2 := &namedTool{name: "tool2"} + tool3 := &namedTool{name: "tool3"} + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + var capturedToolNames []string + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + options := model.GetCommonOptions(&model.Options{}, opts...) + for _, t := range options.Tools { + capturedToolNames = append(capturedToolNames, t.Name) + } + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1, tool2, tool3}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolsFuncHandler{fn: func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) { + filtered := make([]tool.BaseTool, 0) + for _, t := range tools { + info, _ := t.Info(ctx) + if info.Name != "tool2" { + filtered = append(filtered, t) + } + } + return ctx, filtered, returnDirectly, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Contains(t, capturedToolNames, "tool1") + assert.NotContains(t, capturedToolNames, "tool2") + assert.Contains(t, capturedToolNames, "tool3") + }) + + t.Run("ReturnDirectlyModification", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "tool1"} + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: "tool1", Arguments: "{}"}}, + }), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolsFuncHandler{fn: func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) { + for _, t := range tools { + info, _ := t.Info(ctx) + if info.Name == "tool1" { + returnDirectly[info.Name] = true + } + } + return ctx, tools, returnDirectly, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + eventCount := 0 + for { + event, ok := iter.Next() + if !ok { + break + } + eventCount++ + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.Message != nil && + event.Output.MessageOutput.Message.Role == schema.Tool { + assert.Equal(t, "tool1 result", event.Output.MessageOutput.Message.Content) + } + } + assert.Equal(t, 2, eventCount) + }) + + t.Run("DynamicToolCanBeCalledByModel", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + dynamicToolCalled := false + dynamicTool := &callableTool{ + name: "dynamic_tool", + invokeFn: func() { + dynamicToolCalled = true + }, + } + info, _ := dynamicTool.Info(ctx) + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using dynamic tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil).Times(1) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testToolsHandler{tools: []tool.BaseTool{dynamicTool}}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, dynamicToolCalled, "Dynamic tool should have been called") + }) +} + +func TestMessageRewriteHandlers(t *testing.T) { + t.Run("BeforeModelRewriteStatePipeline", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var capturedMsgCount int + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + capturedMsgCount = len(msgs) + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "instruction", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + state.Messages = append(state.Messages, schema.UserMessage("injected1")) + return ctx, state, nil + }}, + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + state.Messages = append(state.Messages, schema.UserMessage("injected2")) + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("original")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, 4, capturedMsgCount) + }) + + t.Run("AfterModelRewriteState", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + afterCalled := false + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + afterCalled = true + assert.True(t, len(state.Messages) > 0) + lastMsg := state.Messages[len(state.Messages)-1] + assert.Equal(t, schema.Assistant, lastMsg.Role) + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, afterCalled) + }) +} + +func TestToolCallWrapperHandlers(t *testing.T) { + t.Run("MultipleToolWrappersPipeline", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &namedTool{name: "test_tool"} + info, _ := testTool.Info(ctx) + + var callOrder []string + var mu sync.Mutex + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil).Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolWrapperHandler{wrapInvokableFn: newTestInvokableToolCallWrapper( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper1-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper1-after") + mu.Unlock() + }, + )}, + &testToolWrapperHandler{wrapInvokableFn: newTestInvokableToolCallWrapper( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper2-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper2-after") + mu.Unlock() + }, + )}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, []string{"wrapper1-before", "wrapper2-before", "wrapper2-after", "wrapper1-after"}, callOrder) + }) + + t.Run("StreamingToolWrappersPipeline", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &streamingNamedTool{name: "streaming_tool"} + info, _ := testTool.Info(ctx) + + var callOrder []string + var mu sync.Mutex + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), + }), nil).Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("done", nil), + }), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolWrapperHandler{wrapStreamableFn: newTestStreamableToolCallWrapper( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper1-stream-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper1-stream-after") + mu.Unlock() + }, + )}, + &testToolWrapperHandler{wrapStreamableFn: newTestStreamableToolCallWrapper( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper2-stream-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper2-stream-after") + mu.Unlock() + }, + )}, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()}) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var hasStreamingToolResult bool + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && + event.Output.MessageOutput.Role == schema.Tool { + hasStreamingToolResult = true + for { + _, err := event.Output.MessageOutput.MessageStream.Recv() + if err != nil { + break + } + } + } + } + + assert.True(t, hasStreamingToolResult, "Should have streaming tool result") + assert.Equal(t, []string{"wrapper1-stream-before", "wrapper2-stream-before", "wrapper2-stream-after", "wrapper1-stream-after"}, callOrder, + "Streaming wrappers should be called in correct order") + }) + + t.Run("ToolWrapperCanModifyResult", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &namedTool{name: "test_tool"} + info, _ := testTool.Info(ctx) + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil).Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolWrapperHandler{wrapInvokableFn: newResultModifyingInvokableToolCallWrapper(func(result string) string { + return "modified: " + result + })}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.Message != nil && + event.Output.MessageOutput.Message.Role == schema.Tool { + assert.Equal(t, "modified: test_tool result", event.Output.MessageOutput.Message.Content) + } + } + }) +} + +func TestToolContextFunctions(t *testing.T) { + t.Run("ModelContextToolsField", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &namedTool{name: "base_tool"} + info, _ := testTool.Info(ctx) + + var wrapperSeenTools []*schema.ToolInfo + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, mc *ModelContext) model.BaseChatModel { + return &toolChainingTestModel{ + inner: m, + mc: mc, + wrapFn: func(ctx context.Context, opts []model.Option) []model.Option { + wrapperSeenTools = mc.Tools + return opts + }, + } + }, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Len(t, wrapperSeenTools, 1, "Wrapper should see base tool") + assert.Equal(t, info.Name, wrapperSeenTools[0].Name, "Wrapper should see base_tool") + }) +} + +type toolChainingTestModel struct { + inner model.BaseChatModel + mc *ModelContext + wrapFn func(ctx context.Context, opts []model.Option) []model.Option +} + +func (m *toolChainingTestModel) Generate(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + if m.wrapFn != nil { + opts = m.wrapFn(ctx, opts) + } + return m.inner.Generate(ctx, msgs, opts...) +} + +func (m *toolChainingTestModel) Stream(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + if m.wrapFn != nil { + opts = m.wrapFn(ctx, opts) + } + return m.inner.Stream(ctx, msgs, opts...) +} + +func (m *toolChainingTestModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +func TestContextPropagation(t *testing.T) { + t.Run("ContextPassedThroughBeforeModelHandlers", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + type ctxKey string + const key1 ctxKey = "key1" + const key2 ctxKey = "key2" + + var handler2ReceivedValue1 interface{} + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + return context.WithValue(ctx, key1, "value1"), state, nil + }}, + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + handler2ReceivedValue1 = ctx.Value(key1) + return context.WithValue(ctx, key2, "value2"), state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, "value1", handler2ReceivedValue1, "Handler 2 should receive context value set by Handler 1") + }) + + t.Run("BeforeAgentContextPropagation", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + type ctxKey string + const key1 ctxKey = "key1" + + var handler2ReceivedValue interface{} + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + return context.WithValue(ctx, key1, "value1"), runCtx, nil + }}, + &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + handler2ReceivedValue = ctx.Value(key1) + return ctx, runCtx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, "value1", handler2ReceivedValue, "Handler 2 should receive context value set by Handler 1 during BeforeAgent") + }) +} + +func TestCustomHandler(t *testing.T) { + t.Run("CustomHandlerWithState", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + customHandler := &countingHandler{} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{customHandler}, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, 1, customHandler.beforeAgentCount) + assert.Equal(t, 1, customHandler.afterAgentCount) + assert.Equal(t, 1, customHandler.beforeModelCount) + assert.Equal(t, 1, customHandler.afterModelCount) + }) +} + +func TestHandlerErrorHandling(t *testing.T) { + t.Run("BeforeAgentErrorStopsRun", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeAgentHandler{fn: func(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + return ctx, runCtx, assert.AnError + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []*schema.Message{schema.UserMessage("test")}, + }) + + var gotErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + gotErr = event.Err + } + } + + assert.Error(t, gotErr) + assert.Contains(t, gotErr.Error(), "BeforeAgent failed") + }) +} + +type namedTool struct { + name string +} + +func (t *namedTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: t.name + " description"}, nil +} + +func (t *namedTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return t.name + " result", nil +} + +type streamingNamedTool struct { + name string +} + +func (t *streamingNamedTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: t.name + " description"}, nil +} + +func (t *streamingNamedTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return t.name + " result", nil +} + +func (t *streamingNamedTool) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) { + return schema.StreamReaderFromArray([]string{t.name + " stream result"}), nil +} + +type callableTool struct { + name string + invokeFn func() +} + +func (t *callableTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: t.name, Desc: t.name + " description"}, nil +} + +func (t *callableTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + if t.invokeFn != nil { + t.invokeFn() + } + return t.name + " result", nil +} + +type countingHandler struct { + *BaseChatModelAgentMiddleware + beforeAgentCount int + afterAgentCount int + beforeModelCount int + afterModelCount int + mu sync.Mutex +} + +func (h *countingHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + h.mu.Lock() + h.beforeAgentCount++ + h.mu.Unlock() + return ctx, runCtx, nil +} + +func (h *countingHandler) AfterAgent(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + h.mu.Lock() + h.afterAgentCount++ + h.mu.Unlock() + return ctx, nil +} + +func (h *countingHandler) BeforeModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + h.mu.Lock() + h.beforeModelCount++ + h.mu.Unlock() + return ctx, state, nil +} + +func (h *countingHandler) AfterModelRewriteState(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + h.mu.Lock() + h.afterModelCount++ + h.mu.Unlock() + return ctx, state, nil +} + +func newTestModelWrapperFn(beforeFn, afterFn func()) func(context.Context, model.BaseChatModel, *ModelContext) model.BaseChatModel { + return func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return &testWrappedModel{ + inner: m, + beforeFn: beforeFn, + afterFn: afterFn, + } + } +} + +type testWrappedModel struct { + inner model.BaseChatModel + beforeFn func() + afterFn func() +} + +func (m *testWrappedModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + if m.beforeFn != nil { + m.beforeFn() + } + result, err := m.inner.Generate(ctx, input, opts...) + if m.afterFn != nil { + m.afterFn() + } + return result, err +} + +func (m *testWrappedModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + if m.beforeFn != nil { + m.beforeFn() + } + result, err := m.inner.Stream(ctx, input, opts...) + if m.afterFn != nil { + m.afterFn() + } + return result, err +} + +func TestModelWrapperHandlers(t *testing.T) { + t.Run("MultipleModelWrappersPipeline", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callOrder []string + var mu sync.Mutex + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testModelWrapperHandler{fn: newTestModelWrapperFn( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper1-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper1-after") + mu.Unlock() + }, + )}, + &testModelWrapperHandler{fn: newTestModelWrapperFn( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper2-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper2-after") + mu.Unlock() + }, + )}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, []string{"wrapper1-before", "wrapper2-before", "wrapper2-after", "wrapper1-after"}, callOrder) + }) + + t.Run("ModelWrapperBeforeAfterCallOrder", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var callOrder []string + var mu sync.Mutex + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + mu.Lock() + callOrder = append(callOrder, "model-generate") + mu.Unlock() + return schema.AssistantMessage("original response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testModelWrapperHandler{fn: newTestModelWrapperFn( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper-after") + mu.Unlock() + }, + )}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, []string{"wrapper-before", "model-generate", "wrapper-after"}, callOrder) + }) + + t.Run("ModelWrapperWithTools", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &namedTool{name: "test_tool"} + info, _ := testTool.Info(ctx) + + var callOrder []string + var mu sync.Mutex + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + mu.Lock() + callOrder = append(callOrder, "model-call") + mu.Unlock() + return schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil + }).Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + mu.Lock() + callOrder = append(callOrder, "model-call") + mu.Unlock() + return schema.AssistantMessage("done", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testModelWrapperHandler{fn: newTestModelWrapperFn( + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper-before") + mu.Unlock() + }, + func() { + mu.Lock() + callOrder = append(callOrder, "wrapper-after") + mu.Unlock() + }, + )}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, []string{ + "wrapper-before", "model-call", "wrapper-after", + "wrapper-before", "model-call", "wrapper-after", + }, callOrder) + }) +} + +func newInputModifyingWrapperFn(inputPrefix string) func(context.Context, model.BaseChatModel, *ModelContext) model.BaseChatModel { + return func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return &inputOutputModifyingModel{ + inner: m, + inputPrefix: inputPrefix, + } + } +} + +type inputOutputModifyingModel struct { + inner model.BaseChatModel + inputPrefix string +} + +func (m *inputOutputModifyingModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + modifiedMessages := make([]*schema.Message, len(input)) + for i, msg := range input { + if msg.Role == schema.User { + modifiedMessages[i] = schema.UserMessage(m.inputPrefix + msg.Content) + } else { + modifiedMessages[i] = msg + } + } + return m.inner.Generate(ctx, modifiedMessages, opts...) +} + +func (m *inputOutputModifyingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + modifiedMessages := make([]*schema.Message, len(input)) + for i, msg := range input { + if msg.Role == schema.User { + modifiedMessages[i] = schema.UserMessage(m.inputPrefix + msg.Content) + } else { + modifiedMessages[i] = msg + } + } + return m.inner.Stream(ctx, modifiedMessages, opts...) +} + +func TestModelWrapper_InputModification(t *testing.T) { + t.Run("ModelWrapperModifiesInput_Generate", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var modelReceivedInput []*schema.Message + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + modelReceivedInput = input + return schema.AssistantMessage("original response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testModelWrapperHandler{fn: newInputModifyingWrapperFn("[WRAPPER]")}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test input")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.NotNil(t, modelReceivedInput) + assert.True(t, len(modelReceivedInput) > 0) + found := false + for _, msg := range modelReceivedInput { + if msg.Content == "[WRAPPER]test input" { + found = true + break + } + } + assert.True(t, found, "Model should receive wrapper-modified input") + }) + + t.Run("ModelWrapperModifiesInput_Stream", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var modelReceivedInput []*schema.Message + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + modelReceivedInput = input + return schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("chunk1", nil), + schema.AssistantMessage("chunk2", nil), + }), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testModelWrapperHandler{fn: newInputModifyingWrapperFn("[WRAPPER]")}, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true, CheckPointStore: newBridgeStore()}) + iter := r.Run(ctx, []Message{schema.UserMessage("test input")}) + + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && + event.Output.MessageOutput.Role == schema.Assistant { + for { + _, err := event.Output.MessageOutput.MessageStream.Recv() + if err != nil { + break + } + } + } + } + + assert.NotNil(t, modelReceivedInput) + assert.True(t, len(modelReceivedInput) > 0) + found := false + for _, msg := range modelReceivedInput { + if msg.Content == "[WRAPPER]test input" { + found = true + break + } + } + assert.True(t, found, "Model should receive wrapper-modified input") + }) +} + +func TestRunLocalValueFunctions(t *testing.T) { + t.Run("SetAndGetRunLocalValue", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var capturedValue any + var capturedFound bool + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + err := SetRunLocalValue(ctx, "test_key", "test_value") + assert.NoError(t, err) + return ctx, state, nil + }}, + &testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + val, found, err := GetRunLocalValue(ctx, "test_key") + assert.NoError(t, err) + capturedValue = val + capturedFound = found + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, capturedFound, "Value should be found") + assert.Equal(t, "test_value", capturedValue, "Value should match what was set") + }) + + t.Run("DeleteRunLocalValue", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var valueAfterDelete any + var foundAfterDelete bool + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + err := SetRunLocalValue(ctx, "delete_key", "delete_value") + assert.NoError(t, err) + + err = DeleteRunLocalValue(ctx, "delete_key") + assert.NoError(t, err) + return ctx, state, nil + }}, + &testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + val, found, err := GetRunLocalValue(ctx, "delete_key") + assert.NoError(t, err) + valueAfterDelete = val + foundAfterDelete = found + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.False(t, foundAfterDelete, "Value should not be found after deletion") + assert.Nil(t, valueAfterDelete, "Value should be nil after deletion") + }) + + t.Run("GetNonExistentKey", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var capturedValue any + var capturedFound bool + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + val, found, err := GetRunLocalValue(ctx, "non_existent_key") + assert.NoError(t, err) + capturedValue = val + capturedFound = found + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.False(t, capturedFound, "Non-existent key should not be found") + assert.Nil(t, capturedValue, "Non-existent key should return nil value") + }) + + t.Run("RunLocalValueGobEncodabilityCheck", func(t *testing.T) { + type unregisteredType struct { + Data string + } + + // unregistered custom struct should fail + err := checkGobEncodability("key", unregisteredType{Data: "hello"}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "not gob-serializable") + assert.Contains(t, err.Error(), "schema.RegisterName") + assert.Contains(t, err.Error(), "context.WithValue") + assert.Contains(t, err.Error(), "unregisteredType") + + // primitives should succeed + assert.NoError(t, checkGobEncodability("key", "hello")) + assert.NoError(t, checkGobEncodability("key", 42)) + assert.NoError(t, checkGobEncodability("key", true)) + assert.NoError(t, checkGobEncodability("key", 3.14)) + }) + + t.Run("RunLocalValueOutsideContext", func(t *testing.T) { + ctx := context.Background() + + err := SetRunLocalValue(ctx, "key", "value") + assert.Error(t, err, "SetRunLocalValue should fail outside agent context") + assert.Contains(t, err.Error(), "SetRunLocalValue failed") + + _, _, err = GetRunLocalValue(ctx, "key") + assert.Error(t, err, "GetRunLocalValue should fail outside agent context") + assert.Contains(t, err.Error(), "GetRunLocalValue failed") + + err = DeleteRunLocalValue(ctx, "key") + assert.Error(t, err, "DeleteRunLocalValue should fail outside agent context") + assert.Contains(t, err.Error(), "DeleteRunLocalValue failed") + }) + + t.Run("RunLocalValuePersistsAcrossModelCalls", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &namedTool{name: "test_tool"} + info, _ := testTool.Info(ctx) + + var firstCallValue any + var secondCallValue any + callCount := 0 + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil).Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + callCount++ + if callCount == 1 { + err := SetRunLocalValue(ctx, "persist_key", "persist_value") + assert.NoError(t, err) + val, _, _ := GetRunLocalValue(ctx, "persist_key") + firstCallValue = val + } else { + val, _, _ := GetRunLocalValue(ctx, "persist_key") + secondCallValue = val + } + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, "persist_value", firstCallValue, "First call should set value") + assert.Equal(t, "persist_value", secondCallValue, "Value should persist to second model call") + }) +} + +func TestHandlerErrorPropagation(t *testing.T) { + t.Run("BeforeModelRewriteStateErrorStopsRun", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + return ctx, state, assert.AnError + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + var gotErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + gotErr = event.Err + } + } + + assert.Error(t, gotErr) + }) + + t.Run("AfterModelRewriteStateErrorStopsRun", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + return ctx, state, assert.AnError + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + + var gotErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + gotErr = event.Err + } + } + + assert.Error(t, gotErr) + }) + + t.Run("MultipleHandlersFirstErrorStops", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + secondHandlerCalled := false + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + return ctx, state, assert.AnError + }}, + &testBeforeModelRewriteStateHandler{fn: func(ctx context.Context, state *ChatModelAgentState, mc *ModelContext) (context.Context, *ChatModelAgentState, error) { + secondHandlerCalled = true + return ctx, state, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.False(t, secondHandlerCalled, "Second handler should not be called after first handler error") + }) +} + +func TestToolContextInWrappers(t *testing.T) { + t.Run("ToolContextHasCorrectNameAndCallID", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + testTool := &namedTool{name: "context_test_tool"} + info, _ := testTool.Info(ctx) + + var capturedToolName string + var capturedCallID string + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "test_call_id_123", Function: schema.FunctionCall{Name: info.Name, Arguments: "{}"}}, + }), nil).Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapInvokableFn: func(_ context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) InvokableToolCallEndpoint { + capturedToolName = tCtx.Name + capturedCallID = tCtx.CallID + return endpoint + }, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, "context_test_tool", capturedToolName, "ToolContext should have correct tool name") + assert.Equal(t, "test_call_id_123", capturedCallID, "ToolContext should have correct call ID") + }) +} + +func TestAfterToolCallsHook(t *testing.T) { + t.Run("CalledAfterToolCalls", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "tool_alpha"} + tool2 := &namedTool{name: "tool_beta"} + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // First call: model returns two tool calls + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling tools", []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "tool_alpha", Arguments: "{}"}}, + {ID: "call_2", Function: schema.FunctionCall{Name: "tool_beta", Arguments: "{}"}}, + }), nil).Times(1) + + // Second call: model returns final response + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + var mu sync.Mutex + callCount := 0 + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1, tool2}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + WithAfterToolCallsHook(func(ctx context.Context) error { + mu.Lock() + callCount++ + mu.Unlock() + return nil + })) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + // Should be called exactly once (one iteration with tool calls) + assert.Equal(t, 1, callCount) + }) + + t.Run("NotCalledWithoutToolCalls", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Model returns a direct response with no tool calls + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("direct response", nil), nil).Times(1) + + callCount := 0 + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + WithAfterToolCallsHook(func(ctx context.Context) error { + callCount++ + return nil + })) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, 0, callCount, "AfterToolCallsHook should not be called when no tool calls happen") + }) + + t.Run("ToolResultsInStateBeforeHookFires", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "mytool"} + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // First call: model returns a tool call + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling", []schema.ToolCall{ + {ID: "c1", Function: schema.FunctionCall{Name: "mytool", Arguments: "{}"}}, + }), nil).Times(1) + + // Second call: final response + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final", nil), nil).Times(1) + + var hookToolResultCount int + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("original")}}, + WithAfterToolCallsHook(func(ctx context.Context) error { + // Verify tool results are already in state when the hook fires + _ = compose.ProcessState(ctx, func(_ context.Context, st *State) error { + for _, msg := range st.Messages { + if msg.Role == schema.Tool { + hookToolResultCount++ + } + } + return nil + }) + return nil + })) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, 1, hookToolResultCount, "Tool results should be in state when hook fires") + }) + + t.Run("HookErrorPropagation", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "mytool"} + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling", []schema.ToolCall{ + {ID: "c1", Function: schema.FunctionCall{Name: "mytool", Arguments: "{}"}}, + }), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + WithAfterToolCallsHook(func(ctx context.Context) error { + return fmt.Errorf("hook failure") + })) + + var sawError bool + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + assert.Contains(t, ev.Err.Error(), "hook failure") + sawError = true + } + } + assert.True(t, sawError, "hook error should propagate as an agent error event") + }) + + t.Run("HookCalledPerIteration", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "mytool"} + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Iteration 1: tool call + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling1", []schema.ToolCall{ + {ID: "c1", Function: schema.FunctionCall{Name: "mytool", Arguments: "{}"}}, + }), nil).Times(1) + + // Iteration 2: tool call again + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling2", []schema.ToolCall{ + {ID: "c2", Function: schema.FunctionCall{Name: "mytool", Arguments: "{}"}}, + }), nil).Times(1) + + // Iteration 3: final answer + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil).Times(1) + + var mu sync.Mutex + hookCount := 0 + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}, + WithAfterToolCallsHook(func(ctx context.Context) error { + mu.Lock() + hookCount++ + mu.Unlock() + return nil + })) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, 2, hookCount, "hook should fire once per tool-call iteration") + }) +} + +func TestToolResultNotDuplicated(t *testing.T) { + t.Run("SecondModelCallHasNoToolResultDuplication", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "mytool"} + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling", []schema.ToolCall{ + {ID: "c1", Function: schema.FunctionCall{Name: "mytool", Arguments: "{}"}}, + }), nil).Times(1) + + var capturedMsgs []*schema.Message + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + capturedMsgs = append([]*schema.Message{}, msgs...) + return schema.AssistantMessage("final", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are helpful.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hello")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.NotNil(t, capturedMsgs) + assert.Equal(t, 4, len(capturedMsgs), + "expected [system, user, assistant, tool_result], got %d messages", len(capturedMsgs)) + assert.Equal(t, schema.System, capturedMsgs[0].Role) + assert.Equal(t, schema.User, capturedMsgs[1].Role) + assert.Equal(t, schema.Assistant, capturedMsgs[2].Role) + assert.Equal(t, schema.Tool, capturedMsgs[3].Role) + + toolResultCount := 0 + for _, msg := range capturedMsgs { + if msg.Role == schema.Tool { + toolResultCount++ + } + } + assert.Equal(t, 1, toolResultCount, + "tool result should appear exactly once, got %d", toolResultCount) + }) + + t.Run("HookInjectedMessagePresentWithoutDuplication", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + tool1 := &namedTool{name: "mytool"} + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("calling", []schema.ToolCall{ + {ID: "c1", Function: schema.FunctionCall{Name: "mytool", Arguments: "{}"}}, + }), nil).Times(1) + + var capturedMsgs []*schema.Message + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + capturedMsgs = append([]*schema.Message{}, msgs...) + return schema.AssistantMessage("final", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are helpful.", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + }, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("hello")}}, + WithAfterToolCallsHook(func(ctx context.Context) error { + return compose.ProcessState(ctx, func(_ context.Context, st *State) error { + st.Messages = append(st.Messages, schema.UserMessage("injected")) + return nil + }) + })) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.NotNil(t, capturedMsgs) + assert.Equal(t, 5, len(capturedMsgs), + "expected [system, user, assistant, tool_result, injected], got %d messages", len(capturedMsgs)) + assert.Equal(t, schema.System, capturedMsgs[0].Role) + assert.Equal(t, schema.User, capturedMsgs[1].Role) + assert.Equal(t, schema.Assistant, capturedMsgs[2].Role) + assert.Equal(t, schema.Tool, capturedMsgs[3].Role) + assert.Equal(t, "injected", capturedMsgs[4].Content) + + toolResultCount := 0 + for _, msg := range capturedMsgs { + if msg.Role == schema.Tool { + toolResultCount++ + } + } + assert.Equal(t, 1, toolResultCount, + "tool result should appear exactly once, got %d", toolResultCount) + }) +} + +type testAfterAgentHandler struct { + *BaseChatModelAgentMiddleware + fn func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) +} + +func (h *testAfterAgentHandler) AfterAgent(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + return h.fn(ctx, state) +} + +type testAgenticAfterAgentHandler struct { + *TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage] + fn func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage]) (context.Context, error) +} + +func (h *testAgenticAfterAgentHandler) AfterAgent(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage]) (context.Context, error) { + return h.fn(ctx, state) +} + +func TestAfterAgent(t *testing.T) { + t.Run("FinalAnswer", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + var called bool + var capturedState *ChatModelAgentState + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + called = true + capturedState = state + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, called, "AfterAgent should be called on final answer") + assert.NotNil(t, capturedState) + assert.GreaterOrEqual(t, len(capturedState.Messages), 2, "state should contain at least user + assistant messages") + }) + + t.Run("ReturnDirectly", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + myTool := &namedTool{name: "myTool"} + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: "myTool", Arguments: "{}"}}, + }), nil).Times(1) + + var called bool + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{myTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testToolsFuncHandler{fn: func(ctx context.Context, tools []tool.BaseTool, returnDirectly map[string]bool) (context.Context, []tool.BaseTool, map[string]bool, error) { + returnDirectly["myTool"] = true + return ctx, tools, returnDirectly, nil + }}, + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + called = true + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, called, "AfterAgent should be called on return-directly tool result") + }) + + t.Run("NotCalledOnModelError", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("model error")).Times(1) + + var called bool + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + called = true + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.False(t, called, "AfterAgent should NOT be called when model errors") + }) + + t.Run("NotCalledOnMaxIterations", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + myTool := &namedTool{name: "myTool"} + + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Using tool", []schema.ToolCall{ + {ID: "call1", Function: schema.FunctionCall{Name: "myTool", Arguments: "{}"}}, + }), nil).AnyTimes() + + var called bool + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + MaxIterations: 1, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{myTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + called = true + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.False(t, called, "AfterAgent should NOT be called on max iterations exceeded") + }) + + t.Run("ErrorStopsRun", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + return ctx, fmt.Errorf("after agent hook error") + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + var gotErr error + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + gotErr = event.Err + } + } + + assert.Error(t, gotErr) + assert.Contains(t, gotErr.Error(), "AfterAgent") + }) + + t.Run("ContextPropagation", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + type ctxKey string + const key1 ctxKey = "afterAgentKey" + + var handler2ReceivedValue interface{} + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + return context.WithValue(ctx, key1, "afterValue"), nil + }}, + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + handler2ReceivedValue = ctx.Value(key1) + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.Equal(t, "afterValue", handler2ReceivedValue, + "Handler 2 should receive context value set by Handler 1 during AfterAgent") + }) + + t.Run("NoToolsPath", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + var called bool + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + called = true + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, called, "AfterAgent should be called on no-tools path") + }) + + t.Run("FailFast", func(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil).Times(1) + + var handler2Called bool + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: cm, + Handlers: []ChatModelAgentMiddleware{ + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + return ctx, fmt.Errorf("first handler error") + }}, + &testAfterAgentHandler{fn: func(ctx context.Context, state *ChatModelAgentState) (context.Context, error) { + handler2Called = true + return ctx, nil + }}, + }, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{Messages: []Message{schema.UserMessage("test")}}) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.False(t, handler2Called, "Handler 2 should NOT be called when Handler 1 errors (fail-fast)") + }) + + t.Run("AgenticFinalAnswer", func(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "agentic response"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + var called bool + var capturedState *TypedChatModelAgentState[*schema.AgenticMessage] + + handler := &testAgenticAfterAgentHandler{fn: func(ctx context.Context, state *TypedChatModelAgentState[*schema.AgenticMessage]) (context.Context, error) { + called = true + capturedState = state + return ctx, nil + }} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticTestAgent", + Description: "test", + Model: m, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&namedTool{name: "dummyTool"}}, + }, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{handler}, + }) + assert.NoError(t, err) + + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("test")}, + }) + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.True(t, called, "AfterAgent should be called on agentic final answer") + assert.NotNil(t, capturedState) + assert.GreaterOrEqual(t, len(capturedState.Messages), 2, "state should contain at least user + assistant messages") + }) +} diff --git a/adk/instruction.go b/adk/instruction.go new file mode 100644 index 0000000..635b65c --- /dev/null +++ b/adk/instruction.go @@ -0,0 +1,64 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino/adk/internal" +) + +const ( + TransferToAgentInstruction = `Available other agents: %s + +Decision rule: +- If you're best suited for the question according to your description: ANSWER +- If another agent is better according its description: CALL '%s' function with their agent name + +When transferring: OUTPUT ONLY THE FUNCTION CALL` + + TransferToAgentInstructionChinese = `可用的其他 agent：%s + +决策规则： +- 如果根据你的职责描述，你最适合回答这个问题：ANSWER +- 如果根据其职责描述，另一个 agent 更适合：调用 %s 函数，并传入该 agent 的名称 + +当进行移交时：只输出函数调用，不要输出其他任何内容` + + agentDescriptionTpl = "\n- Agent name: %s\n Agent description: %s" + agentDescriptionTplChinese = "\n- Agent 名字: %s\n Agent 描述: %s" +) + +func genTransferToAgentInstruction[M MessageType](ctx context.Context, agents []TypedAgent[M]) string { + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: agentDescriptionTpl, + Chinese: agentDescriptionTplChinese, + }) + instruction := internal.SelectPrompt(internal.I18nPrompts{ + English: TransferToAgentInstruction, + Chinese: TransferToAgentInstructionChinese, + }) + + var sb strings.Builder + for _, agent := range agents { + sb.WriteString(fmt.Sprintf(tpl, agent.Name(ctx), agent.Description(ctx))) + } + + return fmt.Sprintf(instruction, sb.String(), TransferToAgentToolName) +} diff --git a/adk/interface.go b/adk/interface.go new file mode 100644 index 0000000..8905950 --- /dev/null +++ b/adk/interface.go @@ -0,0 +1,519 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "fmt" + "io" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/schema" +) + +// ComponentOfAgent is the component type identifier for ADK agents in callbacks. +// Use this to filter callback events to only agent-related events. +const ComponentOfAgent components.Component = "Agent" + +// ComponentOfAgenticAgent is the component type identifier for ADK agents +// that use *schema.AgenticMessage in callbacks. +const ComponentOfAgenticAgent components.Component = "AgenticAgent" + +// MessageType is the sealed type constraint for message types used in ADK. +// Only *schema.Message and *schema.AgenticMessage satisfy this constraint. +// External packages cannot add new types to this union; all generic functions +// in ADK use exhaustive type switches on these two types. +type MessageType interface { + *schema.Message | *schema.AgenticMessage +} + +type Message = *schema.Message +type MessageStream = *schema.StreamReader[Message] + +type AgenticMessage = *schema.AgenticMessage +type AgenticMessageStream = *schema.StreamReader[AgenticMessage] + +// isNilMessage checks whether a generic message value is nil. +// Direct `msg == nil` does not compile for generic pointer types in Go; +// the canonical workaround is to compare through the `any` interface. +func isNilMessage[M MessageType](msg M) bool { + var zero M + return any(msg) == any(zero) +} + +// TypedMessageVariant represents a message output from an agent event. +// It carries either a complete message or a streaming reader, along with +// metadata describing the event's origin. +// +// Role and ToolName are only meaningful for *schema.Message events. For +// *schema.AgenticMessage events (created via EventFromAgenticMessage), these +// fields are always zero-valued because AgenticMessage carries tool results as +// ContentBlocks within the message itself and does not support agent transfer. +// +// For *schema.Message events, Role and ToolName exist independently of the inner +// Message because in streaming mode (IsStreaming=true, Message=nil), the message +// has not materialized yet and the consumer needs metadata without consuming the stream. +type TypedMessageVariant[M MessageType] struct { + IsStreaming bool + + Message M + MessageStream *schema.StreamReader[M] + + // Role indicates the origin of this event within the agent's ReAct loop. + // Only meaningful for *schema.Message events: + // - schema.Assistant: the event carries model output (generation or stream). + // - schema.Tool: the event carries a tool execution result. + // Always zero-valued for *schema.AgenticMessage events; use AgenticRole instead. + Role schema.RoleType + + // AgenticRole indicates the role of the agentic message (assistant, user, system). + // Only meaningful for *schema.AgenticMessage events. + // In streaming mode, this is available before consuming the stream. + // Always zero-valued for *schema.Message events; use Role instead. + AgenticRole schema.AgenticRoleType + + // ToolName is the name of the tool that produced this event. + // Only meaningful for *schema.Message events: non-empty when Role == schema.Tool. + // In streaming mode, this is the only way to identify the source tool before + // the stream is consumed. + // Always empty for *schema.AgenticMessage events. + ToolName string +} + +func (mv *TypedMessageVariant[M]) GetMessage() (M, error) { + if mv.IsStreaming { + return concatMessageStream(mv.MessageStream) + } + return mv.Message, nil +} + +type MessageVariant = TypedMessageVariant[*schema.Message] + +type messageVariantSerialization struct { + IsStreaming bool + Message Message + MessageStream Message + Role schema.RoleType + ToolName string +} + +type agenticMessageVariantSerialization struct { + IsStreaming bool + Message *schema.AgenticMessage + MessageStream *schema.AgenticMessage + Role schema.RoleType + AgenticRole schema.AgenticRoleType + ToolName string +} + +func (mv *TypedMessageVariant[M]) GobEncode() ([]byte, error) { + if mvMsg, ok := any(mv).(*TypedMessageVariant[*schema.Message]); ok { + return gobEncodeMessageVariant(mvMsg) + } + if mvAgentic, ok := any(mv).(*TypedMessageVariant[*schema.AgenticMessage]); ok { + return gobEncodeAgenticMessageVariant(mvAgentic) + } + return nil, fmt.Errorf("gob encoding not supported for this message type") +} + +func (mv *TypedMessageVariant[M]) GobDecode(b []byte) error { + if mvMsg, ok := any(mv).(*TypedMessageVariant[*schema.Message]); ok { + return gobDecodeMessageVariant(mvMsg, b) + } + if mvAgentic, ok := any(mv).(*TypedMessageVariant[*schema.AgenticMessage]); ok { + return gobDecodeAgenticMessageVariant(mvAgentic, b) + } + return fmt.Errorf("gob decoding not supported for this message type") +} + +func gobEncodeMessageVariant(mv *TypedMessageVariant[*schema.Message]) ([]byte, error) { + s := &messageVariantSerialization{ + IsStreaming: mv.IsStreaming, + Message: mv.Message, + Role: mv.Role, + ToolName: mv.ToolName, + } + if mv.IsStreaming { + var messages []Message + for { + frame, err := mv.MessageStream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, fmt.Errorf("error receiving message stream: %w", err) + } + messages = append(messages, frame) + } + m, err := schema.ConcatMessages(messages) + if err != nil { + return nil, fmt.Errorf("failed to encode message: cannot concat message stream: %w", err) + } + s.MessageStream = m + } + buf := &bytes.Buffer{} + err := gob.NewEncoder(buf).Encode(s) + if err != nil { + return nil, fmt.Errorf("failed to gob encode message variant: %w", err) + } + return buf.Bytes(), nil +} + +func gobDecodeMessageVariant(mv *TypedMessageVariant[*schema.Message], b []byte) error { + s := &messageVariantSerialization{} + err := gob.NewDecoder(bytes.NewReader(b)).Decode(s) + if err != nil { + return fmt.Errorf("failed to decoding message variant: %w", err) + } + mv.IsStreaming = s.IsStreaming + mv.Message = s.Message + mv.Role = s.Role + mv.ToolName = s.ToolName + if s.MessageStream != nil { + mv.MessageStream = schema.StreamReaderFromArray([]*schema.Message{s.MessageStream}) + } + return nil +} + +func gobEncodeAgenticMessageVariant(mv *TypedMessageVariant[*schema.AgenticMessage]) ([]byte, error) { + s := &agenticMessageVariantSerialization{ + IsStreaming: mv.IsStreaming, + Message: mv.Message, + Role: mv.Role, + AgenticRole: mv.AgenticRole, + ToolName: mv.ToolName, + } + if mv.IsStreaming { + var messages []*schema.AgenticMessage + for { + frame, err := mv.MessageStream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, fmt.Errorf("error receiving agentic message stream: %w", err) + } + messages = append(messages, frame) + } + m, err := schema.ConcatAgenticMessages(messages) + if err != nil { + return nil, fmt.Errorf("failed to encode agentic message: cannot concat message stream: %w", err) + } + s.MessageStream = m + } + buf := &bytes.Buffer{} + err := gob.NewEncoder(buf).Encode(s) + if err != nil { + return nil, fmt.Errorf("failed to gob encode agentic message variant: %w", err) + } + return buf.Bytes(), nil +} + +func gobDecodeAgenticMessageVariant(mv *TypedMessageVariant[*schema.AgenticMessage], b []byte) error { + s := &agenticMessageVariantSerialization{} + err := gob.NewDecoder(bytes.NewReader(b)).Decode(s) + if err != nil { + return fmt.Errorf("failed to decode agentic message variant: %w", err) + } + mv.IsStreaming = s.IsStreaming + mv.Message = s.Message + mv.Role = s.Role + mv.AgenticRole = s.AgenticRole + mv.ToolName = s.ToolName + if s.MessageStream != nil { + mv.MessageStream = schema.StreamReaderFromArray([]*schema.AgenticMessage{s.MessageStream}) + } + return nil +} + +// typedEventFromMessage creates a TypedAgentEvent containing the given message and optional stream. +func typedEventFromMessage[M MessageType](msg M, msgStream *schema.StreamReader[M], + role schema.RoleType, toolName string) *TypedAgentEvent[M] { + return &TypedAgentEvent[M]{ + Output: &TypedAgentOutput[M]{ + MessageOutput: &TypedMessageVariant[M]{ + IsStreaming: msgStream != nil, + Message: msg, + MessageStream: msgStream, + Role: role, + ToolName: toolName, + }, + }, + } +} + +// typedModelOutputEvent creates a model-output event for the generic path. +// For *schema.Message, Role is set to schema.Assistant. +// For *schema.AgenticMessage, AgenticRole is set to schema.AgenticRoleTypeAssistant. +func typedModelOutputEvent[M MessageType](msg M, msgStream *schema.StreamReader[M]) *TypedAgentEvent[M] { + var role schema.RoleType + var agenticRole schema.AgenticRoleType + var zero M + if _, ok := any(zero).(*schema.Message); ok { + role = schema.Assistant + } else { + agenticRole = schema.AgenticRoleTypeAssistant + } + event := typedEventFromMessage(msg, msgStream, role, "") + event.Output.MessageOutput.AgenticRole = agenticRole + return event +} + +// EventFromMessage creates an AgentEvent containing the given message and optional stream. +// +// role identifies the origin of this event: +// - schema.Assistant: model output (generation or stream). +// - schema.Tool: tool execution result; toolName must be non-empty. +// +// For *schema.AgenticMessage events, use EventFromAgenticMessage instead. +func EventFromMessage(msg Message, msgStream *schema.StreamReader[Message], + role schema.RoleType, toolName string) *AgentEvent { + return typedEventFromMessage(msg, msgStream, role, toolName) +} + +// EventFromAgenticMessage creates a TypedAgentEvent for the AgenticMessage path. +// Unlike EventFromMessage, it does not require role or toolName parameters because +// AgenticMessage carries tool results as ContentBlocks within the message itself, +// and does not support agent transfer. +// +// agenticRole identifies the role of the message (e.g. schema.AgenticRoleTypeAssistant). +// In streaming mode, the role is available on the event before consuming the stream. +func EventFromAgenticMessage(msg AgenticMessage, msgStream AgenticMessageStream, agenticRole schema.AgenticRoleType) *TypedAgentEvent[AgenticMessage] { + return &TypedAgentEvent[AgenticMessage]{ + Output: &TypedAgentOutput[AgenticMessage]{ + MessageOutput: &TypedMessageVariant[AgenticMessage]{ + IsStreaming: msgStream != nil, + Message: msg, + MessageStream: msgStream, + AgenticRole: agenticRole, + }, + }, + } +} + +// TransferToAgentAction represents a transfer-to-agent action. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +type TransferToAgentAction struct { + DestAgentName string +} + +type TypedAgentOutput[M MessageType] struct { + MessageOutput *TypedMessageVariant[M] + + CustomizedOutput any +} + +type AgentOutput = TypedAgentOutput[*schema.Message] + +// NewTransferToAgentAction creates an action to transfer to the specified agent. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func NewTransferToAgentAction(destAgentName string) *AgentAction { + return &AgentAction{TransferToAgent: &TransferToAgentAction{DestAgentName: destAgentName}} +} + +// NewExitAction creates an action that signals the agent to exit. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func NewExitAction() *AgentAction { + return &AgentAction{Exit: true} +} + +// AgentAction represents actions that an agent can emit during execution. +// +// Action Scoping in Agent Tools: +// When an agent is wrapped as an agent tool (via NewAgentTool), actions emitted by the inner agent +// are scoped to the tool boundary: +// - Interrupted: Propagated via CompositeInterrupt to allow proper interrupt/resume across boundaries +// - Exit, TransferToAgent, BreakLoop: Ignored outside the agent tool; these actions only affect +// the inner agent's execution and do not propagate to the parent agent +// +// This scoping ensures that nested agents cannot unexpectedly terminate or transfer control +// of their parent agent's execution flow. +type AgentAction struct { + Exit bool + + Interrupted *InterruptInfo + + TransferToAgent *TransferToAgentAction + + BreakLoop *BreakLoopAction + + CustomizedAction any + + internalInterrupted *core.InterruptSignal +} + +// RunStep represents a step in the agent execution path. +// CheckpointSchema: persisted via serialization.RunCtx (gob). +// +// NOT RECOMMENDED: RunStep is mainly relevant for agent transfer and workflow agents, +// which have not proven to be more effective empirically. Consider using ChatModelAgent +// with AgentTool or DeepAgent instead for most multi-agent scenarios. +type RunStep struct { + agentName string +} + +func init() { + schema.RegisterName[[]RunStep]("eino_run_step_list") +} + +func (r *RunStep) String() string { + return r.agentName +} + +func (r *RunStep) Equals(r1 RunStep) bool { + return r.agentName == r1.agentName +} + +func (r *RunStep) GobEncode() ([]byte, error) { + s := &runStepSerialization{AgentName: r.agentName} + buf := &bytes.Buffer{} + err := gob.NewEncoder(buf).Encode(s) + if err != nil { + return nil, fmt.Errorf("failed to gob encode RunStep: %w", err) + } + return buf.Bytes(), nil +} + +func (r *RunStep) GobDecode(b []byte) error { + s := &runStepSerialization{} + err := gob.NewDecoder(bytes.NewReader(b)).Decode(s) + if err != nil { + return fmt.Errorf("failed to gob decode RunStep: %w", err) + } + r.agentName = s.AgentName + return nil +} + +type runStepSerialization struct { + AgentName string +} + +// TypedAgentEvent represents a single event emitted during agent execution. +// CheckpointSchema: persisted via serialization.RunCtx (gob). +type TypedAgentEvent[M MessageType] struct { + AgentName string + + // RunPath represents the execution path from root agent to the current event source. + // This field is managed entirely by the framework and cannot be set by end-users. + // + // NOT RECOMMENDED: RunPath is mainly relevant for agent transfer and workflow agents, + // which have not proven to be more effective empirically. For ChatModelAgent with + // AgentTool or DeepAgent, RunPath is trivial. Consider those patterns instead. + RunPath []RunStep + + Output *TypedAgentOutput[M] + + Action *AgentAction + + Err error +} + +// AgentEvent is the default event type using *schema.Message. +type AgentEvent = TypedAgentEvent[*schema.Message] + +type TypedAgentInput[M MessageType] struct { + Messages []M + EnableStreaming bool +} + +type AgentInput = TypedAgentInput[*schema.Message] + +// TypedAgent is the base agent interface parameterized by message type. +// +// For M = *schema.Message, the full ADK feature set is supported (multi-agent +// orchestration, cancel monitoring, retry, flowAgent). +// For M = *schema.AgenticMessage, single-agent execution works but cancel +// monitoring on the model stream and retry are not yet wired. +type TypedAgent[M MessageType] interface { + Name(ctx context.Context) string + Description(ctx context.Context) string + + // Run runs the agent. + // The returned AgentEvent within the AsyncIterator must be safe to modify. + // If the returned AgentEvent within the AsyncIterator contains MessageStream, + // the MessageStream MUST be exclusive and safe to be received directly. + // NOTE: it's recommended to use SetAutomaticClose() on the MessageStream of AgentEvents emitted by AsyncIterator, + // so that even the events are not processed, the MessageStream can still be closed. + Run(ctx context.Context, input *TypedAgentInput[M], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] +} + +//go:generate mockgen -destination ../internal/mock/adk/Agent_mock.go --package adk github.com/cloudwego/eino/adk Agent,ResumableAgent +type Agent = TypedAgent[*schema.Message] + +// OnSubAgents is the interface for agents that support sub-agent registration and transfer. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +type OnSubAgents interface { + OnSetSubAgents(ctx context.Context, subAgents []Agent) error + OnSetAsSubAgent(ctx context.Context, parent Agent) error + + OnDisallowTransferToParent(ctx context.Context) error +} + +type TypedResumableAgent[M MessageType] interface { + TypedAgent[M] + + Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] +} + +type ResumableAgent = TypedResumableAgent[*schema.Message] + +func concatMessageStream[M MessageType](stream *schema.StreamReader[M]) (M, error) { + var zero M + switch s := any(stream).(type) { + case *schema.StreamReader[*schema.Message]: + result, err := schema.ConcatMessageStream(s) + if err != nil { + return zero, err + } + return any(result).(M), nil + case *schema.StreamReader[*schema.AgenticMessage]: + defer s.Close() + var msgs []*schema.AgenticMessage + for { + frame, err := s.Recv() + if err == io.EOF { + break + } + if err != nil { + return zero, err + } + msgs = append(msgs, frame) + } + result, err := schema.ConcatAgenticMessages(msgs) + if err != nil { + return zero, err + } + return any(result).(M), nil + default: + panic("unreachable: unknown MessageType") + } +} diff --git a/adk/internal/config.go b/adk/internal/config.go new file mode 100644 index 0000000..9232d87 --- /dev/null +++ b/adk/internal/config.go @@ -0,0 +1,76 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package internal provides adk internal utils. +package internal + +import ( + "fmt" + "sync/atomic" +) + +// Language represents the language setting for the ADK built-in prompts. +type Language uint8 + +const ( + // LanguageEnglish represents English language. + LanguageEnglish Language = iota + // LanguageChinese represents Chinese language. + LanguageChinese +) + +var language atomic.Value + +// SetLanguage sets the language for the ADK built-in prompts. +// The default language is English if not explicitly set. +func SetLanguage(lang Language) error { + if lang != LanguageEnglish && + lang != LanguageChinese { + return fmt.Errorf("invalid language: %v", lang) + } + language.Store(lang) + return nil +} + +// GetLanguage returns the current language setting for the ADK built-in prompts. +// Returns LanguageEnglish if no language has been set. +func getLanguage() Language { + if l, ok := language.Load().(Language); ok { + return l + } + return LanguageEnglish +} + +// I18nPrompts holds prompt strings for different languages. +type I18nPrompts struct { + English string + Chinese string +} + +// SelectPrompt returns the appropriate prompt string based on the current language setting. +// Returns an error if the current language is not supported. +func SelectPrompt(prompts I18nPrompts) string { + lang := getLanguage() + switch lang { + case LanguageEnglish: + return prompts.English + case LanguageChinese: + return prompts.Chinese + default: + // unreachable + panic(fmt.Sprintf("invalid language: %v", lang)) + } +} diff --git a/adk/internal/message_id.go b/adk/internal/message_id.go new file mode 100644 index 0000000..0c7ada1 --- /dev/null +++ b/adk/internal/message_id.go @@ -0,0 +1,58 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import "github.com/google/uuid" + +// EinoMsgIDKey is the Extra key used to store the eino-internal message ID. +const EinoMsgIDKey = "_eino_msg_id" + +// GetMessageID returns the message ID from Extra, or "" if not set. +// Works with any map[string]any (Message.Extra or AgenticMessage.Extra). +func GetMessageID(extra map[string]any) string { + if extra == nil { + return "" + } + id, _ := extra[EinoMsgIDKey].(string) + return id +} + +// SetMessageID sets the message ID in Extra and returns the resulting map. +// +// Copy-on-write: the input map is never mutated, because a message's Extra can +// be shared across fan-out stream readers and an in-place write would race with +// concurrent reads (fatal "concurrent map read and map write"). Reassigning the +// returned map to a shared Extra field is still a field-level race that cannot +// be eliminated while Extra is exported; COW only removes the map-level panic. +func SetMessageID(extra map[string]any, id string) map[string]any { + next := make(map[string]any, len(extra)+1) + for k, v := range extra { + next[k] = v + } + next[EinoMsgIDKey] = id + return next +} + +// EnsureMessageID assigns a UUID v4 if no message ID is present. +// Idempotent: if ID already set, no-op. +// Returns the (possibly newly created) Extra map. +func EnsureMessageID(extra map[string]any) map[string]any { + if GetMessageID(extra) != "" { + return extra + } + return SetMessageID(extra, uuid.NewString()) +} diff --git a/adk/internal/message_id_test.go b/adk/internal/message_id_test.go new file mode 100644 index 0000000..f7c536f --- /dev/null +++ b/adk/internal/message_id_test.go @@ -0,0 +1,87 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestGetMessageID(t *testing.T) { + t.Run("nil extra returns empty", func(t *testing.T) { + assert.Equal(t, "", GetMessageID(nil)) + }) + + t.Run("empty extra returns empty", func(t *testing.T) { + assert.Equal(t, "", GetMessageID(map[string]any{})) + }) + + t.Run("wrong type returns empty", func(t *testing.T) { + extra := map[string]any{EinoMsgIDKey: 123} + assert.Equal(t, "", GetMessageID(extra)) + }) + + t.Run("returns set ID", func(t *testing.T) { + extra := map[string]any{EinoMsgIDKey: "test-id-123"} + assert.Equal(t, "test-id-123", GetMessageID(extra)) + }) +} + +func TestSetMessageID(t *testing.T) { + t.Run("nil extra creates map", func(t *testing.T) { + extra := SetMessageID(nil, "id-1") + assert.NotNil(t, extra) + assert.Equal(t, "id-1", extra[EinoMsgIDKey]) + }) + + t.Run("existing extra preserved", func(t *testing.T) { + extra := map[string]any{"other_key": "other_val"} + result := SetMessageID(extra, "id-2") + assert.Equal(t, "id-2", result[EinoMsgIDKey]) + assert.Equal(t, "other_val", result["other_key"]) + }) +} + +func TestEnsureMessageID(t *testing.T) { + t.Run("nil extra gets ID", func(t *testing.T) { + extra := EnsureMessageID(nil) + id := GetMessageID(extra) + assert.NotEmpty(t, id) + assert.Len(t, id, 36) // UUID v4 format: 8-4-4-4-12 = 36 chars + }) + + t.Run("idempotent - does not overwrite existing ID", func(t *testing.T) { + extra := SetMessageID(nil, "existing-id") + result := EnsureMessageID(extra) + assert.Equal(t, "existing-id", GetMessageID(result)) + }) + + t.Run("empty extra gets new ID", func(t *testing.T) { + extra := map[string]any{} + result := EnsureMessageID(extra) + id := GetMessageID(result) + assert.NotEmpty(t, id) + assert.Len(t, id, 36) + }) + + t.Run("generates unique IDs", func(t *testing.T) { + extra1 := EnsureMessageID(nil) + extra2 := EnsureMessageID(nil) + assert.NotEqual(t, GetMessageID(extra1), GetMessageID(extra2)) + }) +} diff --git a/adk/interrupt.go b/adk/interrupt.go new file mode 100644 index 0000000..3d31054 --- /dev/null +++ b/adk/interrupt.go @@ -0,0 +1,412 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "fmt" + "sync" + + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/schema" +) + +// ResumeInfo holds all the information necessary to resume an interrupted agent execution. +// It is created by the framework and passed to an agent's Resume method. +type ResumeInfo struct { + // EnableStreaming indicates whether the original execution was in streaming mode. + EnableStreaming bool + + // Deprecated: use InterruptContexts from the embedded InterruptInfo for user-facing details, + // and GetInterruptState for internal state retrieval. + *InterruptInfo + + WasInterrupted bool + InterruptState any + IsResumeTarget bool + ResumeData any +} + +// InterruptInfo contains all the information about an interruption event. +// It is created by the framework when an agent returns an interrupt action. +type InterruptInfo struct { + Data any + + // InterruptContexts provides a structured, user-facing view of the interrupt chain. + // Each context represents a step in the agent hierarchy that was interrupted. + InterruptContexts []*InterruptCtx +} + +// TypedInterrupt creates a typed interrupt event that pauses execution to request external input. +// It is the generic counterpart of Interrupt; see Interrupt for full documentation. +func TypedInterrupt[M MessageType](ctx context.Context, info any) *TypedAgentEvent[M] { + var rp []RunStep + rCtx := getRunCtx(ctx) + if rCtx != nil { + rp = rCtx.RunPath + } + + is, err := core.Interrupt(ctx, info, nil, nil, + core.WithLayerPayload(rp)) + if err != nil { + return &TypedAgentEvent[M]{Err: err} + } + + contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes) + + return &TypedAgentEvent[M]{ + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + InterruptContexts: contexts, + }, + internalInterrupted: is, + }, + } +} + +// Interrupt creates a basic interrupt action. +// This is used when an agent needs to pause its execution to request external input or intervention, +// but does not need to save any internal state to be restored upon resumption. +// The `info` parameter is user-facing data that describes the reason for the interrupt. +func Interrupt(ctx context.Context, info any) *AgentEvent { + return TypedInterrupt[*schema.Message](ctx, info) +} + +// TypedStatefulInterrupt creates a typed interrupt event that also saves the agent's internal state. +// It is the generic counterpart of StatefulInterrupt; see StatefulInterrupt for full documentation. +func TypedStatefulInterrupt[M MessageType](ctx context.Context, info any, state any) *TypedAgentEvent[M] { + var rp []RunStep + rCtx := getRunCtx(ctx) + if rCtx != nil { + rp = rCtx.RunPath + } + + is, err := core.Interrupt(ctx, info, state, nil, + core.WithLayerPayload(rp)) + if err != nil { + return &TypedAgentEvent[M]{Err: err} + } + + contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes) + + return &TypedAgentEvent[M]{ + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + InterruptContexts: contexts, + }, + internalInterrupted: is, + }, + } +} + +// StatefulInterrupt creates an interrupt action that also saves the agent's internal state. +// This is used when an agent has internal state that must be restored for it to continue correctly. +// The `info` parameter is user-facing data describing the interrupt. +// The `state` parameter is the agent's internal state object, which will be serialized and stored. +func StatefulInterrupt(ctx context.Context, info any, state any) *AgentEvent { + return TypedStatefulInterrupt[*schema.Message](ctx, info, state) +} + +// TypedCompositeInterrupt creates a typed interrupt event that aggregates sub-interrupt signals. +// It is the generic counterpart of CompositeInterrupt; see CompositeInterrupt for full documentation. +func TypedCompositeInterrupt[M MessageType](ctx context.Context, info any, state any, + subInterruptSignals ...*InterruptSignal) *TypedAgentEvent[M] { + var rp []RunStep + rCtx := getRunCtx(ctx) + if rCtx != nil { + rp = rCtx.RunPath + } + + is, err := core.Interrupt(ctx, info, state, subInterruptSignals, + core.WithLayerPayload(rp)) + if err != nil { + return &TypedAgentEvent[M]{Err: err} + } + + contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes) + + return &TypedAgentEvent[M]{ + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + InterruptContexts: contexts, + }, + internalInterrupted: is, + }, + } +} + +// CompositeInterrupt creates an interrupt event that aggregates sub-interrupt signals. +func CompositeInterrupt(ctx context.Context, info any, state any, + subInterruptSignals ...*InterruptSignal) *AgentEvent { + return TypedCompositeInterrupt[*schema.Message](ctx, info, state, subInterruptSignals...) +} + +// Address represents the unique, hierarchical address of a component within an execution. +// It is a slice of AddressSegments, where each segment represents one level of nesting. +// This is a type alias for core.Address. See the core package for more details. +type Address = core.Address +type AddressSegment = core.AddressSegment +type AddressSegmentType = core.AddressSegmentType + +const ( + AddressSegmentAgent AddressSegmentType = "agent" + AddressSegmentTool AddressSegmentType = "tool" +) + +var allowedAddressSegmentTypes = []AddressSegmentType{AddressSegmentAgent, AddressSegmentTool} + +// AppendAddressSegment adds an address segment for the current execution context. +func AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string) context.Context { + return core.AppendAddressSegment(ctx, segType, segID, "") +} + +// InterruptCtx provides a structured, user-facing view of a single point of interruption. +// It contains the ID and Address of the interrupted component, as well as user-defined info. +// This is a type alias for core.InterruptCtx. See the core package for more details. +type InterruptCtx = core.InterruptCtx +type InterruptSignal = core.InterruptSignal + +// FromInterruptContexts converts user-facing interrupt contexts to an interrupt signal. +func FromInterruptContexts(contexts []*InterruptCtx) *InterruptSignal { + return core.FromInterruptContexts(contexts) +} + +// WithCheckPointID sets the checkpoint ID used for interruption persistence. +func WithCheckPointID(id string) AgentRunOption { + return WrapImplSpecificOptFn(func(t *options) { + t.checkPointID = &id + }) +} + +func init() { + schema.RegisterName[*serialization]("_eino_adk_serialization") + schema.RegisterName[*WorkflowInterruptInfo]("_eino_adk_workflow_interrupt_info") + // Register []byte for gob: the cancel refactor routes bridge store checkpoint + // bytes ([]byte) through InterruptState.State (type any) inside the outer + // serialization struct. Gob requires concrete types behind interface fields + // to be registered. + gob.Register([]byte{}) +} + +// serialization CheckpointSchema: root checkpoint payload (gob). +// Any type tagged with `CheckpointSchema:` is persisted and must remain backward compatible. +type serialization struct { + RunCtx *runContext + // deprecated: still keep it here for backward compatibility + Info *InterruptInfo + EnableStreaming bool + InterruptID2Address map[string]Address + InterruptID2State map[string]core.InterruptState +} + +func runnerLoadCheckPointImpl(store CheckPointStore, ctx context.Context, checkpointID string) ( + context.Context, *runContext, *ResumeInfo, error) { + data, existed, err := store.Get(ctx, checkpointID) + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to get checkpoint from store: %w", err) + } + if !existed { + return nil, nil, nil, fmt.Errorf("checkpoint[%s] not exist", checkpointID) + } + + data = preprocessADKCheckpoint(data) + + s := &serialization{} + err = gob.NewDecoder(bytes.NewReader(data)).Decode(s) + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to decode checkpoint: %w", err) + } + ctx = core.PopulateInterruptState(ctx, s.InterruptID2Address, s.InterruptID2State) + + return ctx, s.RunCtx, &ResumeInfo{ + EnableStreaming: s.EnableStreaming, + InterruptInfo: s.Info, + }, nil +} + +// preprocessADKCheckpoint fixes a gob incompatibility when resuming old ChatModelAgent/DeepAgents checkpoints. +// +// Background +// - ADK checkpoints are gob-encoded. +// - Some values inside checkpoints are stored as `any`, so gob includes a concrete type name +// string in the wire format and uses that name to pick the local Go type to decode into. +// +// Problem (v0.8.0-v0.8.3 checkpoints) +// - In v0.8.0-v0.8.3, *State was registered under the name "_eino_adk_react_state" AND +// implemented GobEncode/GobDecode, so the wire format for that name is "GobEncoder payload" +// (opaque bytes). +// - In v0.7.*, the same name "_eino_adk_react_state" was used but encoded as a normal struct +// (no GobEncode). Gob treats these two wire formats as incompatible. +// - Gob only allows one local Go type per name. Today we register "_eino_adk_react_state" to +// a v0.7-compatible struct decoder (stateV07). If we try to decode a v0.8.0-v0.8.3 +// checkpoint under that same name, gob fails with a "want struct; got non-struct" mismatch. +// +// Solution +// - We keep "_eino_adk_react_state" mapped to the v0.7 decoder. +// - For v0.8.0-v0.8.3 checkpoints only, we rewrite the on-wire name to a same-length alias +// "_eino_adk_state_v080_", which is registered to a GobDecoder-compatible type (stateV080). +// - The alias is the same length as the original, so we can safely replace the length-prefixed +// bytes without re-encoding the whole stream. +func preprocessADKCheckpoint(data []byte) []byte { + const ( + lenPrefixedReactStateName = "\x15" + stateGobNameV07 + lenPrefixedCompatName = "\x15" + stateGobNameV080 + lenPrefixedStateSerializationName = "\x12stateSerialization" + ) + + // the following line checks whether the checkpoint is persisted through v0.8.0-v0.8.3 + if !bytes.Contains(data, []byte(lenPrefixedReactStateName)) || !bytes.Contains(data, []byte(lenPrefixedStateSerializationName)) { + return data + } + return bytes.ReplaceAll(data, + []byte(lenPrefixedReactStateName), + []byte(lenPrefixedCompatName)) +} + +func runnerSaveCheckPointImpl( + enableStreaming bool, + store CheckPointStore, + ctx context.Context, + key string, + info *InterruptInfo, + is *core.InterruptSignal, +) error { + if store == nil { + return nil + } + + runCtx := getRunCtx(ctx) + + id2Addr, id2State := core.SignalToPersistenceMaps(is) + + buf := &bytes.Buffer{} + err := gob.NewEncoder(buf).Encode(&serialization{ + RunCtx: runCtx, + Info: info, + InterruptID2Address: id2Addr, + InterruptID2State: id2State, + EnableStreaming: enableStreaming, + }) + if err != nil { + return fmt.Errorf("failed to encode checkpoint: %w", err) + } + return store.Set(ctx, key, buf.Bytes()) +} + +const bridgeCheckpointID = "adk_react_mock_key" + +func newBridgeStore() *bridgeStore { + return &bridgeStore{data: make(map[string][]byte)} +} + +func newResumeBridgeStore(checkPointID string, data []byte) *bridgeStore { + return &bridgeStore{ + data: map[string][]byte{checkPointID: data}, + } +} + +type bridgeStore struct { + mu sync.Mutex + data map[string][]byte +} + +func (m *bridgeStore) Get(_ context.Context, key string) ([]byte, bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + if v, ok := m.data[key]; ok { + return v, true, nil + } + return nil, false, nil +} + +func (m *bridgeStore) Set(_ context.Context, key string, checkPoint []byte) error { + m.mu.Lock() + defer m.mu.Unlock() + if m.data == nil { + m.data = make(map[string][]byte) + } + m.data[key] = checkPoint + return nil +} + +func getNextResumeAgent(ctx context.Context, _ *ResumeInfo) (string, error) { + nextAgents, err := core.GetNextResumptionPoints(ctx) + if err != nil { + return "", fmt.Errorf("failed to get next agent leading to interruption: %w", err) + } + + if len(nextAgents) == 0 { + return "", errors.New("no child agents leading to interrupted agent were found") + } + + if len(nextAgents) > 1 { + return "", errors.New("agent has multiple child agents leading to interruption, " + + "but concurrent transfer is not supported") + } + + // get the single next agent to delegate to. + var nextAgentID string + for id := range nextAgents { + nextAgentID = id + break + } + + return nextAgentID, nil +} + +func getNextResumeAgents(ctx context.Context, _ *ResumeInfo) (map[string]bool, error) { + nextAgents, err := core.GetNextResumptionPoints(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get next agents leading to interruption: %w", err) + } + + if len(nextAgents) == 0 { + return nil, errors.New("no child agents leading to interrupted agent were found") + } + + return nextAgents, nil +} + +func buildResumeInfo(ctx context.Context, nextAgentID string, info *ResumeInfo) ( + context.Context, *ResumeInfo) { + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, nextAgentID) + nextResumeInfo := &ResumeInfo{ + EnableStreaming: info.EnableStreaming, + InterruptInfo: info.InterruptInfo, + } + + wasInterrupted, hasState, state := core.GetInterruptState[any](ctx) + nextResumeInfo.WasInterrupted = wasInterrupted + if hasState { + nextResumeInfo.InterruptState = state + } + + if wasInterrupted { + isResumeTarget, hasData, data := core.GetResumeContext[any](ctx) + nextResumeInfo.IsResumeTarget = isResumeTarget + if hasData { + nextResumeInfo.ResumeData = data + } + } + + ctx = updateRunPathOnly(ctx, nextAgentID) + + return ctx, nextResumeInfo +} diff --git a/adk/interrupt_test.go b/adk/interrupt_test.go new file mode 100644 index 0000000..480c0f8 --- /dev/null +++ b/adk/interrupt_test.go @@ -0,0 +1,2149 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type interruptTestToolsHandler struct { + *BaseChatModelAgentMiddleware + tools []tool.BaseTool +} + +func TestPreprocessADKCheckpoint(t *testing.T) { + t.Run("no-op when missing markers", func(t *testing.T) { + in := []byte("random") + out := preprocessADKCheckpoint(append([]byte(nil), in...)) + assert.Equal(t, in, out) + }) + + t.Run("rewrite legacy name for v0.8.0-v0.8.3", func(t *testing.T) { + const ( + lenPrefixedReactStateName = "\x15" + stateGobNameV07 + lenPrefixedCompatName = "\x15" + stateGobNameV080 + lenPrefixedStateSerializationName = "\x12stateSerialization" + ) + + in := []byte(lenPrefixedReactStateName + "xxx" + lenPrefixedStateSerializationName + "yyy") + out := preprocessADKCheckpoint(append([]byte(nil), in...)) + assert.True(t, bytes.Contains(out, []byte(lenPrefixedCompatName))) + assert.False(t, bytes.Contains(out, []byte(lenPrefixedReactStateName))) + }) +} + +func (h *interruptTestToolsHandler) BeforeAgent(ctx context.Context, runCtx *ChatModelAgentContext) (context.Context, *ChatModelAgentContext, error) { + runCtx.Tools = append(runCtx.Tools, h.tools...) + return ctx, runCtx, nil +} + +func TestSaveAgentEventWrapper(t *testing.T) { + sr, sw := schema.Pipe[Message](1) + sw.Send(schema.UserMessage("test"), nil) + sw.Close() + sr = sr.Copy(2)[1] + + w := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + RunPath: []RunStep{ + { + "a1", + }, + { + "a2", + }, + }, + }, + mu: sync.Mutex{}, + concatenatedMessage: nil, + } + + _, err := getMessageFromWrappedEvent(w) + assert.NoError(t, err) + + buf, err := w.GobEncode() + assert.NoError(t, err) + assert.NoError(t, err) + + w1 := &agentEventWrapper{} + err = w1.GobDecode(buf) + assert.NoError(t, err) +} + +func TestInterruptFunctionsPopulateInterruptContextsImmediately(t *testing.T) { + ctx := context.Background() + ctx, _ = initRunCtx(ctx, "TestAgent", &AgentInput{Messages: []Message{}}) + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, "TestAgent") + + t.Run("Interrupt populates InterruptContexts", func(t *testing.T) { + event := Interrupt(ctx, "test info") + assert.NotNil(t, event.Action) + assert.NotNil(t, event.Action.Interrupted) + assert.NotNil(t, event.Action.Interrupted.InterruptContexts) + assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts)) + assert.Equal(t, "test info", event.Action.Interrupted.InterruptContexts[0].Info) + assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause) + assert.Equal(t, Address{ + {Type: AddressSegmentAgent, ID: "TestAgent"}, + }, event.Action.Interrupted.InterruptContexts[0].Address) + }) + + t.Run("StatefulInterrupt populates InterruptContexts", func(t *testing.T) { + event := StatefulInterrupt(ctx, "stateful info", "my state") + assert.NotNil(t, event.Action) + assert.NotNil(t, event.Action.Interrupted) + assert.NotNil(t, event.Action.Interrupted.InterruptContexts) + assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts)) + assert.Equal(t, "stateful info", event.Action.Interrupted.InterruptContexts[0].Info) + assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause) + }) + + t.Run("CompositeInterrupt populates InterruptContexts with filtered parent chain", func(t *testing.T) { + subCtx := AppendAddressSegment(ctx, AddressSegmentAgent, "SubAgent") + subEvent := Interrupt(subCtx, "sub info") + event := CompositeInterrupt(ctx, "composite info", "composite state", subEvent.Action.internalInterrupted) + assert.NotNil(t, event.Action) + assert.NotNil(t, event.Action.Interrupted) + assert.NotNil(t, event.Action.Interrupted.InterruptContexts) + assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts)) + + rootCause := event.Action.Interrupted.InterruptContexts[0] + assert.Equal(t, "sub info", rootCause.Info) + assert.True(t, rootCause.IsRootCause) + assert.Equal(t, Address{ + {Type: AddressSegmentAgent, ID: "TestAgent"}, + {Type: AddressSegmentAgent, ID: "SubAgent"}, + }, rootCause.Address) + + assert.NotNil(t, rootCause.Parent, "Parent should not be nil for composite interrupt") + assert.Equal(t, "composite info", rootCause.Parent.Info) + assert.Equal(t, Address{ + {Type: AddressSegmentAgent, ID: "TestAgent"}, + }, rootCause.Parent.Address) + }) + + t.Run("Address only contains agent/tool segments", func(t *testing.T) { + event := Interrupt(ctx, "test info") + addr := event.Action.Interrupted.InterruptContexts[0].Address + for _, seg := range addr { + assert.True(t, seg.Type == AddressSegmentAgent || seg.Type == AddressSegmentTool, + "Address should only contain agent/tool segments, got: %s", seg.Type) + } + }) +} + +func TestSimpleInterrupt(t *testing.T) { + data := "hello world" + agent := &myAgent{ + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + Message: nil, + MessageStream: schema.StreamReaderFromArray([]Message{ + schema.UserMessage("hello "), + schema.UserMessage("world"), + }), + }, + }, + }) + intEvent := Interrupt(ctx, data) + intEvent.Action.Interrupted.Data = data + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.WasInterrupted) + assert.Nil(t, info.InterruptState) + assert.True(t, info.EnableStreaming) + assert.Equal(t, data, info.Data) + + assert.True(t, info.IsResumeTarget) + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } + store := newMyStore() + ctx := context.Background() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + CheckPointStore: store, + }) + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + _, ok := iter.Next() + assert.True(t, ok) + interruptEvent, ok := iter.Next() + assert.True(t, ok) + assert.Equal(t, data, interruptEvent.Action.Interrupted.Data) + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts[0].ID) + assert.True(t, interruptEvent.Action.Interrupted.InterruptContexts[0].IsRootCause) + assert.Equal(t, data, interruptEvent.Action.Interrupted.InterruptContexts[0].Info) + assert.Equal(t, Address{{Type: AddressSegmentAgent, ID: "myAgent"}}, + interruptEvent.Action.Interrupted.InterruptContexts[0].Address) + _, ok = iter.Next() + assert.False(t, ok) + + iter, err := runner.ResumeWithParams(ctx, "1", &ResumeParams{ + Targets: map[string]any{ + interruptEvent.Action.Interrupted.InterruptContexts[0].ID: nil, + }, + }) + assert.NoError(t, err) + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestMultiAgentInterrupt(t *testing.T) { + ctx := context.Background() + sa1 := &myAgent{ + name: "sa1", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "sa1", + Action: &AgentAction{ + TransferToAgent: &TransferToAgentAction{ + DestAgentName: "sa2", + }, + }, + }) + generator.Close() + return iter + }, + } + sa2 := &myAgent{ + name: "sa2", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + intEvent := StatefulInterrupt(ctx, "hello world", "temp state") + intEvent.Action.Interrupted.Data = "hello world" + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.NotNil(t, info) + assert.Equal(t, info.Data, "hello world") + + assert.True(t, info.WasInterrupted) + assert.NotNil(t, info.InterruptState) + assert.Equal(t, "temp state", info.InterruptState) + + assert.True(t, info.IsResumeTarget) + assert.NotNil(t, info.ResumeData) + assert.Equal(t, "resume data", info.ResumeData) + + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "sa2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: schema.UserMessage(info.ResumeData.(string))}, + }, + }) + generator.Close() + return iter + }, + } + a, err := SetSubAgents(ctx, sa1, []Agent{sa2}) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + EnableStreaming: false, + CheckPointStore: newMyStore(), + }) + iter := runner.Query(ctx, "", WithCheckPointID("1")) + event, ok := iter.Next() + assert.True(t, ok) + assert.NotNil(t, event.Action.TransferToAgent) + event, ok = iter.Next() + assert.True(t, ok) + assert.NotNil(t, event.Action.Interrupted) + assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts)) + assert.Equal(t, "hello world", event.Action.Interrupted.InterruptContexts[0].Info) + assert.True(t, event.Action.Interrupted.InterruptContexts[0].IsRootCause) + assert.Equal(t, Address{ + {Type: AddressSegmentAgent, ID: "sa1"}, + {Type: AddressSegmentAgent, ID: "sa2"}, + }, event.Action.Interrupted.InterruptContexts[0].Address) + assert.NotEmpty(t, event.Action.Interrupted.InterruptContexts[0].ID) + + interruptID := event.Action.Interrupted.InterruptContexts[0].ID + _, ok = iter.Next() + assert.False(t, ok) + + iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{ + Targets: map[string]any{ + interruptID: "resume data", + }, + }) + assert.NoError(t, err) + event, ok = iter.Next() + assert.True(t, ok) + assert.Equal(t, event.Output.MessageOutput.Message.Content, "resume data") + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestWorkflowInterrupt(t *testing.T) { + ctx := context.Background() + sa1 := &myAgent{ + name: "sa1", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + + intEvent := Interrupt(ctx, "sa1 interrupt data") + intEvent.Action.Interrupted.Data = "sa1 interrupt data" + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.Equal(t, info.InterruptInfo.Data, "sa1 interrupt data") + assert.True(t, info.WasInterrupted) + assert.Nil(t, info.InterruptState) + assert.True(t, info.IsResumeTarget) + assert.Equal(t, "resume sa1", info.ResumeData) + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } // interrupt once + sa2 := &myAgent{ + name: "sa2", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + + intEvent := StatefulInterrupt(ctx, "sa2 interrupt data", "sa2 interrupt") + intEvent.Action.Interrupted.Data = "sa2 interrupt data" + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.Equal(t, info.InterruptInfo.Data, "sa2 interrupt data") + assert.True(t, info.WasInterrupted) + assert.NotNil(t, info.InterruptState) + assert.Equal(t, "sa2 interrupt", info.InterruptState) + + assert.True(t, info.IsResumeTarget) + assert.NotNil(t, info.ResumeData) + assert.Equal(t, "resume sa2", info.ResumeData) + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } // interrupt once + sa3 := &myAgent{ + name: "sa3", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "sa3", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa3 completed"), + }, + }, + }) + generator.Close() + return iter + }, + } // won't interrupt + sa4 := &myAgent{ + name: "sa4", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "sa4", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa4 completed"), + }, + }, + }) + generator.Close() + return iter + }, + } // won't interrupt + + firstInterruptEvent := &AgentEvent{ + AgentName: "sa1", + RunPath: []RunStep{{"sequential"}, {"sa1"}}, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: &WorkflowInterruptInfo{ + OrigInput: &AgentInput{ + Messages: []Message{schema.UserMessage("hello world")}, + }, + SequentialInterruptIndex: 0, + SequentialInterruptInfo: &InterruptInfo{ + Data: "sa1 interrupt data", + }, + LoopIterations: 0, + }, + InterruptContexts: []*InterruptCtx{ + { + ID: "agent:sequential;agent:sa1", + Info: "sa1 interrupt data", + Address: Address{ + { + ID: "sequential", + Type: AddressSegmentAgent, + }, + { + ID: "sa1", + Type: AddressSegmentAgent, + }, + }, + IsRootCause: true, + Parent: &InterruptCtx{ + ID: "agent:sequential", + Info: "Sequential workflow interrupted", + Address: Address{ + { + ID: "sequential", + Type: AddressSegmentAgent, + }, + }, + }, + }, + }, + }, + }, + } + _ = firstInterruptEvent + secondInterruptEvent := &AgentEvent{ + AgentName: "sa2", + RunPath: []RunStep{{"sequential"}, {"sa1"}, {"sa2"}}, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: &WorkflowInterruptInfo{ + OrigInput: &AgentInput{ + Messages: []Message{schema.UserMessage("hello world")}, + }, + SequentialInterruptIndex: 1, + SequentialInterruptInfo: &InterruptInfo{ + Data: "sa2 interrupt data", + }, + }, + InterruptContexts: []*InterruptCtx{ + { + ID: "agent:sequential;agent:sa1;agent:sa2", + Info: "sa2 interrupt data", + Address: Address{ + { + ID: "sequential", + Type: AddressSegmentAgent, + }, + { + ID: "sa2", + Type: AddressSegmentAgent, + }, + }, + IsRootCause: true, + Parent: &InterruptCtx{ + ID: "agent:sequential", + Info: "Sequential workflow interrupted", + Address: Address{ + { + ID: "sequential", + Type: AddressSegmentAgent, + }, + }, + }, + }, + }, + }, + }, + } + _ = secondInterruptEvent + messageEvents := []*AgentEvent{ + { + AgentName: "sa3", + RunPath: []RunStep{{"sequential"}, {"sa1"}, {"sa2"}, {"sa3"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa3 completed"), + }, + }, + }, + { + AgentName: "sa4", + RunPath: []RunStep{{"sequential"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa4 completed"), + }, + }, + }, + } + _ = messageEvents + + t.Run("test sequential workflow agent", func(t *testing.T) { + + // sequential + a, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "sequential", + Description: "sequential agent", + SubAgents: []Agent{sa1, sa2, sa3, sa4}, + }) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + CheckPointStore: newMyStore(), + }) + var events []*AgentEvent + iter := runner.Query(ctx, "hello world", WithCheckPointID("sequential-1")) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 1, len(events)) + assert.Equal(t, firstInterruptEvent.AgentName, events[0].AgentName) + assert.Equal(t, firstInterruptEvent.RunPath, events[0].RunPath) + assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(firstInterruptEvent.Action.Interrupted.InterruptContexts[0])) + interruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID + events = []*AgentEvent{} + + // Resume after sa1 interrupt + iter, err = runner.ResumeWithParams(ctx, "sequential-1", &ResumeParams{ + Targets: map[string]any{ + interruptID1: "resume sa1", + }, + }) + assert.NoError(t, err) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 1, len(events)) + assert.Equal(t, secondInterruptEvent.AgentName, events[0].AgentName) + assert.Equal(t, secondInterruptEvent.RunPath, events[0].RunPath) + assert.True(t, events[0].Action.Interrupted.InterruptContexts[0]. + EqualsWithoutID(secondInterruptEvent.Action.Interrupted.InterruptContexts[0])) + interruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID + events = []*AgentEvent{} + + // Resume after sa2 interrupt + iter, err = runner.ResumeWithParams(ctx, "sequential-1", &ResumeParams{ + Targets: map[string]any{ + interruptID2: "resume sa2", + }, + }) + assert.NoError(t, err) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 2, len(events)) + assert.Equal(t, messageEvents, events) + }) + + t.Run("test loop workflow agent", func(t *testing.T) { + // loop + a, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "loop", + SubAgents: []Agent{sa1, sa2, sa3, sa4}, + MaxIterations: 2, + }) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + CheckPointStore: newMyStore(), + }) + var events []*AgentEvent + iter := runner.Query(ctx, "hello world", WithCheckPointID("loop-1")) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + loopFirstInterruptEvent := &AgentEvent{ + AgentName: "sa1", + RunPath: []RunStep{{"loop"}, {"sa1"}}, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: &WorkflowInterruptInfo{ + OrigInput: &AgentInput{ + Messages: []Message{schema.UserMessage("hello world")}, + }, + SequentialInterruptIndex: 0, + SequentialInterruptInfo: &InterruptInfo{ + Data: "sa1 interrupt data", + }, + LoopIterations: 0, + }, + InterruptContexts: []*InterruptCtx{ + { + ID: "agent:loop;agent:sa1", + Info: "sa1 interrupt data", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + { + ID: "sa1", + Type: AddressSegmentAgent, + }, + }, + IsRootCause: true, + Parent: &InterruptCtx{ + ID: "agent:loop", + Info: "Loop workflow interrupted", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + }, + }, + }, + }, + }, + }, + } + assert.Equal(t, 1, len(events)) + assert.Equal(t, loopFirstInterruptEvent.AgentName, events[0].AgentName) + assert.Equal(t, loopFirstInterruptEvent.RunPath, events[0].RunPath) + assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFirstInterruptEvent.Action.Interrupted.InterruptContexts[0])) + loopInterruptID1 := events[0].Action.Interrupted.InterruptContexts[0].ID + events = []*AgentEvent{} + + // Resume after sa1 interrupt + iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{ + Targets: map[string]any{ + loopInterruptID1: "resume sa1", + }, + }) + assert.NoError(t, err) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + loopSecondInterruptEvent := &AgentEvent{ + AgentName: "sa2", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}}, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: &WorkflowInterruptInfo{ + OrigInput: &AgentInput{ + Messages: []Message{schema.UserMessage("hello world")}, + }, + SequentialInterruptIndex: 1, + SequentialInterruptInfo: &InterruptInfo{ + Data: "sa2 interrupt data", + }, + LoopIterations: 0, + }, + InterruptContexts: []*InterruptCtx{ + { + ID: "agent:loop;agent:sa1;agent:sa2", + Info: "sa2 interrupt data", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + { + ID: "sa2", + Type: AddressSegmentAgent, + }, + }, + IsRootCause: true, + Parent: &InterruptCtx{ + ID: "agent:loop", + Info: "Loop workflow interrupted", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + }, + }, + }, + }, + }, + }, + } + assert.Equal(t, 1, len(events)) + assert.Equal(t, loopSecondInterruptEvent.AgentName, events[0].AgentName) + assert.Equal(t, loopSecondInterruptEvent.RunPath, events[0].RunPath) + assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopSecondInterruptEvent.Action.Interrupted.InterruptContexts[0])) + loopInterruptID2 := events[0].Action.Interrupted.InterruptContexts[0].ID + events = []*AgentEvent{} + + // Resume after sa2 interrupt + iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{ + Targets: map[string]any{ + loopInterruptID2: "resume sa2", + }, + }) + assert.NoError(t, err) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + loopThirdInterruptEvent := &AgentEvent{ + AgentName: "sa1", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}}, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: &WorkflowInterruptInfo{ + OrigInput: &AgentInput{ + Messages: []Message{schema.UserMessage("hello world")}, + }, + SequentialInterruptIndex: 0, + SequentialInterruptInfo: &InterruptInfo{ + Data: "sa1 interrupt data", + }, + LoopIterations: 1, + }, + InterruptContexts: []*InterruptCtx{ + { + ID: "agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1", + Info: "sa1 interrupt data", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + { + ID: "sa1", + Type: AddressSegmentAgent, + }, + }, + IsRootCause: true, + Parent: &InterruptCtx{ + ID: "agent:loop", + Info: "Loop workflow interrupted", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + }, + }, + }, + }, + }, + }, + } + + loopFourthInterruptEvent := &AgentEvent{ + AgentName: "sa2", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}}, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: &WorkflowInterruptInfo{ + OrigInput: &AgentInput{ + Messages: []Message{schema.UserMessage("hello world")}, + }, + SequentialInterruptIndex: 1, + SequentialInterruptInfo: &InterruptInfo{ + Data: "sa2 interrupt data", + }, + LoopIterations: 1, + }, + InterruptContexts: []*InterruptCtx{ + { + ID: "agent:loop;agent:sa1;agent:sa2;agent:sa3;agent:sa4;agent:sa1;agent:sa2", + Info: "sa2 interrupt data", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + { + ID: "sa2", + Type: AddressSegmentAgent, + }, + }, + IsRootCause: true, + Parent: &InterruptCtx{ + ID: "agent:loop", + Info: "Loop workflow interrupted", + Address: Address{ + { + ID: "loop", + Type: AddressSegmentAgent, + }, + }, + }, + }, + }, + }, + }, + } + + loopMessageEvents := []*AgentEvent{ + { + AgentName: "sa3", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa3 completed"), + }, + }, + }, + { + AgentName: "sa4", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa4 completed"), + }, + }, + }, + loopThirdInterruptEvent, + } + assert.Equal(t, 3, len(events)) + // Check the first two message events + assert.Equal(t, loopMessageEvents[0].AgentName, events[0].AgentName) + assert.Equal(t, loopMessageEvents[0].RunPath, events[0].RunPath) + assert.Equal(t, loopMessageEvents[0].Output.MessageOutput.Message.Content, events[0].Output.MessageOutput.Message.Content) + + assert.Equal(t, loopMessageEvents[1].AgentName, events[1].AgentName) + assert.Equal(t, loopMessageEvents[1].RunPath, events[1].RunPath) + assert.Equal(t, loopMessageEvents[1].Output.MessageOutput.Message.Content, events[1].Output.MessageOutput.Message.Content) + + // Check the third interrupt event using EqualsWithoutID + assert.Equal(t, loopMessageEvents[2].AgentName, events[2].AgentName) + assert.Equal(t, loopMessageEvents[2].RunPath, events[2].RunPath) + assert.True(t, events[2].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopMessageEvents[2].Action.Interrupted.InterruptContexts[0])) + loopInterruptID3 := events[2].Action.Interrupted.InterruptContexts[0].ID + events = []*AgentEvent{} + + // Resume after third interrupt + iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{ + Targets: map[string]any{ + loopInterruptID3: "resume sa1", + }, + }) + assert.NoError(t, err) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + assert.Equal(t, 1, len(events)) + assert.Equal(t, loopFourthInterruptEvent.AgentName, events[0].AgentName) + assert.Equal(t, loopFourthInterruptEvent.RunPath, events[0].RunPath) + assert.True(t, events[0].Action.Interrupted.InterruptContexts[0].EqualsWithoutID(loopFourthInterruptEvent.Action.Interrupted.InterruptContexts[0])) + loopInterruptID4 := events[0].Action.Interrupted.InterruptContexts[0].ID + events = []*AgentEvent{} + + // Resume after fourth interrupt + iter, err = runner.ResumeWithParams(ctx, "loop-1", &ResumeParams{ + Targets: map[string]any{ + loopInterruptID4: "resume sa2", + }, + }) + assert.NoError(t, err) + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + loopFinalMessageEvents := []*AgentEvent{ + { + AgentName: "sa3", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}, {"sa3"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa3 completed"), + }, + }, + }, + { + AgentName: "sa4", + RunPath: []RunStep{{"loop"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}, {"sa1"}, {"sa2"}, {"sa3"}, {"sa4"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa4 completed"), + }, + }, + }, + } + assert.Equal(t, 2, len(events)) + assert.Equal(t, loopFinalMessageEvents, events) + }) + + t.Run("test parallel workflow agent", func(t *testing.T) { + // parallel + a, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "parallel agent", + SubAgents: []Agent{sa1, sa2, sa3, sa4}, + }) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + CheckPointStore: newMyStore(), + }) + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + var ( + events []*AgentEvent + interruptEvent *AgentEvent + ) + + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + continue + } + events = append(events, event) + } + assert.Equal(t, 2, len(events)) + + // Debug: Print actual events to see what we're getting + for i, event := range events { + t.Logf("Event %d: AgentName=%s, RunPath=%v, Output=%v", i, event.AgentName, event.RunPath, event.Output) + } + + // Define parallel message events separately + parallelMessageEvents := []*AgentEvent{ + { + AgentName: "sa4", + RunPath: []RunStep{{"parallel agent"}, {"sa4"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa4 completed"), + }, + }, + }, + { + AgentName: "sa3", + RunPath: []RunStep{{"parallel agent"}, {"sa3"}}, + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa3 completed"), + }, + }, + }, + } + + assert.Contains(t, events, parallelMessageEvents[0]) + assert.Contains(t, events, parallelMessageEvents[1]) + + assert.NotNil(t, interruptEvent) + assert.Equal(t, "parallel agent", interruptEvent.AgentName) + assert.Equal(t, []RunStep{{"parallel agent"}}, interruptEvent.RunPath) + assert.NotNil(t, interruptEvent.Action.Interrupted) + wii, ok := interruptEvent.Action.Interrupted.Data.(*WorkflowInterruptInfo) + assert.True(t, ok) + assert.Equal(t, 2, len(wii.ParallelInterruptInfo)) + + var sa1Found, sa2Found bool + for _, info := range wii.ParallelInterruptInfo { + switch info.Data { + case "sa1 interrupt data": + sa1Found = true + case "sa2 interrupt data": + sa2Found = true + } + } + assert.True(t, sa1Found) + assert.True(t, sa2Found) + + var sa1InfoFound, sa2InfoFound bool + for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + switch ctx.Info { + case "sa1 interrupt data": + sa1InfoFound = true + case "sa2 interrupt data": + sa2InfoFound = true + } + } + + assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts)) + assert.True(t, sa1InfoFound) + assert.True(t, sa2InfoFound) + + var parallelInterruptID1, parallelInterruptID2 string + for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + switch ctx.Info { + case "sa1 interrupt data": + parallelInterruptID1 = ctx.ID + case "sa2 interrupt data": + parallelInterruptID2 = ctx.ID + } + } + assert.NotEmpty(t, parallelInterruptID1) + assert.NotEmpty(t, parallelInterruptID2) + + iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{ + Targets: map[string]any{ + parallelInterruptID1: "resume sa1", + parallelInterruptID2: "resume sa2", + }, + }) + assert.NoError(t, err) + _, ok = iter.Next() + assert.False(t, ok) + }) +} + +func TestChatModelInterrupt(t *testing.T) { + ctx := context.Background() + a, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "name", + Description: "description", + Instruction: "instruction", + Model: &myModel{ + validator: func(i int, messages []*schema.Message) bool { + if i > 0 && (len(messages) != 4 || messages[2].Content != "new user message") { + return false + } + return true + }, + messages: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "1", + Function: schema.FunctionCall{ + Name: "tool1", + Arguments: "arguments", + }, + }, + }), + schema.AssistantMessage("completed", nil), + }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&myTool1{}}, + }, + }, + }) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + CheckPointStore: newMyStore(), + }) + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + _, ok := iter.Next() + assert.True(t, ok) + event, ok := iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.NotNil(t, event.Action.Interrupted) + assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts)) + assert.Equal(t, Address{ + {Type: AddressSegmentAgent, ID: "name"}, + {Type: AddressSegmentTool, ID: "tool1", SubID: "1"}, + }, event.Action.Interrupted.InterruptContexts[0].Address) + + var ( + chatModelAgentID string + toolID string + ) + + intCtx := event.Action.Interrupted.InterruptContexts[0] + for intCtx != nil { + if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentTool { + toolID = intCtx.ID + } else if intCtx.Address[len(intCtx.Address)-1].Type == AddressSegmentAgent { + chatModelAgentID = intCtx.ID + } + intCtx = intCtx.Parent + } + + _, ok = iter.Next() + assert.False(t, ok) + + iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{ + Targets: map[string]any{ + chatModelAgentID: &ChatModelAgentResumeData{ + HistoryModifier: func(ctx context.Context, history []Message) []Message { + history[2].Content = "new user message" + return history + }, + }, + toolID: "tool resume result", + }, + }) + assert.NoError(t, err) + event, ok = iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.Equal(t, event.Output.MessageOutput.Message.Content, "tool resume result") + event, ok = iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.Equal(t, event.Output.MessageOutput.Message.Content, "completed") +} + +func TestChatModelAgentToolInterrupt(t *testing.T) { + sa := &myAgent{ + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + intAct := Interrupt(ctx, "hello world") + intAct.Action.Interrupted.Data = "hello world" + generator.Send(intAct) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.NotNil(t, info) + assert.False(t, info.EnableStreaming) + + if !info.IsResumeTarget { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + intAct := Interrupt(ctx, "interrupt again") + intAct.Action.Interrupted.Data = "interrupt again" + generator.Send(intAct) + generator.Close() + return iter + } + + assert.NotNil(t, info.ResumeData) + assert.Equal(t, "resume sa", info.ResumeData) + + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{Output: &AgentOutput{MessageOutput: &MessageVariant{Message: schema.UserMessage(fmt.Sprintf("my agent completed with data %s", info.ResumeData))}}}) + generator.Close() + return iter + }, + } + ctx := context.Background() + a, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "name", + Description: "description", + Instruction: "instruction", + Model: &myModel{ + messages: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "1", + Function: schema.FunctionCall{ + Name: "myAgent", + Arguments: "{\"request\":\"123\"}", + }, + }, + }), + schema.AssistantMessage("completed", nil), + }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{NewAgentTool(ctx, sa)}, + }, + }, + }) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + CheckPointStore: newMyStore(), + }) + + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + _, ok := iter.Next() + assert.True(t, ok) + event, ok := iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.NotNil(t, event.Action.Interrupted) + _, ok = iter.Next() + assert.False(t, ok) + + iter, err = runner.Resume(ctx, "1") + assert.NoError(t, err) + event, ok = iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.NotNil(t, event.Action.Interrupted) + assert.Equal(t, 1, len(event.Action.Interrupted.InterruptContexts)) + for _, ctx := range event.Action.Interrupted.InterruptContexts { + if ctx.IsRootCause { + assert.Equal(t, Address{ + {Type: AddressSegmentAgent, ID: "name"}, + {Type: AddressSegmentTool, ID: "myAgent", SubID: "1"}, + {Type: AddressSegmentAgent, ID: "myAgent"}, + }, ctx.Address) + assert.Equal(t, "interrupt again", ctx.Info) + } + } + + var toolInterruptID string + for _, ctx := range event.Action.Interrupted.InterruptContexts { + if ctx.IsRootCause { + toolInterruptID = ctx.ID + break + } + } + assert.NotEmpty(t, toolInterruptID) + + _, ok = iter.Next() + assert.False(t, ok) + + iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{ + Targets: map[string]any{ + toolInterruptID: "resume sa", + }, + }) + assert.NoError(t, err) + event, ok = iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.Equal(t, event.Output.MessageOutput.Message.Content, "my agent completed with data resume sa") + event, ok = iter.Next() + assert.True(t, ok) + assert.NoError(t, event.Err) + assert.Equal(t, event.Output.MessageOutput.Message.Content, "completed") + _, ok = iter.Next() + assert.False(t, ok) +} + +func newMyStore() *myStore { + return &myStore{ + m: map[string][]byte{}, + } +} + +type myStore struct { + m map[string][]byte +} + +func (m *myStore) Set(_ context.Context, key string, value []byte) error { + m.m[key] = value + return nil +} + +func (m *myStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := m.m[key] + return v, ok, nil +} + +type myAgentOptions struct { + value string +} + +func withValue(value string) AgentRunOption { + return WrapImplSpecificOptFn(func(t *myAgentOptions) { + t.value = value + }) +} + +type myAgent struct { + name string + runFn func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] + resumeFn func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] +} + +func (m *myAgent) Name(_ context.Context) string { + if len(m.name) > 0 { + return m.name + } + return "myAgent" +} + +func (m *myAgent) Description(_ context.Context) string { + return "myAgent description" +} + +func (m *myAgent) Run(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return m.runFn(ctx, input, options...) +} + +func (m *myAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + return m.resumeFn(ctx, info, opts...) +} + +type myModel struct { + times int + messages []*schema.Message + validator func(int, []*schema.Message) bool +} + +func (m *myModel) Generate(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.Message, error) { + if m.validator != nil && !m.validator(m.times, input) { + return nil, errors.New("invalid input") + } + if m.times >= len(m.messages) { + return nil, errors.New("exceeded max number of messages") + } + t := m.times + m.times++ + return m.messages[t], nil +} + +func (m *myModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("implement me") +} + +func (m *myModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +type myTool1 struct{} + +func (m *myTool1) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "tool1", + Desc: "desc", + }, nil +} + +func (m *myTool1) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + if wasInterrupted, _, _ := tool.GetInterruptState[any](ctx); !wasInterrupted { + return "", tool.Interrupt(ctx, nil) + } + + if isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx); !isResumeFlow { + return "", tool.Interrupt(ctx, nil) + } else if hasResumeData { + return data, nil + } + + return "result", nil +} + +func TestCyclicalAgentInterrupt(t *testing.T) { + ctx := context.Background() + + var agentA, agentB, agentC Agent + + // agentC interrupts + agentC = &myAgent{ + name: "C", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + intAct := Interrupt(ctx, "interrupt from C") + generator.Send(intAct) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.IsResumeTarget) + assert.NotNil(t, info.ResumeData) + assert.Equal(t, "resume C", info.ResumeData) + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "C", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{Message: schema.UserMessage("C completed")}, + }, + }) + generator.Close() + return iter + }, + } + + // agentB transfers back to its parent A + agentB = &myAgent{ + name: "B", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "B", + Action: &AgentAction{ + TransferToAgent: &TransferToAgentAction{ + DestAgentName: "A", // Transfer back to parent + }, + }, + }) + generator.Close() + return iter + }, + } + + // agentA is the parent, orchestrating the A->B->A->C flow + agentA = &myAgent{ + name: "A", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + runCtx := getRunCtx(ctx) + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + + // If the last agent was B, we are in the A->B->A path, so transfer to C. + // Otherwise, it's the first run, transfer to B. + dest := "B" + if len(runCtx.RunPath) > 1 && runCtx.RunPath[len(runCtx.RunPath)-2].agentName == "B" { + dest = "C" + } + + generator.Send(&AgentEvent{ + AgentName: "A", + Action: &AgentAction{ + TransferToAgent: &TransferToAgentAction{ + DestAgentName: dest, + }, + }, + }) + generator.Close() + return iter + }, + } + + // Set up the hierarchy: A is parent of B and C. + agentA, err := SetSubAgents(ctx, agentA, []Agent{agentB, agentC}) + assert.NoError(t, err) + + // Run the test + runner := NewRunner(ctx, RunnerConfig{ + Agent: agentA, + CheckPointStore: newMyStore(), + }) + iter := runner.Query(ctx, "start", WithCheckPointID("cyclical-1")) + + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + // We expect 3 transfer events (A->B, B->A, A->C) and 1 interrupt event from C. + assert.Equal(t, 4, len(events)) + + interruptEvent := events[3] + assert.NotNil(t, interruptEvent.Action.Interrupted) + assert.Equal(t, "C", interruptEvent.AgentName) + + // Check the interrupt context + assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts)) + interruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0] + assert.True(t, interruptCtx.IsRootCause) + assert.Equal(t, "interrupt from C", interruptCtx.Info) + + expectedAddr := Address{ + {Type: AddressSegmentAgent, ID: "A"}, + {Type: AddressSegmentAgent, ID: "B"}, + {Type: AddressSegmentAgent, ID: "A"}, + {Type: AddressSegmentAgent, ID: "C"}, + } + assert.Equal(t, expectedAddr, interruptCtx.Address) + assert.NotEmpty(t, interruptCtx.ID) + + // Resume the execution + iter, err = runner.ResumeWithParams(ctx, "cyclical-1", &ResumeParams{ + Targets: map[string]any{ + interruptCtx.ID: "resume C", + }, + }) + assert.NoError(t, err) + + events = []*AgentEvent{} + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + + // We expect one output event from C + assert.Equal(t, 1, len(events)) + assert.Equal(t, "C completed", events[0].Output.MessageOutput.Message.Content) +} + +// myStatefulTool is a tool that can interrupt and has internal state to track invocations. + +type myStatefulTool struct { + name string + t *testing.T +} + +func (m *myStatefulTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: m.name, + Desc: "desc", + }, nil +} + +type myStatefulToolState struct { + InterruptCount int +} + +func init() { + schema.Register[myStatefulToolState]() +} + +func (m *myStatefulTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + wasInterrupted, hasState, state := tool.GetInterruptState[myStatefulToolState](ctx) + if !wasInterrupted { + return "", tool.StatefulInterrupt(ctx, fmt.Sprintf("interrupt from %s", m.name), myStatefulToolState{InterruptCount: 1}) + } + + isResumeFlow, hasResumeData, data := tool.GetResumeContext[string](ctx) + if !isResumeFlow || !hasResumeData { + assert.True(m.t, hasState, "tool %s should have interrupt state on resume", m.name) + return "", tool.StatefulInterrupt(ctx, fmt.Sprintf("interrupt from %s", m.name), myStatefulToolState{InterruptCount: state.InterruptCount + 1}) + } + + return data, nil +} + +func TestChatModelParallelToolInterruptAndResume(t *testing.T) { + ctx := context.Background() + + toolA := &myStatefulTool{name: "toolA", t: t} + toolB := &myStatefulTool{name: "toolB", t: t} + + chatModel, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "ParallelToolAgent", + Description: "An agent that uses parallel tools", + Model: &myModel{ + messages: []*schema.Message{ + // 1. First model response: call toolA and toolB in parallel + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "toolA", Arguments: "{}"}}, + {ID: "2", Function: schema.FunctionCall{Name: "toolB", Arguments: "{}"}}, + }), + // 2. Second model response (after tools are resumed): call them again to check state + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "3", Function: schema.FunctionCall{Name: "toolA", Arguments: "{}"}}, + {ID: "4", Function: schema.FunctionCall{Name: "toolB", Arguments: "{}"}}, + }), + // 3. Final completion + schema.AssistantMessage("all done", nil), + }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{toolA, toolB}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: chatModel, + CheckPointStore: newMyStore(), + }) + + // 1. Initial query -> parallel interrupt from toolA and toolB + iter := runner.Query(ctx, "start", WithCheckPointID("parallel-tool-test-1")) + normalEvents, interruptEvent := consumeUntilInterrupt(iter) + + assert.Equal(t, 1, len(normalEvents)) + assert.NotNil(t, interruptEvent) + assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts), + "should have 2 interrupts") + + var toolAInterruptID, toolBInterruptID string + for _, info := range interruptEvent.Action.Interrupted.InterruptContexts { + switch info.Info { + case "interrupt from toolA": + toolAInterruptID = info.ID + assert.True(t, info.IsRootCause) + case "interrupt from toolB": + toolBInterruptID = info.ID + assert.True(t, info.IsRootCause) + } + } + assert.NotEmpty(t, toolAInterruptID) + assert.NotEmpty(t, toolBInterruptID) + + // 2. Resume, targeting only toolA. toolB should re-interrupt. + iter, err = runner.ResumeWithParams(ctx, "parallel-tool-test-1", &ResumeParams{ + Targets: map[string]any{ + toolAInterruptID: "toolA resumed", + }, + }) + assert.NoError(t, err) + _, interruptEvent = consumeUntilInterrupt(iter) + + assert.NotNil(t, interruptEvent, "expected a re-interrupt from toolB") + assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts), + "should have 1 remaining interrupts") + + var rootCause *InterruptCtx + for _, info := range interruptEvent.Action.Interrupted.InterruptContexts { + if info.IsRootCause { + rootCause = info + break + } + } + + if rootCause == nil { + t.Fatal("expected a root cause interrupt from toolB") + } + assert.Equal(t, "interrupt from toolB", rootCause.Info) + toolBReInterruptID := rootCause.ID + + // 3. Resume the re-interrupted toolB. The agent should then call the tools again. + iter, err = runner.ResumeWithParams(ctx, "parallel-tool-test-1", &ResumeParams{ + Targets: map[string]any{ + toolBReInterruptID: "toolB resumed", + }, + }) + assert.NoError(t, err) + + // 4. Consume all final events. The internal assertions in the tools will check the wasInterrupted flag. + // We expect to see the results of the second tool calls, and then the final agent completion. + finalEvents, interruptEvent := consumeUntilInterrupt(iter) + assert.Equal(t, 2, len(finalEvents)) + assert.NotNil(t, interruptEvent) +} + +// TestNestedChatModelAgentWithAgentTool verifies that the shouldFire method correctly prevents +// duplicate event firing in nested ChatModelAgent scenarios (ChatModelAgent -> AgentTool -> ChatModelAgent). +// This ensures that only the inner agent's cbHandler fires, not the outer agent's. +func TestNestedChatModelAgentWithAgentTool(t *testing.T) { + ctx := context.Background() + + // Create an interruptible tool for the inner agent + innerTool := &myStatefulTool{name: "innerTool", t: t} + + // Create the inner ChatModelAgent that will be wrapped by AgentTool + innerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "InnerAgent", + Description: "Inner agent with interruptible tool", + Model: &myModel{ + messages: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "innerTool", Arguments: "{}"}}, + }), + schema.AssistantMessage("inner agent completed", nil), + }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{innerTool}, + }, + }, + }) + assert.NoError(t, err) + + // Wrap the inner agent in an AgentTool + agentTool := NewAgentTool(ctx, innerAgent) + + // Create the outer ChatModelAgent that uses the AgentTool + outerAgent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "OuterAgent", + Description: "Outer agent with AgentTool containing inner agent", + Model: &myModel{ + messages: []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "InnerAgent", Arguments: "{}"}}, + }), + schema.AssistantMessage("outer agent completed", nil), + }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{agentTool}, + }, + }, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: outerAgent, + CheckPointStore: newMyStore(), + }) + + // Run the query - this should trigger the nested agent structure + iter := runner.Query(ctx, "start", WithCheckPointID("nested-agent-test-1")) + + // Collect all events to verify no duplicates + var allEvents []*AgentEvent + var interruptEvent *AgentEvent + + for { + event, ok := iter.Next() + if !ok { + break + } + + if event.Action != nil && event.Action.Interrupted != nil { + assert.Nil(t, interruptEvent) + interruptEvent = event + } + + allEvents = append(allEvents, event) + } + + if interruptEvent == nil { + t.Fatal("expected an interrupt event") + } + + // Verify we got exactly one interrupt event (not duplicated) + assert.NotNil(t, interruptEvent, "should have an interrupt event") + assert.Equal(t, 1, len(interruptEvent.Action.Interrupted.InterruptContexts), + "should have exactly one interrupt context") + + // Verify the interrupt comes from the inner tool, not duplicated + interruptCtx := interruptEvent.Action.Interrupted.InterruptContexts[0] + assert.True(t, interruptCtx.IsRootCause, "interrupt should be root cause") + assert.Equal(t, "interrupt from innerTool", interruptCtx.Info) + + // Verify the address path shows the correct nested structure + expectedAddress := Address{ + {Type: AddressSegmentAgent, ID: "OuterAgent"}, + {Type: AddressSegmentTool, ID: "InnerAgent", SubID: "1"}, + {Type: AddressSegmentAgent, ID: "InnerAgent"}, + {Type: AddressSegmentTool, ID: "innerTool", SubID: "1"}, + } + assert.Equal(t, expectedAddress, interruptCtx.Address, + "interrupt address should show correct nested structure") + + // Verify no duplicate events by checking agent names in events + var agentNames []string + for _, event := range allEvents { + if event.AgentName != "" { + agentNames = append(agentNames, event.AgentName) + } + } + + // Should only have events from the outer agent (the inner agent's events should be handled + // by the AgentTool and not duplicated by the outer agent's cbHandler) + for _, name := range agentNames { + assert.Equal(t, "OuterAgent", name, + "all events should come from OuterAgent, not duplicated from InnerAgent") + } + + // Now resume the interrupt + interruptID := interruptCtx.ID + iter, err = runner.ResumeWithParams(ctx, "nested-agent-test-1", &ResumeParams{ + Targets: map[string]any{ + interruptID: "resume inner tool", + }, + }) + assert.NoError(t, err) + + // Collect final events after resume + var finalEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + finalEvents = append(finalEvents, event) + } + + // Verify completion events + assert.Greater(t, len(finalEvents), 0, "should have completion events after resume") + + // Check that we get the expected completion messages + var foundInnerCompletion, foundOuterCompletion bool + for _, event := range finalEvents { + if event.Output != nil && event.Output.MessageOutput != nil { + if event.Output.MessageOutput.Message != nil { + content := event.Output.MessageOutput.Message.Content + switch content { + case "inner agent completed": + foundInnerCompletion = true + case "outer agent completed": + foundOuterCompletion = true + } + } + } + } + + assert.True(t, foundInnerCompletion, "should have inner agent completion") + assert.True(t, foundOuterCompletion, "should have outer agent completion") +} + +// consumeUntilInterrupt consumes events from the iterator until an interrupt is found or it's exhausted. +func consumeUntilInterrupt(iter *AsyncIterator[*AgentEvent]) (normalEvents []*AgentEvent, interruptEvent *AgentEvent) { + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + continue + } + normalEvents = append(normalEvents, event) + } + return +} + +type returnDirectlyTool struct { + name string +} + +func (t *returnDirectlyTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "A tool that returns directly", + }, nil +} + +func (t *returnDirectlyTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return "return directly result", nil +} + +type interruptingTool struct { + name string +} + +func (i *interruptingTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: i.name, + Desc: "A tool that interrupts", + }, nil +} + +func (i *interruptingTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + if wasInterrupted, _, _ := compose.GetInterruptState[any](ctx); !wasInterrupted { + return "", compose.Interrupt(ctx, "interrupt data") + } + + if isResumeFlow, hasResumeData, data := compose.GetResumeContext[string](ctx); isResumeFlow && hasResumeData { + return data, nil + } + + return "resumed without data", nil +} + +type twoToolCallModel struct { + returnDirectlyToolName string + interruptingToolName string + callCount int + receivedTools []*schema.ToolInfo + mu sync.Mutex +} + +func (m *twoToolCallModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) { + m.mu.Lock() + m.callCount++ + callNum := m.callCount + options := model.GetCommonOptions(&model.Options{}, opts...) + if options.Tools != nil { + m.receivedTools = options.Tools + } + m.mu.Unlock() + + if callNum == 1 { + return &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: "call_return_directly", + Type: "function", + Function: schema.FunctionCall{ + Name: m.returnDirectlyToolName, + Arguments: "{}", + }, + }, + { + ID: "call_interrupting", + Type: "function", + Function: schema.FunctionCall{ + Name: m.interruptingToolName, + Arguments: "{}", + }, + }, + }, + }, nil + } + return schema.AssistantMessage("final response", nil), nil +} + +func (m *twoToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("not implemented") +} + +func (m *twoToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func (m *twoToolCallModel) GetReceivedTools() []*schema.ToolInfo { + m.mu.Lock() + defer m.mu.Unlock() + return m.receivedTools +} + +type dynamicTool struct { + name string +} + +func (t *dynamicTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "A dynamically added tool", + }, nil +} + +func (t *dynamicTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return "dynamic tool result", nil +} + +func TestReturnDirectlyEventSentAfterResume(t *testing.T) { + ctx := context.Background() + + returnDirectlyToolName := "return_directly_tool" + interruptingToolName := "interrupting_tool" + dynamicToolName := "dynamic_tool" + + mdl := &twoToolCallModel{ + returnDirectlyToolName: returnDirectlyToolName, + interruptingToolName: interruptingToolName, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent for return directly + interrupt", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{ + &returnDirectlyTool{name: returnDirectlyToolName}, + &interruptingTool{name: interruptingToolName}, + }, + }, + ReturnDirectly: map[string]bool{ + returnDirectlyToolName: true, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &interruptTestToolsHandler{tools: []tool.BaseTool{&dynamicTool{name: dynamicToolName}}}, + }, + }) + assert.NoError(t, err) + + store := newMyStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + CheckPointStore: store, + }) + + iter := runner.Query(ctx, "test input", WithCheckPointID("test_checkpoint")) + + var interruptEvent *AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + } + + assert.NotNil(t, interruptEvent, "Should have an interrupt event") + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts) + + receivedToolsBeforeResume := mdl.GetReceivedTools() + var hasDynamicToolBeforeResume bool + for _, ti := range receivedToolsBeforeResume { + if ti.Name == dynamicToolName { + hasDynamicToolBeforeResume = true + } + } + assert.True(t, hasDynamicToolBeforeResume, "Dynamic tool should be in tool list before interrupt") + + interruptID := interruptEvent.Action.Interrupted.InterruptContexts[0].ID + resumeIter, err := runner.ResumeWithParams(ctx, "test_checkpoint", &ResumeParams{ + Targets: map[string]any{ + interruptID: "resume data", + }, + }) + assert.NoError(t, err) + + var resumeEvents []*AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + resumeEvents = append(resumeEvents, event) + } + + var hasReturnDirectlyEvent bool + for _, e := range resumeEvents { + if e.Output != nil && e.Output.MessageOutput != nil { + if e.Output.MessageOutput.Role == schema.Tool && e.Output.MessageOutput.ToolName == returnDirectlyToolName { + hasReturnDirectlyEvent = true + } + } + } + assert.True(t, hasReturnDirectlyEvent, "ReturnDirectlyEvent should be sent after resume") + + receivedToolsAfterResume := mdl.GetReceivedTools() + var hasDynamicToolAfterResume bool + for _, ti := range receivedToolsAfterResume { + if ti.Name == dynamicToolName { + hasDynamicToolAfterResume = true + } + } + assert.True(t, hasDynamicToolAfterResume, "Dynamic tool should be in tool list after resume (bc.toolUpdated path)") +} + +// streamErrorThenToolCallModel simulates a model that: +// - On the first Stream call: emits several good chunks then an error (triggering retry) +// - On the second Stream call (retry): returns a tool call message (success) +type streamErrorThenToolCallModel struct { + callCount int32 + toolCallName string +} + +func (m *streamErrorThenToolCallModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("final answer", nil), nil +} + +func (m *streamErrorThenToolCallModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + count := atomic.AddInt32(&m.callCount, 1) + + sr, sw := schema.Pipe[*schema.Message](10) + go func() { + defer sw.Close() + if count == 1 { + // First call: emit good chunks then error + sw.Send(schema.AssistantMessage("chunk1", nil), nil) + sw.Send(schema.AssistantMessage("chunk2", nil), nil) + sw.Send(schema.AssistantMessage("chunk3", nil), nil) + sw.Send(nil, errRetryAble) + return + } + // Second call (retry): return tool call + sw.Send(schema.AssistantMessage("", []schema.ToolCall{{ + ID: "call-1", + Function: schema.FunctionCall{Name: m.toolCallName, Arguments: "{}"}, + }}), nil) + }() + return sr, nil +} + +func (m *streamErrorThenToolCallModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +// TestStreamRetryThenToolInterruptCheckpoint reproduces a bug where: +// 1. ChatModelAgent with ModelRetryConfig.MaxRetries = 2 +// 2. First model Stream call emits good chunks then a retryable error +// 3. Retry succeeds, model returns a tool call +// 4. The tool triggers an interrupt, causing the Runner to save a checkpoint +// 5. The checkpoint save fails because the first (failed) model call's stream event +// is in the session, and when MessageVariant.GobEncode consumes the stream, +// it hits the error chunk and returns an encoding error. +func TestStreamRetryThenToolInterruptCheckpoint(t *testing.T) { + ctx := context.Background() + + interruptToolName := "interrupt_tool" + mdl := &streamErrorThenToolCallModel{toolCallName: interruptToolName} + + interruptTool := &interruptingTool{name: interruptToolName} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "RetryInterruptAgent", + Description: "Agent that retries model then tool interrupts", + Instruction: "You are a test agent.", + Model: mdl, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 2, + IsRetryAble: func(ctx context.Context, err error) bool { return errors.Is(err, errRetryAble) }, + }, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{interruptTool}, + }, + }, + }) + assert.NoError(t, err) + + store := newMyStore() + runner := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, []Message{schema.UserMessage("test query")}, WithCheckPointID("retry_interrupt_ckpt")) + + var events []*AgentEvent + var checkpointErr error + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + if event.Err != nil { + checkpointErr = event.Err + t.Logf("event error: %v", event.Err) + } + } + + // The bug: checkpoint save fails because the failed stream's error chunk + // is encountered during gob encoding of the session events. + // If the bug is fixed, checkpointErr should be nil and we should see an interrupt event. + assert.NoError(t, checkpointErr, "checkpoint save should not fail due to failed stream's error in session") + + var hasInterrupt bool + for _, event := range events { + if event.Action != nil && event.Action.Interrupted != nil { + hasInterrupt = true + } + } + assert.True(t, hasInterrupt, "should receive an interrupt event from the tool") + + // Verify the model was called twice (first call errored, second succeeded) + assert.Equal(t, int32(2), atomic.LoadInt32(&mdl.callCount), "model should be called exactly twice (1 failure + 1 retry)") +} diff --git a/adk/message_id_test.go b/adk/message_id_test.go new file mode 100644 index 0000000..ef85335 --- /dev/null +++ b/adk/message_id_test.go @@ -0,0 +1,1094 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +func isValidUUID(s string) bool { + // UUID v4 format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx (8-4-4-4-12 = 36 chars) + if len(s) != 36 { + return false + } + for i, c := range s { + if i == 8 || i == 13 || i == 18 || i == 23 { + if c != '-' { + return false + } + } else if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) { + return false + } + } + return true +} + +// collectEvents drains all events from the iterator (non-streaming). +func collectEvents(t *testing.T, iter *AsyncIterator[*AgentEvent]) []*AgentEvent { + t.Helper() + var events []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + events = append(events, event) + } + return events +} + +// Scenario 1: AgentEvent messages have IDs (Generate mode) +func TestMessageID_EventHasID_Generate(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("hello world", nil), nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + }) + + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + require.NotNil(t, events[0].Output.MessageOutput) + + msg := events[0].Output.MessageOutput.Message + require.NotNil(t, msg) + msgID := GetMessageID(msg) + assert.NotEmpty(t, msgID, "event message should have an ID") + assert.True(t, isValidUUID(msgID), "message ID should be a valid UUID, got: %s", msgID) +} + +// Scenario 2: Event and state messages share the same ID +func TestMessageID_EventAndStateShareSameID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("response", nil), nil). + Times(1) + + var stateMessagesAfterModel []*schema.Message + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + Middlewares: []AgentMiddleware{ + { + AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + // Capture state messages after model call (including the model output) + stateMessagesAfterModel = make([]*schema.Message, len(state.Messages)) + copy(stateMessagesAfterModel, state.Messages) + return nil + }, + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + }) + + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + + eventMsg := events[0].Output.MessageOutput.Message + eventMsgID := GetMessageID(eventMsg) + assert.NotEmpty(t, eventMsgID) + + // The last message in state should be the model output with the same ID + require.NotEmpty(t, stateMessagesAfterModel) + lastStateMsg := stateMessagesAfterModel[len(stateMessagesAfterModel)-1] + stateMsgID := GetMessageID(lastStateMsg) + + assert.Equal(t, eventMsgID, stateMsgID, + "event msg ID (%s) and state msg ID (%s) must match", eventMsgID, stateMsgID) +} + +// Scenario 3: Stream — first chunk carries ID, concatenated message has correct ID +func TestMessageID_Stream_FirstChunkOnly(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("chunk1", nil), + schema.AssistantMessage("chunk2", nil), + schema.AssistantMessage("chunk3", nil), + }), nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + EnableStreaming: true, + }) + + event, ok := iter.Next() + require.True(t, ok) + require.Nil(t, event.Err) + require.NotNil(t, event.Output.MessageOutput) + require.True(t, event.Output.MessageOutput.IsStreaming) + + stream := event.Output.MessageOutput.MessageStream + require.NotNil(t, stream) + + var chunks []*schema.Message + for { + msg, err := stream.Recv() + if err != nil { + break + } + chunks = append(chunks, msg) + } + require.GreaterOrEqual(t, len(chunks), 1) + + // First chunk should have the ID + firstChunkID := GetMessageID(chunks[0]) + assert.NotEmpty(t, firstChunkID, "first chunk should carry the message ID") + assert.True(t, isValidUUID(firstChunkID)) + + // Subsequent chunks should NOT have the ID in Extra (first-chunk-only injection) + for i := 1; i < len(chunks); i++ { + chunkID := GetMessageID(chunks[i]) + assert.Empty(t, chunkID, "chunk %d should not have message ID (first-chunk-only)", i) + } + + // No more events + _, ok = iter.Next() + assert.False(t, ok) +} + +// Scenario 4: Tool messages have IDs +func TestMessageID_ToolMessagesHaveID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolForTest{tarCount: 1} + info, err := fakeTool.Info(ctx) + require.NoError(t, err) + + generateCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tool", + []schema.ToolCall{{ + ID: "tc-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "tester"}`, + }, + }}), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Capture tool result messages from state via BeforeChatModel on the 2nd model call. + var toolMsgIDInState string + beforeModelCount := 0 + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + }, + Middlewares: []AgentMiddleware{ + { + BeforeChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + beforeModelCount++ + if beforeModelCount == 2 { + // 2nd model call: state.Messages contains tool result messages + for _, m := range state.Messages { + if m.Role == schema.Tool && m.ToolCallID == "tc-1" { + toolMsgIDInState = GetMessageID(m) + } + } + } + return nil + }, + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("use tool")}, + }) + + events := collectEvents(t, iter) + // Expect 3 events: model(tool_call) + tool(result) + model(final) + require.Len(t, events, 3) + + // Tool event (index 1) + toolEvent := events[1] + require.Nil(t, toolEvent.Err) + require.NotNil(t, toolEvent.Output.MessageOutput) + assert.Equal(t, schema.Tool, toolEvent.Output.MessageOutput.Role) + + toolMsg := toolEvent.Output.MessageOutput.Message + require.NotNil(t, toolMsg) + toolMsgID := GetMessageID(toolMsg) + assert.NotEmpty(t, toolMsgID, "tool message should have an ID") + assert.True(t, isValidUUID(toolMsgID)) + + // All events should have IDs + for i, ev := range events { + require.Nil(t, ev.Err) + require.NotNil(t, ev.Output.MessageOutput) + msg := ev.Output.MessageOutput.Message + require.NotNil(t, msg) + assert.NotEmpty(t, GetMessageID(msg), "event[%d] should have a message ID", i) + } + + // The tool message in state should share the same ID as the event tool message. + assert.NotEmpty(t, toolMsgIDInState, "tool message in state should have an ID") + assert.Equal(t, toolMsgID, toolMsgIDInState, + "tool event msg ID (%s) and state msg ID (%s) must match", toolMsgID, toolMsgIDInState) +} + +// Scenario 5: Retry — the final accepted result carries a message ID +func TestMessageID_Retry_FinalResultHasID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + retryErr := errors.New("retryable error") + + var callCount int32 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + count := atomic.AddInt32(&callCount, 1) + if count < 3 { + return nil, retryErr + } + return schema.AssistantMessage("Success after retry", nil), nil + }).Times(3) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + ShouldRetry: func(ctx context.Context, retryCtx *RetryContext) *RetryDecision { + return &RetryDecision{Retry: errors.Is(retryCtx.Err, retryErr)} + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hello")}, + }) + + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + + msg := events[0].Output.MessageOutput.Message + msgID := GetMessageID(msg) + assert.NotEmpty(t, msgID, "surviving message should have an ID") + assert.True(t, isValidUUID(msgID)) + assert.Equal(t, int32(3), atomic.LoadInt32(&callCount)) +} + +// Scenario 6: WrapModel handler sees model output with ID +func TestMessageID_WrapModelSeesID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("model output", nil), nil). + Times(1) + + var capturedMsgID string + + handler := &wrapModelIDCheckHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + onGenerate: func(result *schema.Message) { + capturedMsgID = GetMessageID(result) + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + Handlers: []ChatModelAgentMiddleware{handler}, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + }) + + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + + assert.NotEmpty(t, capturedMsgID, "WrapModel handler should see message ID on model output") + assert.True(t, isValidUUID(capturedMsgID)) + + // The event should carry the same ID + eventMsgID := GetMessageID(events[0].Output.MessageOutput.Message) + assert.Equal(t, capturedMsgID, eventMsgID, + "WrapModel-captured ID (%s) should match event ID (%s)", capturedMsgID, eventMsgID) +} + +// wrapModelIDCheckHandler wraps the model to inspect the output for message ID. +type wrapModelIDCheckHandler struct { + *BaseChatModelAgentMiddleware + onGenerate func(result *schema.Message) +} + +func (h *wrapModelIDCheckHandler) WrapModel(_ context.Context, m model.BaseChatModel, _ *ModelContext) (model.BaseChatModel, error) { + return &idCheckModelWrapper{inner: m, onGenerate: h.onGenerate}, nil +} + +type idCheckModelWrapper struct { + inner model.BaseChatModel + onGenerate func(result *schema.Message) +} + +func (w *idCheckModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + result, err := w.inner.Generate(ctx, input, opts...) + if err == nil && w.onGenerate != nil { + w.onGenerate(result) + } + return result, err +} + +func (w *idCheckModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return w.inner.Stream(ctx, input, opts...) +} + +// Scenario 7: User input messages do NOT get automatic IDs (they are external, not framework-created). +// Only framework-created messages (model output, tool results, TypedSendEvent) get auto-assigned IDs. +func TestMessageID_UserInputNoAutoID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + var stateMessagesBeforeModel []*schema.Message + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + // Capture input messages + stateMessagesBeforeModel = make([]*schema.Message, len(input)) + copy(stateMessagesBeforeModel, input) + return schema.AssistantMessage("response", nil), nil + }).Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hello")}, + }) + + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + + // User input messages should NOT have auto-assigned IDs. + // Framework only assigns IDs to messages it creates (model output, tool results, SendEvent). + require.NotEmpty(t, stateMessagesBeforeModel) + + for i, msg := range stateMessagesBeforeModel { + msgID := GetMessageID(msg) + assert.Empty(t, msgID, "input message[%d] (role=%s) should NOT have auto-assigned ID", i, msg.Role) + } +} + +// Scenario 8: Middleware must call EnsureMessageID before SendEvent; pointer identity ensures state consistency +// TestMessageID_SendEvent_MiddlewareMustEnsureID verifies that TypedSendEvent is a pure +// transport and does NOT auto-assign message IDs. Middleware authors must call +// EnsureMessageID themselves before sending. +func TestMessageID_SendEvent_MiddlewareMustEnsureID(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("model response", nil), nil). + Times(1) + + // Track the message pointer that the middleware creates and writes to both state and event + var middlewareMsg *schema.Message + var stateMsgIDAfterSendEvent string + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMsgID", + Instruction: "test", + Model: cm, + Middlewares: []AgentMiddleware{ + { + AfterChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + // Middleware creates a new message and writes the SAME pointer to both state and event + middlewareMsg = schema.AssistantMessage("middleware injected", nil) + + // Middleware is responsible for assigning the ID before sending + EnsureMessageID(middlewareMsg) + + // Write to state + state.Messages = append(state.Messages, middlewareMsg) + + // Send as event — TypedSendEvent does NOT auto-assign ID + event := EventFromMessage(middlewareMsg, nil, schema.Assistant, "") + err := SendEvent(ctx, event) + if err != nil { + return err + } + + // Because we called EnsureMessageID on the shared pointer, + // the state copy also has the ID (pointer identity) + stateMsgIDAfterSendEvent = internal.GetMessageID(middlewareMsg.Extra) + + return nil + }, + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + }) + + var allEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + allEvents = append(allEvents, event) + } + + // We expect at least 2 events: model response + middleware injected message + require.GreaterOrEqual(t, len(allEvents), 2) + + // The middleware message pointer should have an ID (assigned by middleware via EnsureMessageID) + require.NotNil(t, middlewareMsg) + middlewareMsgID := GetMessageID(middlewareMsg) + assert.NotEmpty(t, middlewareMsgID, "middleware should have assigned an ID via EnsureMessageID") + assert.True(t, isValidUUID(middlewareMsgID)) + + // The ID captured right after SendEvent (via pointer identity) should be the same + assert.Equal(t, middlewareMsgID, stateMsgIDAfterSendEvent, + "pointer identity: ID read from state pointer (%s) should match message ID (%s)", + stateMsgIDAfterSendEvent, middlewareMsgID) + + // Find the middleware event in the collected events + var middlewareEventMsgID string + for _, ev := range allEvents { + if ev.Err != nil || ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + msg := ev.Output.MessageOutput.Message + if msg != nil && msg.Content == "middleware injected" { + middlewareEventMsgID = GetMessageID(msg) + break + } + } + assert.Equal(t, middlewareMsgID, middlewareEventMsgID, + "event message ID (%s) should match the middleware message ID (%s)", + middlewareEventMsgID, middlewareMsgID) +} + +func TestAttack_ConcatCorruptsIDIfMultipleChunksCarryIt(t *testing.T) { + id := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + msgs := []*schema.Message{ + {Role: schema.Assistant, Content: "chunk1", Extra: map[string]any{internal.EinoMsgIDKey: id}}, + {Role: schema.Assistant, Content: "chunk2", Extra: map[string]any{internal.EinoMsgIDKey: id}}, + {Role: schema.Assistant, Content: "chunk3", Extra: map[string]any{internal.EinoMsgIDKey: id}}, + } + concatenated, err := schema.ConcatMessages(msgs) + require.NoError(t, err) + + resultID := internal.GetMessageID(concatenated.Extra) + // ConcatMessages string-concatenates duplicate Extra keys, corrupting the ID + assert.NotEqual(t, id, resultID, "ConcatMessages should corrupt the ID when multiple chunks carry it") + assert.NotEqual(t, 36, len(resultID), "corrupted ID should not be 36 chars") + assert.Equal(t, "chunk1chunk2chunk3", concatenated.Content) +} + +func TestAttack_ConcatPreservesIDIfOnlyFirstChunkHasIt(t *testing.T) { + id := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + msgs := []*schema.Message{ + {Role: schema.Assistant, Content: "chunk1", Extra: map[string]any{internal.EinoMsgIDKey: id}}, + {Role: schema.Assistant, Content: "chunk2"}, + {Role: schema.Assistant, Content: "chunk3"}, + } + concatenated, err := schema.ConcatMessages(msgs) + require.NoError(t, err) + + resultID := internal.GetMessageID(concatenated.Extra) + assert.Equal(t, id, resultID, "ID should be preserved when only first chunk carries it") + assert.Equal(t, "chunk1chunk2chunk3", concatenated.Content) +} + +func TestAttack_ConcurrentGenerate_NoSharedExtraMutation(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Shared singleton message - same pointer returned every time + sharedMsg := schema.AssistantMessage("shared response", nil) + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(sharedMsg, nil). + AnyTimes() + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAttack", + Instruction: "test", + Model: cm, + }) + require.NoError(t, err) + + const N = 10 + ids := make([]string, N) + var wg sync.WaitGroup + wg.Add(N) + for i := 0; i < N; i++ { + go func(idx int) { + defer wg.Done() + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + }) + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + msg := events[0].Output.MessageOutput.Message + require.NotNil(t, msg) + ids[idx] = GetMessageID(msg) + }(i) + } + wg.Wait() + + // All IDs should be unique and valid + seen := make(map[string]bool) + for i, id := range ids { + assert.NotEmpty(t, id, "goroutine %d should have an ID", i) + assert.True(t, isValidUUID(id), "goroutine %d ID should be valid UUID: %s", i, id) + assert.False(t, seen[id], "goroutine %d has duplicate ID: %s", i, id) + seen[id] = true + } + + // The original shared message should NOT have been mutated (or if it was, it should still be valid) + // The important thing is no panic and unique IDs +} + +func TestAttack_GenerateCopyDoesNotAffectOriginal(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + originalMsg := schema.AssistantMessage("original", nil) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(originalMsg, nil). + Times(1) + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAttack", + Instruction: "test", + Model: cm, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("hi")}, + }) + + events := collectEvents(t, iter) + require.Len(t, events, 1) + require.Nil(t, events[0].Err) + + eventMsg := events[0].Output.MessageOutput.Message + eventMsgID := GetMessageID(eventMsg) + assert.NotEmpty(t, eventMsgID) + + // The ORIGINAL message returned by the model should NOT have an ID + // because wrapGenerateEndpoint copies before mutating + originalID := GetMessageID(originalMsg) + assert.Empty(t, originalID, "original model output should NOT be mutated by ID assignment") +} + +// ============================================================ +// AgenticMessage Integration Tests +// ============================================================ + +// TestMessageID_AgenticGenerate verifies that AgenticMessage-typed agents +// get message IDs assigned on Generate output, covering the *schema.AgenticMessage +// branches in EnsureMessageID, GetMessageID, and copyMessage. +func TestMessageID_AgenticGenerate(t *testing.T) { + ctx := context.Background() + + agenticResponse := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "agentic response"}), + }, + } + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticResponse, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticMsgID", + Instruction: "test", + Model: m, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("hi")}, + }) + + event, ok := iter.Next() + require.True(t, ok) + require.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + require.NotNil(t, msg) + + // Verify via the AgenticMessage-specific public API + msgID := GetMessageID(msg) + assert.NotEmpty(t, msgID, "agentic model output should have message ID") + assert.True(t, isValidUUID(msgID), "agentic message ID should be valid UUID: %s", msgID) + + // Original message should NOT be mutated (copyMessage for AgenticMessage branch) + originalID := GetMessageID(agenticResponse) + assert.Empty(t, originalID, "original agentic model output should NOT be mutated") + + // Drain iterator + for { + _, ok := iter.Next() + if !ok { + break + } + } +} + +// TestMessageID_AgenticStream verifies first-chunk-only ID injection for AgenticMessage streams. +func TestMessageID_AgenticStream(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return nil, errors.New("should not be called") + }, + streamFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + r, w := schema.Pipe[*schema.AgenticMessage](3) + go func() { + defer w.Close() + for i := 0; i < 3; i++ { + w.Send(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "chunk"}), + }, + }, nil) + } + }() + return r, nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "AgenticStreamMsgID", + Instruction: "test", + Model: m, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &TypedAgentInput[*schema.AgenticMessage]{ + Messages: []*schema.AgenticMessage{schema.UserAgenticMessage("hi")}, + EnableStreaming: true, + }) + + event, ok := iter.Next() + require.True(t, ok) + require.Nil(t, event.Err) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.True(t, event.Output.MessageOutput.IsStreaming) + + stream := event.Output.MessageOutput.MessageStream + require.NotNil(t, stream) + + var streamMsgID string + for { + chunk, err := stream.Recv() + if err != nil { + break + } + chunkID := GetMessageID(chunk) + if streamMsgID == "" && chunkID != "" { + streamMsgID = chunkID + } else if chunkID != "" { + // Subsequent chunks should not have ID (first-chunk-only) + t.Errorf("expected only first chunk to have ID, got ID on later chunk: %s", chunkID) + } + } + + // Drain remaining events + for { + _, ok := iter.Next() + if !ok { + break + } + } + + assert.NotEmpty(t, streamMsgID, "first stream chunk should have message ID") + assert.True(t, isValidUUID(streamMsgID), "stream message ID should be valid UUID: %s", streamMsgID) +} + +// TestMessageID_AgenticPublicAPIHelpers tests the batch helpers and ensures +// the AgenticMessage public API variants work correctly. +func TestMessageID_AgenticPublicAPIHelpers(t *testing.T) { + t.Run("EnsureMessageID_idempotent", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "test"}), + }, + } + assert.Empty(t, GetMessageID(msg)) + + EnsureMessageID(msg) + id1 := GetMessageID(msg) + assert.NotEmpty(t, id1) + assert.True(t, isValidUUID(id1)) + + // Idempotent: second call should not change the ID + EnsureMessageID(msg) + id2 := GetMessageID(msg) + assert.Equal(t, id1, id2) + }) + + t.Run("EnsureMessageIDs_batch", func(t *testing.T) { + msgs := []*schema.AgenticMessage{ + {Role: schema.AgenticRoleTypeAssistant}, + {Role: schema.AgenticRoleTypeUser}, + {Role: schema.AgenticRoleTypeAssistant}, + } + for _, msg := range msgs { + EnsureMessageID(msg) + } + + seen := make(map[string]bool) + for i, msg := range msgs { + id := GetMessageID(msg) + assert.NotEmpty(t, id, "msg[%d] should have ID", i) + assert.True(t, isValidUUID(id), "msg[%d] ID should be valid UUID: %s", i, id) + assert.False(t, seen[id], "msg[%d] has duplicate ID: %s", i, id) + seen[id] = true + } + }) +} + +// --- Adversarial attack tests for message ID system --- + +// TestAttack_PopToolMsgID_DoublePop tests that calling popToolMsgID twice for the +// same key returns "" on second call. +func TestAttack_PopToolMsgID_DoublePop(t *testing.T) { + st := &typedState[*schema.Message]{} + st.setToolMsgID("myTool", "call-1", "uuid-abc") + + // First pop returns the ID + id1 := st.popToolMsgID("myTool", "call-1") + assert.Equal(t, "uuid-abc", id1) + + // Second pop returns empty + id2 := st.popToolMsgID("myTool", "call-1") + assert.Empty(t, id2, "double-pop should return empty") + + // Inner map should be cleaned up + assert.Nil(t, st.ToolMsgIDs["myTool"], "inner map should be removed when empty") +} + +// namedFakeToolForTest is a variant of fakeToolForTest with a configurable name. +type namedFakeToolForTest struct { + name string +} + +func (t *namedFakeToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: t.name + " tool for testing", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Desc: "user name for testing", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (t *namedFakeToolForTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return `{"say": "ok"}`, nil +} + +// TestAttack_ToolMsgIDConsistency_MultipleTools is an integration test: when an agent +// has multiple tools called in one turn, verify that EACH tool's event message ID +// matches its corresponding state message ID. +func TestAttack_ToolMsgIDConsistency_MultipleTools(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + tool1 := &namedFakeToolForTest{name: "greet"} + tool2 := &namedFakeToolForTest{name: "farewell"} + + info1, err := tool1.Info(ctx) + require.NoError(t, err) + info2, err := tool2.Info(ctx) + require.NoError(t, err) + + var generateCount int + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + generateCount++ + if generateCount == 1 { + return schema.AssistantMessage("calling tools", []schema.ToolCall{ + {ID: "tc-1", Function: schema.FunctionCall{Name: info1.Name, Arguments: `{"name": "alice"}`}}, + {ID: "tc-2", Function: schema.FunctionCall{Name: info2.Name, Arguments: `{"name": "bob"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Capture state message IDs + var stateMsgIDs map[string]string // callID -> msgID + beforeModelCount := 0 + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestMultiTool", + Instruction: "test", + Model: cm, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1, tool2}, + }, + }, + Middlewares: []AgentMiddleware{ + { + BeforeChatModel: func(ctx context.Context, state *ChatModelAgentState) error { + beforeModelCount++ + if beforeModelCount == 2 { + stateMsgIDs = make(map[string]string) + for _, m := range state.Messages { + if m.Role == schema.Tool { + stateMsgIDs[m.ToolCallID] = GetMessageID(m) + } + } + } + return nil + }, + }, + }, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &AgentInput{ + Messages: []Message{schema.UserMessage("use tools")}, + }) + + events := collectEvents(t, iter) + // Expect: model(tool_calls) + tool1(result) + tool2(result) + model(final) = 4 events + require.GreaterOrEqual(t, len(events), 4) + + // Collect tool event IDs + eventMsgIDs := make(map[string]string) // callID -> msgID + for _, ev := range events { + if ev.Err != nil { + continue + } + if ev.Output != nil && ev.Output.MessageOutput != nil { + msg := ev.Output.MessageOutput.Message + if msg != nil && msg.Role == schema.Tool { + eventMsgIDs[msg.ToolCallID] = GetMessageID(msg) + } + } + } + + // Each tool call should have an ID in both event and state, and they must match + require.NotEmpty(t, stateMsgIDs, "state should have tool message IDs") + for callID, stateID := range stateMsgIDs { + assert.NotEmpty(t, stateID, "state msg for %s should have ID", callID) + assert.True(t, isValidUUID(stateID), "state msg ID should be UUID: %s", stateID) + eventID, ok := eventMsgIDs[callID] + assert.True(t, ok, "event should have msg for callID %s", callID) + assert.Equal(t, stateID, eventID, + "event and state msg IDs for callID %s must match: event=%s state=%s", callID, eventID, stateID) + } +} + +// TestAttack_ToolResultToBlocks_EdgeCases verifies toolResultToBlocks handles +// nil ToolResult, empty Parts, and Parts with nil media fields. +func TestAttack_ToolResultToBlocks_EdgeCases(t *testing.T) { + t.Run("nil ToolResult", func(t *testing.T) { + blocks := toolResultToBlocks(nil) + assert.Nil(t, blocks, "nil ToolResult should produce nil blocks") + }) + + t.Run("empty Parts", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{}} + blocks := toolResultToBlocks(tr) + assert.Nil(t, blocks, "empty Parts should produce nil blocks") + }) + + t.Run("text part with empty text", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: ""}, + }} + blocks := toolResultToBlocks(tr) + require.Len(t, blocks, 1) + assert.NotNil(t, blocks[0].Text) + assert.Equal(t, "", blocks[0].Text.Text) + }) + + t.Run("image part with nil Image field", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeImage, Image: nil}, + }} + blocks := toolResultToBlocks(tr) + assert.Empty(t, blocks) + }) + + t.Run("audio part with nil Audio field", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeAudio, Audio: nil}, + }} + blocks := toolResultToBlocks(tr) + assert.Empty(t, blocks) + }) + + t.Run("video part with nil Video field", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeVideo, Video: nil}, + }} + blocks := toolResultToBlocks(tr) + assert.Empty(t, blocks) + }) + + t.Run("file part with nil File field", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeFile, File: nil}, + }} + blocks := toolResultToBlocks(tr) + assert.Empty(t, blocks) + }) + + t.Run("mixed: valid text + nil image + valid text", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "hello"}, + {Type: schema.ToolPartTypeImage, Image: nil}, + {Type: schema.ToolPartTypeText, Text: "world"}, + }} + blocks := toolResultToBlocks(tr) + require.Len(t, blocks, 2) + assert.Equal(t, "hello", blocks[0].Text.Text) + assert.Equal(t, "world", blocks[1].Text.Text) + }) + + t.Run("image part with nil URL pointers", func(t *testing.T) { + tr := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeImage, Image: &schema.ToolOutputImage{ + MessagePartCommon: schema.MessagePartCommon{ + URL: nil, + Base64Data: nil, + MIMEType: "image/png", + }, + }}, + }} + blocks := toolResultToBlocks(tr) + require.Len(t, blocks, 1) + assert.NotNil(t, blocks[0].Image) + assert.Equal(t, "", blocks[0].Image.URL, "nil URL pointer should deref to empty string") + assert.Equal(t, "", blocks[0].Image.Base64Data) + assert.Equal(t, "image/png", blocks[0].Image.MIMEType) + }) +} diff --git a/adk/middlewares/agentsmd/agentsmd.go b/adk/middlewares/agentsmd/agentsmd.go new file mode 100644 index 0000000..5339998 --- /dev/null +++ b/adk/middlewares/agentsmd/agentsmd.go @@ -0,0 +1,213 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package agentsmd provides a middleware that automatically injects Agents.md +// file contents into model input messages. The injection is transient — content +// is prepended at model call time and never persisted to conversation state, +// so it is naturally excluded from summarization / compression. +package agentsmd + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// Config defines the configuration for the agentsmd middleware. +type Config struct { + // Backend provides file access for loading Agents.md files. + // Implementations can use local filesystem, remote storage, or any other backend. + // Required. + Backend Backend + + // AgentsMDFiles specifies the ordered list of Agents.md file paths to load. + // Files are loaded and injected in the given order. + // Supports @import syntax inside files for recursive inclusion (max depth 5). + AgentsMDFiles []string + + // AllAgentsMDMaxBytes limits the total byte size of all loaded Agents.md content. + // Files are loaded in order; once the cumulative size exceeds this limit, + // remaining files are skipped. Each individual file is always loaded in full. + // 0 means no limit. + AllAgentsMDMaxBytes int + + // OnLoadWarning is an optional callback invoked when a non-fatal error occurs + // during Agents.md file loading (e.g. file not found, circular @import, depth + // exceeded). If nil, warnings are logged via log.Printf. + // + // Note: Backend.Read errors other than os.ErrNotExist (e.g. permission denied, + // I/O errors) are NOT treated as warnings and will abort the loading process. + OnLoadWarning func(filePath string, err error) +} + +// NewTyped creates a generic agentsmd middleware that injects Agents.md content into every +// model call. The content is loaded from the configured file paths via Backend +// on each model invocation. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +// +// Recommended: place this middleware AFTER the summarization middleware, so that +// Agents.md content is excluded from summarization/compression. +func NewTyped[M adk.MessageType](_ context.Context, cfg *Config) (adk.TypedChatModelAgentMiddleware[M], error) { + if err := cfg.validate(); err != nil { + return nil, err + } + + return &typedMiddleware[M]{ + loader: newLoaderConfig(cfg.Backend, cfg.AgentsMDFiles, cfg.AllAgentsMDMaxBytes, cfg.OnLoadWarning), + }, nil +} + +// New creates an agentsmd middleware that injects Agents.md content into every +// model call. The content is loaded from the configured file paths via Backend +// on each model invocation. +// +// Recommended: place this middleware AFTER the summarization middleware, so that +// Agents.md content is excluded from summarization/compression. +func New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, cfg) +} + +type typedMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + loader *loaderConfig +} + +const agentsMDCacheKey = "__agentsmd_content_cache__" +const agentsMDExtraKey = "__agentsmd_content__" + +// BeforeModelRewriteState injects Agents.md content as a User message before +// the first User message in the conversation. The injected message is tagged +// with an Extra key so that repeated invocations are idempotent. +func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { + // Idempotent: if we already injected, return early. + for _, msg := range state.Messages { + if hasAgentsMDExtra(msg) { + return ctx, state, nil + } + } + + content, err := m.loadContent(ctx) + if err != nil { + return ctx, nil, err + } + if content == "" { + return ctx, state, nil + } + + nState := *state + nState.Messages = typedInsertBeforeFirstUser(state.Messages, content) + return ctx, &nState, nil +} + +// hasAgentsMDExtra checks whether a message has the agentsmd extra key set. +func hasAgentsMDExtra[M adk.MessageType](msg M) bool { + switch v := any(msg).(type) { + case *schema.Message: + if v.Extra != nil { + if _, ok := v.Extra[agentsMDExtraKey]; ok { + return true + } + } + case *schema.AgenticMessage: + if v.Extra != nil { + if _, ok := v.Extra[agentsMDExtraKey]; ok { + return true + } + } + } + return false +} + +// typedInsertBeforeFirstUser inserts a user message with agentsmd content before the first User message. +func typedInsertBeforeFirstUser[M adk.MessageType](msgs []M, content string) []M { + newMsg := makeUserMsgWithExtra[M](content) + result := make([]M, 0, len(msgs)+1) + for i, msg := range msgs { + if isUserRole(msg) { + result = append(result, newMsg) + result = append(result, msgs[i:]...) + return result + } + result = append(result, msg) + } + result = append(result, newMsg) + return result +} + +func isUserRole[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.User + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeUser + } + return false +} + +func makeUserMsgWithExtra[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.UserMessage(content) + msg.Extra = map[string]any{agentsMDExtraKey: true} + return any(msg).(M) + case *schema.AgenticMessage: + msg := schema.UserAgenticMessage(content) + msg.Extra = map[string]any{agentsMDExtraKey: true} + return any(msg).(M) + } + panic("unreachable") +} + +// loadContent retrieves the Agents.md content, using a per-Run cache to avoid +// reloading on every model call within the same Run(). +func (m *typedMiddleware[M]) loadContent(ctx context.Context) (string, error) { + if cached, found, err := adk.GetRunLocalValue(ctx, agentsMDCacheKey); err == nil && found { + if s, ok := cached.(string); ok { + return s, nil + } + } + + content, err := m.loader.load(ctx) + if err != nil { + return "", fmt.Errorf("[agentsmd]: failed to load agent files: %w", err) + } + + if content != "" { + _ = adk.SetRunLocalValue(ctx, agentsMDCacheKey, content) + } + + return content, nil +} + +func (c *Config) validate() error { + if c == nil { + return fmt.Errorf("[agentsmd]: config is required") + } + if c.Backend == nil { + return fmt.Errorf("[agentsmd]: backend is required") + } + if len(c.AgentsMDFiles) == 0 { + return fmt.Errorf("[agentsmd]: at least one agent file path is required") + } + if c.AllAgentsMDMaxBytes < 0 { + return fmt.Errorf("[agentsmd]: AllAgentMDDocsMaxBytes must be non-negative") + } + return nil +} diff --git a/adk/middlewares/agentsmd/agentsmd_generic_test.go b/adk/middlewares/agentsmd/agentsmd_generic_test.go new file mode 100644 index 0000000..28e2569 --- /dev/null +++ b/adk/middlewares/agentsmd/agentsmd_generic_test.go @@ -0,0 +1,351 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package agentsmd + +import ( + "context" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// --- generic table-driven test helpers --- + +func makeUserMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{Role: schema.User, Content: content}).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(content)).(M) + default: + panic("unreachable") + } +} + +func makeSystemMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{Role: schema.System, Content: content}).(M) + case *schema.AgenticMessage: + return any(schema.SystemAgenticMessage(content)).(M) + default: + panic("unreachable") + } +} + +func makeAssistantMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{Role: schema.Assistant, Content: content}).(M) + case *schema.AgenticMessage: + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: content})}, + }).(M) + default: + panic("unreachable") + } +} + +func getMsgRole[M adk.MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return string(v.Role) + case *schema.AgenticMessage: + return string(v.Role) + default: + panic("unreachable") + } +} + +func getMsgContent[M adk.MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return v.Content + case *schema.AgenticMessage: + for _, block := range v.ContentBlocks { + if block == nil { + continue + } + if block.UserInputText != nil { + return block.UserInputText.Text + } + if block.AssistantGenText != nil { + return block.AssistantGenText.Text + } + } + return "" + default: + panic("unreachable") + } +} + +func getMsgExtra[M adk.MessageType](msg M) map[string]any { + switch v := any(msg).(type) { + case *schema.Message: + return v.Extra + case *schema.AgenticMessage: + return v.Extra + default: + panic("unreachable") + } +} + +// --- generic table-driven test --- + +type agentsMDTestCase struct { + name string + run func(t *testing.T) +} + +func testAgentsMDGeneric[M adk.MessageType](t *testing.T) { + tests := []agentsMDTestCase{ + { + name: "BasicInjection", + run: func(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "You are a helpful assistant.") + + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(state.Messages)) + } + if getMsgRole(state.Messages[0]) != "user" { + t.Fatalf("expected first message role user, got %s", getMsgRole(state.Messages[0])) + } + if !strings.Contains(getMsgContent(state.Messages[0]), "You are a helpful assistant.") { + t.Fatalf("expected agent.md content in first message, got %q", getMsgContent(state.Messages[0])) + } + if !strings.Contains(getMsgContent(state.Messages[0]), "") { + t.Fatalf("expected system-reminder tag, got %q", getMsgContent(state.Messages[0])) + } + if count := strings.Count(getMsgContent(state.Messages[0]), ""); count != 1 { + t.Fatalf("expected exactly one opening system-reminder tag, got %d in %q", count, getMsgContent(state.Messages[0])) + } + if count := strings.Count(getMsgContent(state.Messages[0]), ""); count != 1 { + t.Fatalf("expected exactly one closing system-reminder tag, got %d in %q", count, getMsgContent(state.Messages[0])) + } + if getMsgContent(state.Messages[1]) != "hello" { + t.Fatalf("expected original message preserved, got %q", getMsgContent(state.Messages[1])) + } + }, + }, + { + name: "InsertBeforeFirstUserMessage", + run: func(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + input := []M{ + makeSystemMsg[M]("system prompt"), + makeUserMsg[M]("hello"), + } + state := &adk.TypedChatModelAgentState[M]{Messages: input} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + if len(state.Messages) != 3 { + t.Fatalf("expected 3 messages, got %d", len(state.Messages)) + } + if getMsgRole(state.Messages[0]) != "system" { + t.Fatalf("expected first message role system, got %s", getMsgRole(state.Messages[0])) + } + if getMsgContent(state.Messages[0]) != "system prompt" { + t.Fatalf("expected system prompt preserved, got %q", getMsgContent(state.Messages[0])) + } + if getMsgRole(state.Messages[1]) != "user" || !strings.Contains(getMsgContent(state.Messages[1]), "agent instructions") { + t.Fatalf("expected agentmd message at index 1, got role=%s content=%q", getMsgRole(state.Messages[1]), getMsgContent(state.Messages[1])) + } + if getMsgRole(state.Messages[2]) != "user" || getMsgContent(state.Messages[2]) != "hello" { + t.Fatalf("expected original user message at index 2, got role=%s content=%q", getMsgRole(state.Messages[2]), getMsgContent(state.Messages[2])) + } + }, + }, + { + name: "InsertWithNoUserMessage", + run: func(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + input := []M{ + makeSystemMsg[M]("system prompt"), + makeAssistantMsg[M]("assistant reply"), + } + state := &adk.TypedChatModelAgentState[M]{Messages: input} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + if len(state.Messages) != 3 { + t.Fatalf("expected 3 messages, got %d", len(state.Messages)) + } + if getMsgRole(state.Messages[0]) != "system" { + t.Fatalf("expected System at index 0, got %s", getMsgRole(state.Messages[0])) + } + if getMsgRole(state.Messages[1]) != "assistant" { + t.Fatalf("expected Assistant at index 1, got %s", getMsgRole(state.Messages[1])) + } + if getMsgRole(state.Messages[2]) != "user" || !strings.Contains(getMsgContent(state.Messages[2]), "agent instructions") { + t.Fatalf("expected agentmd appended at end, got role=%s content=%q", getMsgRole(state.Messages[2]), getMsgContent(state.Messages[2])) + } + }, + }, + { + name: "AllFilesEmpty", + run: func(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "") + + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 1 { + t.Fatalf("expected 1 message (no agentmd prepended), got %d", len(state.Messages)) + } + if getMsgContent(state.Messages[0]) != "hello" { + t.Fatalf("expected original message unchanged, got %q", getMsgContent(state.Messages[0])) + } + }, + }, + { + name: "Idempotency", + run: func(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after first call, got %d", len(state.Messages)) + } + + // Verify the marker is set in Extra. + extra := getMsgExtra(state.Messages[0]) + if extra == nil { + t.Fatal("expected Extra to be set on injected message") + } + if _, ok := extra[agentsMDExtraKey]; !ok { + t.Fatalf("expected agentsMDExtraKey in Extra, got %v", extra) + } + + // Call again with the same state (which now contains the marker message). + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after second call (idempotent), got %d", len(state.Messages)) + } + if !strings.Contains(getMsgContent(state.Messages[0]), "agent instructions") { + t.Fatalf("expected agentmd content preserved, got %q", getMsgContent(state.Messages[0])) + } + }, + }, + { + name: "ReinsertAfterRemoval", + run: func(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after first call, got %d", len(state.Messages)) + } + + // Simulate removal of the marker message (e.g., by summarization). + state = &adk.TypedChatModelAgentState[M]{Messages: []M{makeUserMsg[M]("hello")}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after re-insert, got %d", len(state.Messages)) + } + if !strings.Contains(getMsgContent(state.Messages[0]), "agent instructions") { + t.Fatalf("expected agentmd content re-inserted, got %q", getMsgContent(state.Messages[0])) + } + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, tc.run) + } +} + +func TestAgentsMDGeneric(t *testing.T) { + t.Run("Message", testAgentsMDGeneric[*schema.Message]) + t.Run("AgenticMessage", testAgentsMDGeneric[*schema.AgenticMessage]) +} diff --git a/adk/middlewares/agentsmd/agentsmd_test.go b/adk/middlewares/agentsmd/agentsmd_test.go new file mode 100644 index 0000000..5d016e1 --- /dev/null +++ b/adk/middlewares/agentsmd/agentsmd_test.go @@ -0,0 +1,1367 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package agentsmd + +import ( + "context" + "fmt" + "os" + "strings" + "testing" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/schema" +) + +// --- test helpers --- + +type memBackend struct { + files map[string]string +} + +func newMemBackend() *memBackend { + return &memBackend{files: make(map[string]string)} +} + +func (b *memBackend) set(path string, content string) { + b.files[path] = content +} + +func (b *memBackend) Read(_ context.Context, req *ReadRequest) (*filesystem.FileContent, error) { + content, ok := b.files[req.FilePath] + if !ok { + return nil, fmt.Errorf("file not found: %s: %w", req.FilePath, os.ErrNotExist) + } + return &filesystem.FileContent{Content: content}, nil +} + +// errBackend always returns a non-ErrNotExist error on Read, simulating I/O failures. +type errBackend struct{} + +func (b *errBackend) Read(_ context.Context, req *ReadRequest) (*filesystem.FileContent, error) { + return nil, fmt.Errorf("permission denied: %s", req.FilePath) +} + +// partialErrBackend returns content for known files and I/O error for others. +type partialErrBackend struct { + files map[string]string +} + +func (b *partialErrBackend) Read(_ context.Context, req *ReadRequest) (*filesystem.FileContent, error) { + content, ok := b.files[req.FilePath] + if !ok { + return nil, fmt.Errorf("I/O error reading %s", req.FilePath) + } + return &filesystem.FileContent{Content: content}, nil +} + +// --- tests --- + +func TestNew_Validation(t *testing.T) { + ctx := context.Background() + b := newMemBackend() + + _, err := New(ctx, nil) + if err == nil { + t.Fatal("expected error for nil config") + } + + _, err = New(ctx, &Config{}) + if err == nil { + t.Fatal("expected error for empty config") + } + + _, err = New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/test.md"}, AllAgentsMDMaxBytes: -1}) + if err == nil { + t.Fatal("expected error for negative max bytes") + } + + _, err = New(ctx, &Config{AgentsMDFiles: []string{"/test.md"}}) + if err == nil { + t.Fatal("expected error for nil backend") + } +} + +func TestMiddleware_BasicInjection(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "You are a helpful assistant.") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + userMsg := &schema.Message{Role: schema.User, Content: "hello"} + state := &adk.ChatModelAgentState{Messages: []*schema.Message{userMsg}} + + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(state.Messages)) + } + if state.Messages[0].Role != schema.User { + t.Fatalf("expected first message role User, got %s", state.Messages[0].Role) + } + if !strings.Contains(state.Messages[0].Content, "You are a helpful assistant.") { + t.Fatalf("expected agent.md content in first message, got %q", state.Messages[0].Content) + } + if !strings.Contains(state.Messages[0].Content, "") { + t.Fatalf("expected system-reminder tag, got %q", state.Messages[0].Content) + } + if count := strings.Count(state.Messages[0].Content, ""); count != 1 { + t.Fatalf("expected exactly one opening system-reminder tag, got %d in %q", count, state.Messages[0].Content) + } + if count := strings.Count(state.Messages[0].Content, ""); count != 1 { + t.Fatalf("expected exactly one closing system-reminder tag, got %d in %q", count, state.Messages[0].Content) + } + if state.Messages[1].Content != "hello" { + t.Fatalf("expected original message preserved, got %q", state.Messages[1].Content) + } +} + +func TestMiddleware_MultipleFiles(t *testing.T) { + b := newMemBackend() + b.set("/a.md", "instruction A") + b.set("/b.md", "instruction B") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/a.md", "/b.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + content := state.Messages[0].Content + idxA := strings.Index(content, "instruction A") + idxB := strings.Index(content, "instruction B") + if idxA < 0 || idxB < 0 { + t.Fatalf("both files should be included, content: %q", content) + } + if idxA >= idxB { + t.Fatal("file A should appear before file B") + } +} + +func TestMiddleware_ImportResolution(t *testing.T) { + b := newMemBackend() + b.set("/project/agent.md", "main instructions\n@sub/rules.md\nend") + b.set("/project/sub/rules.md", "imported rule") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/project/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + content := state.Messages[0].Content + // Original text should be preserved with @path intact. + if !strings.Contains(content, "main instructions") { + t.Fatalf("should contain original text, got %q", content) + } + if !strings.Contains(content, "@sub/rules.md") { + t.Fatalf("@import reference should be preserved in original text, got %q", content) + } + if !strings.Contains(content, "end") { + t.Fatalf("should contain original trailing text, got %q", content) + } + // Imported file should appear as a separate section. + if !strings.Contains(content, "Contents of /project/sub/rules.md") { + t.Fatalf("imported file should have its own section, got %q", content) + } + if !strings.Contains(content, "imported rule") { + t.Fatalf("imported file content should be present, got %q", content) + } +} + +func TestMiddleware_RecursiveImport(t *testing.T) { + b := newMemBackend() + b.set("/a.md", "top\n@/b.md") + b.set("/b.md", "middle\n@/c.md") + b.set("/c.md", "leaf content") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/a.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + content := state.Messages[0].Content + // All three files should appear as separate sections. + for _, section := range []string{"Contents of /a.md", "Contents of /b.md", "Contents of /c.md"} { + if !strings.Contains(content, section) { + t.Fatalf("expected section %q in content, got %q", section, content) + } + } + for _, text := range []string{"top", "middle", "leaf content"} { + if !strings.Contains(content, text) { + t.Fatalf("expected %q in content, got %q", text, content) + } + } + // Sections should appear in order: a, b, c. + idxA := strings.Index(content, "Contents of /a.md") + idxB := strings.Index(content, "Contents of /b.md") + idxC := strings.Index(content, "Contents of /c.md") + if !(idxA < idxB && idxB < idxC) { + t.Fatalf("sections should appear in order a < b < c, got a=%d b=%d c=%d", idxA, idxB, idxC) + } +} + +func TestMiddleware_MaxImportDepth(t *testing.T) { + b := newMemBackend() + for i := 0; i < 7; i++ { + var content string + if i < 6 { + content = fmt.Sprintf("level %d\n@/level%d.md", i, i+1) + } else { + content = fmt.Sprintf("level %d", i) + } + b.set(fmt.Sprintf("/level%d.md", i), content) + } + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/level0.md"}}) + if err != nil { + t.Fatal(err) + } + + // Import failure at depth > 5 is logged, not returned as error. + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatalf("expected no error (depth exceeded is logged), got %v", err) + } + // Levels 0-5 should be present as sections; level 6 fails silently. + content := state.Messages[0].Content + for i := 0; i <= 5; i++ { + want := fmt.Sprintf("Contents of /level%d.md", i) + if !strings.Contains(content, want) { + t.Fatalf("expected %q in content, got %q", want, content) + } + } + if strings.Contains(content, "Contents of /level6.md") { + t.Fatalf("level6 should not be present (depth exceeded), got %q", content) + } +} + +func TestMiddleware_CircularImport(t *testing.T) { + b := newMemBackend() + b.set("/a.md", "@/b.md") + b.set("/b.md", "@/a.md") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/a.md"}}) + if err != nil { + t.Fatal(err) + } + + // Circular import failure is logged, not returned as error. + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatalf("expected no error (circular import is logged), got %v", err) + } + // /a.md and /b.md should both be present; the circular ref from b->a is skipped. + content := state.Messages[0].Content + if !strings.Contains(content, "Contents of /a.md") { + t.Fatalf("expected /a.md section, got %q", content) + } + if !strings.Contains(content, "Contents of /b.md") { + t.Fatalf("expected /b.md section, got %q", content) + } +} + +func TestMiddleware_MaxBytesLimit(t *testing.T) { + b := newMemBackend() + b.set("/a.md", "AAAA") // 4 bytes + b.set("/b.md", "BBBB") // 4 bytes + + ctx := context.Background() + mw, err := New(ctx, &Config{ + Backend: b, + AgentsMDFiles: []string{"/a.md", "/b.md"}, + AllAgentsMDMaxBytes: 5, // file a (4) fits, file b (4) would exceed + }) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + content := state.Messages[0].Content + if !strings.Contains(content, "AAAA") { + t.Fatal("first file should be included") + } + if strings.Contains(content, "BBBB") { + t.Fatal("second file should be excluded due to max bytes") + } +} + +func TestMiddleware_InjectedInState(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + originalMsgs := []*schema.Message{{Role: schema.User, Content: "hello"}} + state := &adk.ChatModelAgentState{Messages: originalMsgs} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + // The original slice should not be modified (new slice is returned). + if len(originalMsgs) != 1 { + t.Fatalf("original messages slice should not be modified, got %d messages", len(originalMsgs)) + } + if originalMsgs[0].Content != "hello" { + t.Fatalf("original message should be unchanged, got %q", originalMsgs[0].Content) + } + // The returned state should have the injected message. + if len(newState.Messages) != 2 { + t.Fatalf("new state should have 2 messages (injected + original), got %d", len(newState.Messages)) + } + if !strings.Contains(newState.Messages[0].Content, "agent instructions") { + t.Fatalf("expected agentmd content in first message, got %q", newState.Messages[0].Content) + } + if newState.Messages[1].Content != "hello" { + t.Fatalf("expected original user message preserved, got %q", newState.Messages[1].Content) + } +} + +func TestMiddleware_AbsoluteImportPath(t *testing.T) { + b := newMemBackend() + b.set("/project/main.md", "start\n@/shared/imported.md\nend") + b.set("/shared/imported.md", "absolute import content") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/project/main.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + content := state.Messages[0].Content + // @path preserved in original text. + if !strings.Contains(content, "@/shared/imported.md") { + t.Fatalf("@import reference should be preserved, got %q", content) + } + // Imported content in separate section. + if !strings.Contains(content, "Contents of /shared/imported.md") { + t.Fatalf("expected separate section for imported file, got %q", content) + } + if !strings.Contains(content, "absolute import content") { + t.Fatalf("expected absolute import content, got %q", content) + } +} + +func TestMiddleware_ImportAsSeparateSection(t *testing.T) { + b := newMemBackend() + b.set("/project/agent.md", "Please read @sub/rules.md and also @sub/style.md for guidance.") + b.set("/project/sub/rules.md", "RULE_CONTENT") + b.set("/project/sub/style.md", "STYLE_CONTENT") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/project/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + content := state.Messages[0].Content + // Original text preserved with @paths intact. + if !strings.Contains(content, "Please read @sub/rules.md and also @sub/style.md for guidance.") { + t.Fatalf("original text with @paths should be preserved, got %q", content) + } + // Imported files appear as separate sections. + if !strings.Contains(content, "Contents of /project/sub/rules.md") { + t.Fatalf("expected rules.md section, got %q", content) + } + if !strings.Contains(content, "RULE_CONTENT") { + t.Fatalf("expected imported rule content, got %q", content) + } + if !strings.Contains(content, "Contents of /project/sub/style.md") { + t.Fatalf("expected style.md section, got %q", content) + } + if !strings.Contains(content, "STYLE_CONTENT") { + t.Fatalf("expected imported style content, got %q", content) + } + + // Sections should be ordered: agent.md, rules.md, style.md. + idxAgent := strings.Index(content, "Contents of /project/agent.md") + idxRules := strings.Index(content, "Contents of /project/sub/rules.md") + idxStyle := strings.Index(content, "Contents of /project/sub/style.md") + if !(idxAgent < idxRules && idxRules < idxStyle) { + t.Fatalf("sections should appear in order agent < rules < style, got agent=%d rules=%d style=%d", idxAgent, idxRules, idxStyle) + } +} + +// --- loader-specific tests --- + +func TestLoader_NoImportsPassthrough(t *testing.T) { + // Content without any @path should be returned as-is in its section. + b := newMemBackend() + b.set("/agent.md", "plain text without imports\nline two") + + l := newLoaderConfig(b, []string{"/agent.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "plain text without imports") { + t.Fatalf("expected plain content, got %q", content) + } + if !strings.Contains(content, "line two") { + t.Fatalf("expected second line, got %q", content) + } +} + +func TestLoader_ImportAsSeparateSection(t *testing.T) { + // @path in the middle of a sentence should be preserved; imported file is a separate section. + b := newMemBackend() + b.set("/doc.md", "before @/snippet.md after") + b.set("/snippet.md", "INJECTED") + + l := newLoaderConfig(b, []string{"/doc.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + // Original text preserved. + if !strings.Contains(content, "before @/snippet.md after") { + t.Fatalf("original text should be preserved with @path, got %q", content) + } + // Imported file in separate section. + if !strings.Contains(content, "Contents of /snippet.md") { + t.Fatalf("expected separate section for snippet.md, got %q", content) + } + if !strings.Contains(content, "INJECTED") { + t.Fatalf("expected imported content, got %q", content) + } +} + +func TestLoader_MultipleImportsSameLine(t *testing.T) { + // Multiple @path on one line should each get a separate section. + b := newMemBackend() + b.set("/doc.md", "see @/a.txt and @/b.txt here") + b.set("/a.txt", "AAA") + b.set("/b.txt", "BBB") + + l := newLoaderConfig(b, []string{"/doc.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + // Original text preserved. + if !strings.Contains(content, "see @/a.txt and @/b.txt here") { + t.Fatalf("original text should be preserved, got %q", content) + } + // Each imported file has its own section. + if !strings.Contains(content, "Contents of /a.txt") { + t.Fatalf("expected section for a.txt, got %q", content) + } + if !strings.Contains(content, "AAA") { + t.Fatalf("expected a.txt content, got %q", content) + } + if !strings.Contains(content, "Contents of /b.txt") { + t.Fatalf("expected section for b.txt, got %q", content) + } + if !strings.Contains(content, "BBB") { + t.Fatalf("expected b.txt content, got %q", content) + } +} + +func TestLoader_SameFileTwiceOnSameLine(t *testing.T) { + // The same file referenced twice should appear only once as a section (deduped). + b := newMemBackend() + b.set("/doc.md", "@/shared.md and @/shared.md again") + b.set("/shared.md", "SHARED") + + l := newLoaderConfig(b, []string{"/doc.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + // Original text preserved. + if !strings.Contains(content, "@/shared.md and @/shared.md again") { + t.Fatalf("original text should be preserved, got %q", content) + } + // shared.md content should appear only once (deduped). + count := strings.Count(content, "Contents of /shared.md") + if count != 1 { + t.Fatalf("expected shared.md section to appear once (deduped), got %d in %q", count, content) + } +} + +func TestLoader_ImportFileNotFound(t *testing.T) { + b := newMemBackend() + b.set("/doc.md", "load @/missing.md please") + + l := newLoaderConfig(b, []string{"/doc.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error (missing import is logged), got %v", err) + } + // Original text preserved; missing file simply has no section. + if !strings.Contains(content, "load @/missing.md please") { + t.Fatalf("expected original text preserved, got %q", content) + } + if strings.Contains(content, "Contents of /missing.md") { + t.Fatalf("missing file should not have a section, got %q", content) + } +} + +func TestLoader_RelativePathResolution(t *testing.T) { + // Relative path should resolve relative to the host file's directory. + b := newMemBackend() + b.set("/a/b/host.md", "ref @../c/target.md done") + b.set("/a/c/target.md", "TARGET") + + l := newLoaderConfig(b, []string{"/a/b/host.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + // Original text preserved. + if !strings.Contains(content, "ref @../c/target.md done") { + t.Fatalf("original text should be preserved, got %q", content) + } + // Imported file as separate section. + if !strings.Contains(content, "Contents of /a/c/target.md") { + t.Fatalf("expected section for target.md, got %q", content) + } + if !strings.Contains(content, "TARGET") { + t.Fatalf("expected imported content, got %q", content) + } +} + +func TestLoader_RelativeTopLevelPath(t *testing.T) { + // Top-level file uses relative path; imports with ./ resolve correctly. + b := newMemBackend() + b.set("sub/agents.md", "start @./other.md end") + b.set("sub/other.md", "OTHER CONTENT") + + l := newLoaderConfig(b, []string{"sub/agents.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "start @./other.md end") { + t.Fatalf("expected original text preserved, got %q", content) + } + if !strings.Contains(content, "OTHER CONTENT") { + t.Fatalf("expected imported content, got %q", content) + } +} + +func TestLoader_RelativeTopLevelWithDotDotImport(t *testing.T) { + // Top-level file uses relative path; import with ../ resolves correctly. + b := newMemBackend() + b.set("sub/agents.md", "see @../shared/x.md here") + b.set("shared/x.md", "SHARED X") + + l := newLoaderConfig(b, []string{"sub/agents.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "SHARED X") { + t.Fatalf("expected imported content, got %q", content) + } + // filepath.Clean should normalize "sub/../shared/x.md" to "shared/x.md" + if !strings.Contains(content, "Contents of shared/x.md") { + t.Fatalf("expected normalized path in section header, got %q", content) + } +} + +func TestLoader_RelativeTopLevelDedup(t *testing.T) { + // Two top-level relative paths that resolve to the same file via filepath.Clean + // should be deduped (loaded only once). + b := newMemBackend() + b.set("sub/a.md", "CONTENT A") + + l := newLoaderConfig(b, []string{"sub/a.md", "./sub/a.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + count := strings.Count(content, "CONTENT A") + if count != 1 { + t.Fatalf("expected file loaded once (deduped), got %d occurrences in %q", count, content) + } +} + +func TestLoader_AbsoluteTopLevelWithRelativeImport(t *testing.T) { + // Absolute top-level path with relative @import resolves correctly. + b := newMemBackend() + b.set("/project/agents.md", "ref @./lib/helper.md done") + b.set("/project/lib/helper.md", "HELPER") + + l := newLoaderConfig(b, []string{"/project/agents.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "HELPER") { + t.Fatalf("expected imported content, got %q", content) + } + if !strings.Contains(content, "Contents of /project/lib/helper.md") { + t.Fatalf("expected section header, got %q", content) + } +} + +func TestLoader_AbsoluteTopLevelWithDotDotImport(t *testing.T) { + // Absolute top-level path; @import with ../ resolves and normalizes. + b := newMemBackend() + b.set("/project/sub/agents.md", "load @../shared/x.md here") + b.set("/project/shared/x.md", "SHARED") + + l := newLoaderConfig(b, []string{"/project/sub/agents.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "SHARED") { + t.Fatalf("expected imported content, got %q", content) + } + // filepath.Clean normalizes "/project/sub/../shared/x.md" to "/project/shared/x.md" + if !strings.Contains(content, "Contents of /project/shared/x.md") { + t.Fatalf("expected normalized path in section header, got %q", content) + } +} + +func TestLoader_RelativeImportDedup(t *testing.T) { + // Two different relative @import paths that resolve to the same file + // should be deduped via filepath.Clean. + b := newMemBackend() + b.set("/a/main.md", "first @/a/b/shared.md second @../a/b/shared.md end") + b.set("/a/b/shared.md", "SHARED ONCE") + + l := newLoaderConfig(b, []string{"/a/main.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + count := strings.Count(content, "SHARED ONCE") + if count != 1 { + t.Fatalf("expected shared file loaded once (deduped), got %d in %q", count, content) + } +} + +func TestLoader_NestedRelativeImport(t *testing.T) { + // File A imports B via relative path, B imports C via relative path. + // All three should appear as separate sections. + b := newMemBackend() + b.set("/root/main.md", "start @sub/mid.md end") + b.set("/root/sub/mid.md", "mid @deep/leaf.md mid_end") + b.set("/root/sub/deep/leaf.md", "LEAF") + + l := newLoaderConfig(b, []string{"/root/main.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + for _, section := range []string{"Contents of /root/main.md", "Contents of /root/sub/mid.md", "Contents of /root/sub/deep/leaf.md"} { + if !strings.Contains(content, section) { + t.Fatalf("expected section %q, got %q", section, content) + } + } + if !strings.Contains(content, "LEAF") { + t.Fatalf("expected leaf content, got %q", content) + } +} + +func TestLoader_TransitiveImport(t *testing.T) { + // Imported file itself contains @imports; all should appear as separate sections. + b := newMemBackend() + b.set("/main.md", "header @/mid.md footer") + b.set("/mid.md", "mid-start @/leaf.md mid-end") + b.set("/leaf.md", "LEAF_VALUE") + + l := newLoaderConfig(b, []string{"/main.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + for _, section := range []string{"Contents of /main.md", "Contents of /mid.md", "Contents of /leaf.md"} { + if !strings.Contains(content, section) { + t.Fatalf("expected section %q, got %q", section, content) + } + } + if !strings.Contains(content, "LEAF_VALUE") { + t.Fatalf("expected leaf value, got %q", content) + } +} + +func TestLoader_EmptyFile(t *testing.T) { + b := newMemBackend() + b.set("/empty.md", "") + + l := newLoaderConfig(b, []string{"/empty.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + // Empty file is treated as non-existent, so output should be empty. + if content != "" { + t.Fatalf("expected empty output for empty file, got %q", content) + } +} + +func TestLoader_MaxBytesFirstFileFull(t *testing.T) { + // Even if the first file alone exceeds maxBytes, it should still be loaded in full. + b := newMemBackend() + b.set("/big.md", "ABCDEFGHIJ") // 10 bytes + + l := newLoaderConfig(b, []string{"/big.md"}, 3, nil) + content, err := l.load(context.Background()) // maxBytes=3, but first file always loads + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "ABCDEFGHIJ") { + t.Fatalf("first file should always load in full, got %q", content) + } +} + +func TestLoader_CircularImportInline(t *testing.T) { + // Circular reference via @import should be detected, logged, and skipped. + b := newMemBackend() + b.set("/a.md", "text @/b.md more") + b.set("/b.md", "ref @/a.md back") + + l := newLoaderConfig(b, []string{"/a.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error (circular import is logged), got %v", err) + } + // Both a and b should have sections; circular back-reference a from b is skipped. + if !strings.Contains(content, "Contents of /a.md") { + t.Fatalf("expected /a.md section, got %q", content) + } + if !strings.Contains(content, "Contents of /b.md") { + t.Fatalf("expected /b.md section, got %q", content) + } +} + +func TestLoader_MaxDepthInline(t *testing.T) { + // Deep chain via @import should be logged at depth > 5, not returned as error. + b := newMemBackend() + for i := 0; i < 7; i++ { + var content string + if i < 6 { + content = fmt.Sprintf("level%d @/level%d.md tail", i, i+1) + } else { + content = fmt.Sprintf("level%d", i) + } + b.set(fmt.Sprintf("/level%d.md", i), content) + } + + l := newLoaderConfig(b, []string{"/level0.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error (depth exceeded is logged), got %v", err) + } + // Levels 0-5 should have sections. + for i := 0; i <= 5; i++ { + want := fmt.Sprintf("Contents of /level%d.md", i) + if !strings.Contains(content, want) { + t.Fatalf("expected %q in content, got %q", want, content) + } + } + // Level 6 should not be present. + if strings.Contains(content, "Contents of /level6.md") { + t.Fatalf("level6 should not be present (depth exceeded), got %q", content) + } +} + +func TestLoader_DiamondDependency(t *testing.T) { + // A imports B and D; B imports C; D also imports C. + // C should appear only once (deduped across the whole load). + b := newMemBackend() + b.set("/a.md", "start @/b.md middle @/d.md end") + b.set("/b.md", "B(@/c.md)") + b.set("/d.md", "D(@/c.md)") + b.set("/c.md", "SHARED") + + l := newLoaderConfig(b, []string{"/a.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("diamond dependency should not be circular, got error: %v", err) + } + + // C should appear only once as a section (deduped). + count := strings.Count(content, "Contents of /c.md") + if count != 1 { + t.Fatalf("expected /c.md section once (deduped), got %d in %q", count, content) + } + // All files should have sections. + for _, section := range []string{"Contents of /a.md", "Contents of /b.md", "Contents of /c.md", "Contents of /d.md"} { + if !strings.Contains(content, section) { + t.Fatalf("expected section %q, got %q", section, content) + } + } +} + +func TestLoader_AtSignInNormalText(t *testing.T) { + // Bare @word without "/" or file extension should not trigger import. + // Email-like patterns (@example.com) with non-allowed extensions should also be ignored. + b := newMemBackend() + b.set("/agent.md", "contact me @ anytime or @ spaces and @someone mentioned and user@example.com and @company.org") + + l := newLoaderConfig(b, []string{"/agent.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(content, "contact me @ anytime") { + t.Fatalf("bare @ should not trigger import, got %q", content) + } + if !strings.Contains(content, "@someone mentioned") { + t.Fatalf("@someone without / or extension should not trigger import, got %q", content) + } + if !strings.Contains(content, "@example.com") { + t.Fatalf("email-like @example.com should not trigger import, got %q", content) + } + if !strings.Contains(content, "@company.org") { + t.Fatalf("email-like @company.org should not trigger import, got %q", content) + } +} + +func TestLoader_MaxBytesWithImports(t *testing.T) { + // Two top-level files that both import the same shared file. + // Budget should account for imported file bytes. + b := newMemBackend() + b.set("/a.md", "A(@/shared.md)") + b.set("/b.md", "B(@/shared.md)") + b.set("/shared.md", strings.Repeat("X", 100)) // 100 bytes + + l := newLoaderConfig(b, []string{"/a.md", "/b.md"}, 120, nil) + // /a.md = 14 bytes + /shared.md = 100 bytes => 114 total after /a.md. + // Budget = 120: /b.md (14 bytes) would push to 128, exceeding budget. + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("load failed: %v", err) + } + + // /a.md and its import should be included. + if !strings.Contains(content, strings.Repeat("X", 100)) { + t.Fatal("expected /a.md with shared content to be included") + } + + // /b.md should be excluded because totalBytes exceeded budget after loading /a.md. + if strings.Contains(content, "B(") { + t.Fatalf("expected /b.md to be excluded due to budget, got %q", content) + } +} + +func TestNew_Validation_EmptyAgentFiles(t *testing.T) { + ctx := context.Background() + b := newMemBackend() + + _, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{}}) + if err == nil { + t.Fatal("expected error for empty agent files") + } + if !strings.Contains(err.Error(), "at least one agent file path is required") { + t.Fatalf("unexpected error message: %v", err) + } +} + +func TestMiddleware_GenerateError(t *testing.T) { + // Non-ErrNotExist errors (e.g. permission denied) should propagate. + b := &errBackend{} + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/file.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hi"}}} + _, _, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err == nil { + t.Fatal("expected error when backend read fails with non-ErrNotExist") + } + if !strings.Contains(err.Error(), "failed to load agent files") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestLoader_DuplicateTopLevelFiles(t *testing.T) { + // Same file listed twice in AgentFiles; second should be deduped via seen map. + b := newMemBackend() + b.set("/agent.md", "unique content") + + l := newLoaderConfig(b, []string{"/agent.md", "/agent.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + + count := strings.Count(content, "Contents of /agent.md") + if count != 1 { + t.Fatalf("expected /agent.md section once (deduped), got %d", count) + } +} + +func TestLoader_LoadFileError(t *testing.T) { + // Missing file (ErrNotExist) is silently skipped. + b := newMemBackend() + l := newLoaderConfig(b, []string{"/missing.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected missing file to be skipped, got error: %v", err) + } + if content != "" { + t.Fatalf("expected empty output, got %q", content) + } +} + +func TestLoader_MaxBytesStopsImports(t *testing.T) { + // When budget is exhausted, further imports in collectImports should be skipped. + b := newMemBackend() + b.set("/main.md", "@/big.md @/small.md") + b.set("/big.md", strings.Repeat("B", 200)) + b.set("/small.md", "SMALL") + + l := newLoaderConfig(b, []string{"/main.md"}, 50, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + + // main.md itself is loaded (always), big.md pushes over budget, + // small.md should be skipped. + if !strings.Contains(content, "Contents of /main.md") { + t.Fatal("main.md should be present") + } + if strings.Contains(content, "SMALL") { + t.Fatal("small.md should be skipped after budget exhausted") + } +} + +func TestFormatContent_Empty(t *testing.T) { + // formatContent with nil/empty slice should return empty string. + if got := formatContent(nil); got != "" { + t.Fatalf("expected empty string for nil, got %q", got) + } + if got := formatContent([]loadedFile{}); got != "" { + t.Fatalf("expected empty string for empty slice, got %q", got) + } +} + +func TestMiddleware_AllFilesEmpty(t *testing.T) { + // When all agent files have empty content, loader returns "" and + // BeforeModelRewriteState returns the original state unchanged. + b := newMemBackend() + b.set("/agent.md", "") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + userMsg := []*schema.Message{{Role: schema.User, Content: "hello"}} + state := &adk.ChatModelAgentState{Messages: userMsg} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + // Empty file produces no agentmd content, so original messages pass through unchanged. + if len(state.Messages) != 1 { + t.Fatalf("expected 1 message (no agentmd prepended), got %d", len(state.Messages)) + } + if state.Messages[0].Content != "hello" { + t.Fatalf("expected original message unchanged, got %q", state.Messages[0].Content) + } +} + +func TestLoader_ExactOutput(t *testing.T) { + // Verify the exact output format matches the expected structure: + // each file (top-level and imported) gets its own "Contents of ..." section, + // @path references are preserved in the original text. + b := newMemBackend() + b.set("/project/CLAUDE.md", "this is project claude.md\n\n- git workflow @git/git-instructions.md") + b.set("/project/git/git-instructions.md", "this is git-instructions.md") + + l := newLoaderConfig(b, []string{"/project/CLAUDE.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatal(err) + } + + expected := ` +As you answer the user's questions, you can use the following context: +Codebase and user instructions are shown below. Be sure to adhere to these instructions. IMPORTANT: These instructions OVERRIDE any default behavior and you MUST follow them exactly as written. + +Contents of /project/CLAUDE.md (instructions): + +this is project claude.md + +- git workflow @git/git-instructions.md + +Contents of /project/git/git-instructions.md (instructions): + +this is git-instructions.md +IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task. +` + + if content != expected { + t.Fatalf("output mismatch.\n\ngot:\n%s\n\nexpected:\n%s", content, expected) + } +} + +func TestLoader_MissingFileSkipped(t *testing.T) { + b := newMemBackend() + b.set("/good.md", "GOOD CONTENT") + // /missing.md is not set + + l := newLoaderConfig(b, []string{"/missing.md", "/good.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error for missing file, got %v", err) + } + if !strings.Contains(content, "GOOD CONTENT") { + t.Fatal("expected good.md content in output") + } +} + +func TestLoader_AllMissingFilesSkipped(t *testing.T) { + b := newMemBackend() + + l := newLoaderConfig(b, []string{"/missing1.md", "/missing2.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error for missing files, got %v", err) + } + if content != "" { + t.Fatalf("expected empty output when all files missing, got %q", content) + } +} + +func TestLoader_CircularImportSkipped(t *testing.T) { + b := newMemBackend() + b.set("/a.md", "A content @/b.md") + b.set("/b.md", "B content @/a.md") + + // Circular import in collectImports is logged via onWarning and skipped. + l := newLoaderConfig(b, []string{"/a.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if !strings.Contains(content, "A content") { + t.Fatal("expected a.md content") + } + if !strings.Contains(content, "B content") { + t.Fatal("expected b.md content") + } +} + +func TestLoader_DepthExceededSkipped(t *testing.T) { + b := newMemBackend() + // Create a chain that exceeds maxImportDepth (5) + b.set("/l0.md", "@/l1.md") + b.set("/l1.md", "@/l2.md") + b.set("/l2.md", "@/l3.md") + b.set("/l3.md", "@/l4.md") + b.set("/l4.md", "@/l5.md") + b.set("/l5.md", "@/l6.md") + b.set("/l6.md", "DEEP") + + l := newLoaderConfig(b, []string{"/l0.md"}, 0, nil) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error for depth exceeded, got %v", err) + } + // Should have content up to the depth limit, deep file skipped. + if !strings.Contains(content, "/l0.md") { + t.Fatal("expected l0.md in output") + } +} + +func TestLoader_OnLoadWarningCallback(t *testing.T) { + b := newMemBackend() + b.set("/good.md", "GOOD CONTENT") + + var warnings []error + onWarning := func(filePath string, err error) { + warnings = append(warnings, fmt.Errorf("%s: %w", filePath, err)) + } + + l := newLoaderConfig(b, []string{"/missing.md", "/good.md"}, 0, onWarning) + content, err := l.load(context.Background()) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if !strings.Contains(content, "GOOD CONTENT") { + t.Fatal("expected good.md content in output") + } + if len(warnings) == 0 { + t.Fatal("expected at least one warning for missing file") + } + if !strings.Contains(warnings[0].Error(), "file not found") { + t.Fatalf("expected file not found warning, got %v", warnings[0]) + } +} + +func TestMiddleware_MissingFile(t *testing.T) { + b := newMemBackend() + // /missing.md not set — will fail to read + + ctx := context.Background() + mw, err := New(ctx, &Config{ + Backend: b, + AgentsMDFiles: []string{"/missing.md"}, + }) + if err != nil { + t.Fatal(err) + } + + userMsg := []*schema.Message{{Role: schema.User, Content: "hello"}} + state := &adk.ChatModelAgentState{Messages: userMsg} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatalf("expected no error for missing file, got %v", err) + } + // No agent.md content, so original messages should be passed through unchanged. + if len(state.Messages) != 1 { + t.Fatalf("expected 1 message (no agentmd prepended), got %d", len(state.Messages)) + } +} + +func TestMiddleware_InsertBeforeFirstUserMessage(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + // Input has a System message before the User message. + input := []*schema.Message{ + {Role: schema.System, Content: "system prompt"}, + {Role: schema.User, Content: "hello"}, + } + state := &adk.ChatModelAgentState{Messages: input} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + if len(state.Messages) != 3 { + t.Fatalf("expected 3 messages, got %d", len(state.Messages)) + } + if state.Messages[0].Role != schema.System { + t.Fatalf("expected first message role System, got %s", state.Messages[0].Role) + } + if state.Messages[0].Content != "system prompt" { + t.Fatalf("expected system prompt preserved, got %q", state.Messages[0].Content) + } + if state.Messages[1].Role != schema.User || !strings.Contains(state.Messages[1].Content, "agent instructions") { + t.Fatalf("expected agentmd message before user message, got role=%s content=%q", state.Messages[1].Role, state.Messages[1].Content) + } + if state.Messages[2].Role != schema.User || state.Messages[2].Content != "hello" { + t.Fatalf("expected original user message at index 2, got role=%s content=%q", state.Messages[2].Role, state.Messages[2].Content) + } +} + +func TestMiddleware_InsertWithNoUserMessage(t *testing.T) { + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + // Input has no User message at all. + input := []*schema.Message{ + {Role: schema.System, Content: "system prompt"}, + {Role: schema.Assistant, Content: "assistant reply"}, + } + state := &adk.ChatModelAgentState{Messages: input} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + + if len(state.Messages) != 3 { + t.Fatalf("expected 3 messages, got %d", len(state.Messages)) + } + if state.Messages[0].Role != schema.System { + t.Fatalf("expected System at index 0, got %s", state.Messages[0].Role) + } + if state.Messages[1].Role != schema.Assistant { + t.Fatalf("expected Assistant at index 1, got %s", state.Messages[1].Role) + } + if state.Messages[2].Role != schema.User || !strings.Contains(state.Messages[2].Content, "agent instructions") { + t.Fatalf("expected agentmd appended at end, got role=%s content=%q", state.Messages[2].Role, state.Messages[2].Content) + } +} + +func TestLoader_ImportIOError(t *testing.T) { + // When an imported file returns a non-ErrNotExist error (e.g. I/O error), + // the load should propagate the error (covers collectImports and loadFile error paths). + b := &partialErrBackend{ + files: map[string]string{ + "/main.md": "content @/broken.md", + }, + // /broken.md is NOT in the map, so Read returns I/O error (not ErrNotExist) + } + + l := newLoaderConfig(b, []string{"/main.md"}, 0, nil) + _, err := l.load(context.Background()) + if err == nil { + t.Fatal("expected error from I/O failure on imported file") + } + if !strings.Contains(err.Error(), "I/O error") { + t.Fatalf("expected I/O error, got: %v", err) + } +} + +func TestMiddleware_Idempotency(t *testing.T) { + // Calling BeforeModelRewriteState twice should NOT duplicate the agentsmd message. + // The marker in msg.Extra[agentsMDExtraKey] prevents re-injection. + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hello"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after first call, got %d", len(state.Messages)) + } + + // Call again with the same state (which now contains the marker message). + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after second call (idempotent), got %d", len(state.Messages)) + } + if !strings.Contains(state.Messages[0].Content, "agent instructions") { + t.Fatalf("expected agentmd content preserved, got %q", state.Messages[0].Content) + } +} + +func TestMiddleware_ReinsertAfterRemoval(t *testing.T) { + // If the marker message is removed from state.Messages, calling + // BeforeModelRewriteState should re-insert it. + b := newMemBackend() + b.set("/agent.md", "agent instructions") + + ctx := context.Background() + mw, err := New(ctx, &Config{Backend: b, AgentsMDFiles: []string{"/agent.md"}}) + if err != nil { + t.Fatal(err) + } + + state := &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hello"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after first call, got %d", len(state.Messages)) + } + + // Simulate removal of the marker message (e.g., by summarization). + // Keep only the original user message. + state = &adk.ChatModelAgentState{Messages: []*schema.Message{{Role: schema.User, Content: "hello"}}} + _, state, err = mw.BeforeModelRewriteState(ctx, state, nil) + if err != nil { + t.Fatal(err) + } + if len(state.Messages) != 2 { + t.Fatalf("expected 2 messages after re-insert, got %d", len(state.Messages)) + } + if !strings.Contains(state.Messages[0].Content, "agent instructions") { + t.Fatalf("expected agentmd content re-inserted, got %q", state.Messages[0].Content) + } +} + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + b := newMemBackend() + b.set("/agent.md", "You are a helpful assistant.") + + mw, err := NewTyped[*schema.AgenticMessage](ctx, &Config{ + Backend: b, + AgentsMDFiles: []string{"/agent.md"}, + }) + if err != nil { + t.Fatal(err) + } + if mw == nil { + t.Fatal("expected non-nil middleware") + } + + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} diff --git a/adk/middlewares/agentsmd/loader.go b/adk/middlewares/agentsmd/loader.go new file mode 100644 index 0000000..db73338 --- /dev/null +++ b/adk/middlewares/agentsmd/loader.go @@ -0,0 +1,299 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package agentsmd + +import ( + "context" + "errors" + "fmt" + "log" + "os" + "path/filepath" + "regexp" + "strings" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" +) + +// importRegex matches @path/to/file anywhere in text. +// The path must start with a letter, digit, dot, underscore, slash, or tilde, followed by +// path characters (letters, digits, dots, slashes, hyphens, underscores). +// A post-match filter further requires the path to contain "/" or end with +// an allowed extension (see allowedImportExts), so bare words like @someone +// and email-like patterns like @example.com are ignored. +var importRegex = regexp.MustCompile(`@([a-zA-Z0-9_.~/][a-zA-Z0-9_.~/\-]*)`) + +// allowedImportExts is the set of file extensions recognised as @import targets. +// Paths without "/" must end with one of these extensions to be treated as imports; +// this avoids false positives on email addresses (@example.com) and mentions (@foo.bar). +var allowedImportExts = map[string]bool{ + ".md": true, + ".txt": true, + ".mdx": true, + ".yaml": true, + ".yml": true, + ".json": true, + ".toml": true, +} + +const maxImportDepth = 5 + +// ReadRequest is an alias for filesystem.ReadRequest. +type ReadRequest = filesystem.ReadRequest +type FileContent = filesystem.FileContent + +// Backend defines the file access interface for loading Agents.md files. +// Implementations can use local filesystem, remote storage, or any other backend. +type Backend interface { + // Read reads the content of a file. + // If the file does not exist, implementations should return an error wrapping + // os.ErrNotExist (so that errors.Is(err, os.ErrNotExist) returns true). This allows the loader + // to silently skip missing files and notify via OnLoadWarning callback. + // Other errors (e.g. permission denied, I/O errors) will abort the loading process. + Read(ctx context.Context, req *ReadRequest) (*FileContent, error) +} + +// loaderConfig holds the immutable configuration for creating loaders. +// It is safe for concurrent use by multiple goroutines. +type loaderConfig struct { + backend Backend + files []string // ordered file paths from config + maxBytes int // cumulative read budget; 0 means unlimited + onWarning func(filePath string, err error) // callback for non-fatal loading warnings +} + +func newLoaderConfig(backend Backend, files []string, maxBytes int, onWarning func(filePath string, err error)) *loaderConfig { + if onWarning == nil { + onWarning = func(filePath string, err error) { + log.Printf("[agentsmd] warning: %s: %v", filePath, err) + } + } + return &loaderConfig{ + backend: backend, + files: files, + maxBytes: maxBytes, + onWarning: onWarning, + } +} + +// loader handles loading and @import resolution for agents.md files. +// A new loader is created for each load() call to avoid sharing mutable state +// (totalBytes) across concurrent invocations. +type loader struct { + *loaderConfig + totalBytes int // accumulated bytes during this load call +} + +func (cfg *loaderConfig) newLoader() *loader { + return &loader{loaderConfig: cfg} +} + +// load reads all agents.md files and returns the formatted content. +// Each top-level file and its @imported files appear as separate sections. +func (cfg *loaderConfig) load(ctx context.Context) (string, error) { + l := cfg.newLoader() + + var parts []loadedFile + seen := make(map[string]bool) // dedup across all files and imports + + for i, filePath := range l.files { + files, err := l.loadFile(ctx, filePath, 0, make(map[string]bool), seen) + if err != nil { + return "", fmt.Errorf("failed to load %q: %w", filePath, err) + } + + // If loading this file caused the budget to be exceeded, skip it + // (but always include the first file). + if i > 0 && l.maxBytes > 0 && l.totalBytes > l.maxBytes { + l.onWarning(filePath, fmt.Errorf("skipped: cumulative size %d exceeds max bytes %d", l.totalBytes, l.maxBytes)) + break + } + + parts = append(parts, files...) + } + + return formatContent(parts), nil +} + +// loadFile reads a file via Backend and collects @imported files as separate entries. +// Returns a slice where the first element is this file itself, followed by all +// transitively imported files (in encounter order, preserving @path in original text). +// visited tracks the current ancestor chain to detect circular imports. +// seen tracks globally loaded files to avoid duplicate reads and byte counting. +func (l *loader) loadFile(ctx context.Context, filePath string, depth int, visited map[string]bool, seen map[string]bool) ([]loadedFile, error) { + filePath = filepath.Clean(filePath) + + if depth > maxImportDepth { + l.onWarning(filePath, fmt.Errorf("@import depth exceeds maximum of %d", maxImportDepth)) + return nil, nil + } + + if visited[filePath] { + l.onWarning(filePath, fmt.Errorf("circular @import detected")) + return nil, nil + } + + if seen[filePath] { + return nil, nil + } + + visited[filePath] = true + defer delete(visited, filePath) + + fileContent, err := l.backend.Read(ctx, &ReadRequest{FilePath: filePath, Offset: 1}) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + l.onWarning(filePath, fmt.Errorf("file not found, skipping")) + return nil, nil + } + return nil, err + } + content := "" + if fileContent != nil { + content = fileContent.Content + } + + l.totalBytes += len(content) + seen[filePath] = true + + if content == "" { + return nil, nil + } + + // Collect imported files as separate sections (content stays untouched). + imports, err := l.collectImports(ctx, filePath, content, depth, visited, seen) + if err != nil { + return nil, err + } + + // This file first, then its imports. + result := make([]loadedFile, 0, 1+len(imports)) + result = append(result, loadedFile{path: filePath, content: content}) + result = append(result, imports...) + return result, nil +} + +// collectImports scans content for @path/to/file references and loads each +// imported file (plus its transitive imports). The original content is NOT modified. +// Returns the list of imported loadedFile entries in encounter order. +// seen is shared across the entire load call to avoid duplicate reads. +// Non-fatal errors (file not found, depth exceeded, circular import) are reported +// via onWarning and skipped. Fatal errors (e.g. I/O) are returned. +func (l *loader) collectImports(ctx context.Context, hostPath, content string, depth int, visited map[string]bool, seen map[string]bool) ([]loadedFile, error) { + dir := filepath.Dir(hostPath) + var imports []loadedFile + + matches := importRegex.FindAllStringSubmatch(content, -1) + for _, match := range matches { + rawPath := match[1] + + // Only treat as import if path contains "/" or ends with an allowed extension. + // This avoids false positives on email addresses and social mentions. + if !strings.Contains(rawPath, "/") && !allowedImportExts[filepath.Ext(rawPath)] { + continue + } + + // If budget is exhausted, skip further imports. + if l.maxBytes > 0 && l.totalBytes > l.maxBytes { + break + } + + importPath := rawPath + if !filepath.IsAbs(importPath) { + importPath = filepath.Join(dir, importPath) + } + + if seen[importPath] { + continue + } + + files, err := l.loadFile(ctx, importPath, depth+1, visited, seen) + if err != nil { + return nil, fmt.Errorf("failed to import %q from %q: %w", rawPath, hostPath, err) + } + + imports = append(imports, files...) + } + + return imports, nil +} + +type loadedFile struct { + path string + content string +} + +const formatHeaderEn = ` +As you answer the user's questions, you can use the following context: +Codebase and user instructions are shown below. Be sure to adhere to these instructions. IMPORTANT: These instructions OVERRIDE any default behavior and you MUST follow them exactly as written. +` + +const formatHeaderCn = ` +在回答用户问题时，你可以使用以下上下文： +代码库和用户指令如下。请务必遵守这些指令。重要提示：这些指令会覆盖任何默认行为，你必须严格按照要求执行。 +` + +const formatFileHeaderEn = "\nContents of " + +const formatFileHeaderCn = "\n文件内容：" + +const formatFileLabelEn = " (instructions):\n\n" + +const formatFileLabelCn = "（指令）：\n\n" + +const formatFooterEn = `IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task. +` + +const formatFooterCn = `重要提示：此上下文可能与你的任务相关，也可能不相关。除非此上下文与你的任务高度相关，否则不要响应此上下文。 +` + +func formatContent(files []loadedFile) string { + if len(files) == 0 { + return "" + } + + header := internal.SelectPrompt(internal.I18nPrompts{ + English: formatHeaderEn, + Chinese: formatHeaderCn, + }) + fileHeader := internal.SelectPrompt(internal.I18nPrompts{ + English: formatFileHeaderEn, + Chinese: formatFileHeaderCn, + }) + fileLabel := internal.SelectPrompt(internal.I18nPrompts{ + English: formatFileLabelEn, + Chinese: formatFileLabelCn, + }) + footer := internal.SelectPrompt(internal.I18nPrompts{ + English: formatFooterEn, + Chinese: formatFooterCn, + }) + + var sb strings.Builder + sb.WriteString(header) + + for _, f := range files { + sb.WriteString(fileHeader) + sb.WriteString(f.path) + sb.WriteString(fileLabel) + sb.WriteString(f.content) + sb.WriteString("\n") + } + sb.WriteString(footer) + return sb.String() +} diff --git a/adk/middlewares/dynamictool/toolsearch/prompt.go b/adk/middlewares/dynamictool/toolsearch/prompt.go new file mode 100644 index 0000000..5aaa56a --- /dev/null +++ b/adk/middlewares/dynamictool/toolsearch/prompt.go @@ -0,0 +1,162 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package toolsearch + +const ( + toolDescription = `Search for or select deferred tools to make them available for use. + +MANDATORY PREREQUISITE - THIS IS A HARD REQUIREMENT + +You MUST use this tool to load deferred tools BEFORE calling them directly. + +This is a BLOCKING REQUIREMENT - deferred tools are NOT available until you load them using this tool. Look for messages in the conversation for the list of tools you can discover. Both query modes (keyword search and direct selection) load the returned tools — once a tool appears in the results, it is immediately available to call. + +Why this is non-negotiable: +- Deferred tools are not loaded until discovered via this tool +- Calling a deferred tool without first loading it will fail +Query modes: + +1. Keyword search - Use keywords when you're unsure which tool to use or need to discover multiple tools at once: + - "list directory" - find tools for listing directories + - "notebook jupyter" - find notebook editing tools + - "slack message" - find slack messaging tools + - Returns up to 5 matching tools ranked by relevance + - All returned tools are immediately available to call — no further selection step needed +2. Direct selection - Use select: when you know the exact tool name: + - "select:mcp__slack__read_channel" + - "select:NotebookEdit" + - "select:Read,Edit,Grep" - load multiple tools at once with comma separation + - Returns the named tool(s) if they exist +IMPORTANT: Both modes load tools equally. Do NOT follow up a keyword search with select: calls for tools already returned — they are already loaded. + +3. Required keyword - Prefix with + to require a match: + - "+linear create issue" - only tools from "linear", ranked by "create"/"issue" + - "+slack send" - only "slack" tools, ranked by "send" + - Useful when you know the service name but not the exact tool +CORRECT Usage Patterns: + + +User: I need to work with slack somehow +Assistant: Let me search for slack tools. +[Calls tool_search with query: "slack"] +Assistant: Found several options including mcp__slack__read_channel. +[Calls mcp__slack__read_channel directly — it was loaded by the keyword search] + + + +User: Edit the Jupyter notebook +Assistant: Let me load the notebook editing tool. +[Calls tool_search with query: "select:NotebookEdit"] +[Calls NotebookEdit] + + + +User: List files in the src directory +Assistant: I can see mcp__filesystem__list_directory in the available tools. Let me select it. +[Calls tool_search with query: "select:mcp__filesystem__list_directory"] +[Calls the tool] + + +INCORRECT Usage Patterns - NEVER DO THESE: + + +User: Read my slack messages +Assistant: [Directly calls mcp__slack__read_channel without loading it first] +WRONG - You must load the tool FIRST using this tool + + + +Assistant: [Calls tool_search with query: "slack", gets back mcp__slack__read_channel] +Assistant: [Calls tool_search with query: "select:mcp__slack__read_channel"] +WRONG - The keyword search already loaded the tool. The select call is redundant. +` + + toolDescriptionChinese = `搜索或选择延迟加载（deferred）的工具，使其可供调用。 + +强制前提条件（MANDATORY PREREQUISITE）— 硬性要求 + +在直接调用任何 延迟加载工具（deferred tools） 之前，你 必须先使用此工具将其加载。 + +这是一个 阻塞性要求（BLOCKING REQUIREMENT) — 延迟加载工具在被加载之前是 不可用的。你需要在对话中查找 消息，以获取可以发现的工具列表。无论使用哪种查询方式（关键字搜索 或 直接选择），只要工具出现在返回结果中，它们就会自动被加载并立即可调用。 + +为什么这是不可协商的规则： +- 延迟加载工具在被发现之前不会被加载 +- 如果你在加载之前直接调用延迟工具，调用将会失败 +查询模式： + +1. 关键字搜索（Keyword search）- 当你不确定具体需要哪个工具，或希望一次发现多个工具时使用关键字搜索： +- "list directory" — 查找用于列出目录的工具 +- "notebook jupyter" — 查找 Jupyter Notebook 编辑工具 +- "slack message" — 查找 Slack 消息相关工具 +- 返回最多 5 个最相关的工具 +- 所有返回的工具都会立即加载并可直接调用 — 不需要额外执行 select 步骤 + +2. 直接选择（Direct selection）— 当你已经知道工具的确切名称时使用 select:： +- "select:mcp__slack__read_channel" +- "select:NotebookEdit" +- "select:Read,Edit,Grep" — 一次加载多个工具 +- 如果工具存在，将被加载并返回 +重要说明：两种模式的加载效果完全相同。不要在关键词搜索之后，对返回的工具再次进行 select: 选择 — 它们已经加载好了。 + +3. 必须匹配关键字（Required keyword）— 在关键字前添加 + 可以 强制匹配特定服务或来源。 +- "+linear create issue" — 仅返回名字中包含 "linear" 的工具，按 "create" / "issue" 排序 +- "+slack send" — 仅返回名字中包含 "slack" 的工具，按 "send" 排序 +- 适用于你知道服务名称但不知道具体工具名称 + +正确使用示例： + + +User: 我需要处理 Slack 相关的事情 +Assistant: 让我搜索 Slack 工具。 +[调用 tool_search，query: "slack"] +Assistant: 找到多个选项，包括 mcp__slack__read_channel。 +[直接调用 mcp__slack__read_channel — 关键字搜索已经加载了该工具] + + + +User: 编辑这个 Jupyter Notebook +Assistant: 让我加载 Notebook 编辑工具。 +[调用 tool_search，query: "select:NotebookEdit"] +[调用 NotebookEdit] + + + +User: 列出 src 目录中的文件 +Assistant: 我看到可用工具中有 mcp__filesystem__list_directory，让我加载它。 +[调用 tool_search，query: "select:mcp__filesystem__list_directory"] +[调用该工具] + + +错误用法（严禁） + + +User: 读取我的 Slack 消息 +Assistant: [不调用 tool_search 工具加载，直接调用 mcp__slack__read_channel] +错误 — 在调用工具之前没有先使用 tool_search 加载该工具。 + + + +Assistant:[调用 tool_search，query: "slack"，返回 mcp__slack__read_channel] +Assistant:[再次调用 tool_search，query: "select:mcp__slack__read_channel"] +错误 — 关键字搜索 已经加载了该工具，再次 select 是冗余操作。` + + systemReminderTpl = ` +{{- range .Tools }} +{{ . }} +{{- end }} +` +) diff --git a/adk/middlewares/dynamictool/toolsearch/toolsearch.go b/adk/middlewares/dynamictool/toolsearch/toolsearch.go new file mode 100644 index 0000000..9215b19 --- /dev/null +++ b/adk/middlewares/dynamictool/toolsearch/toolsearch.go @@ -0,0 +1,647 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package toolsearch provides tool search middleware. +package toolsearch + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "text/template" + "unicode" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// Config is the configuration for the tool search middleware. +type Config struct { + // DynamicTools is a list of tools that can be dynamically searched and loaded by the agent. + DynamicTools []tool.BaseTool + + // UseModelToolSearch indicates whether the ChatModel natively supports tool search. + // + // When true, the middleware delegates tool search to the model's native capability. + // + // When false (default), the middleware manages tool visibility by filtering the tool list + // based on tool_search results before each model call. Note that this approach may + // invalidate the model's KV-cache (as the tool list changes between calls), and effectiveness + // depends on the model's ability to work with a dynamically changing tool set. + UseModelToolSearch bool +} + +// NewTyped constructs and returns the generic tool search middleware. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +func NewTyped[M adk.MessageType](ctx context.Context, config *Config) (adk.TypedChatModelAgentMiddleware[M], error) { + if config == nil { + return nil, fmt.Errorf("config is required") + } + if len(config.DynamicTools) == 0 { + return nil, fmt.Errorf("tools is required") + } + + tpl, err := template.New("").Parse(systemReminderTpl) + if err != nil { + return nil, err + } + + dynamicToolInfos := make([]*schema.ToolInfo, 0, len(config.DynamicTools)) + mapOfDynamicTools := make(map[string]*schema.ToolInfo, len(config.DynamicTools)) + toolNames := make([]string, 0, len(config.DynamicTools)) + for _, t := range config.DynamicTools { + info, infoErr := t.Info(ctx) + if infoErr != nil { + return nil, fmt.Errorf("failed to get dynamic tool info: %w", infoErr) + } + + if _, ok := mapOfDynamicTools[info.Name]; ok { + return nil, fmt.Errorf("duplicate dynamic tool name: %s", info.Name) + } + + toolNames = append(toolNames, info.Name) + mapOfDynamicTools[info.Name] = info + dynamicToolInfos = append(dynamicToolInfos, info) + } + + buf := &bytes.Buffer{} + err = tpl.Execute(buf, systemReminder{Tools: toolNames}) + if err != nil { + return nil, fmt.Errorf("failed to format system reminder template: %w", err) + } + + return &typedMiddleware[M]{ + dynamicTools: config.DynamicTools, + mapOfDynamicTools: mapOfDynamicTools, + dynamicToolInfos: dynamicToolInfos, + useModelToolSearch: config.UseModelToolSearch, + sr: buf.String(), + }, nil +} + +// New constructs and returns the tool search middleware. +// +// The tool search middleware enables dynamic tool selection for agents with large tool libraries. +// Instead of passing all tools to the model at once (which can overwhelm context limits), +// this middleware: +// +// 1. Adds a "tool_search" meta-tool that accepts keyword queries to search tools +// 2. Initially hides all dynamic tools from the model's tool list +// 3. When the model calls tool_search, matching tools become available for subsequent calls +// +// Example usage: +// +// middleware, _ := toolsearch.New(ctx, &toolsearch.Config{ +// DynamicTools: []tool.BaseTool{weatherTool, stockTool, currencyTool, ...}, +// }) +// agent, _ := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ +// // ... +// Handlers: []adk.ChatModelAgentMiddleware{middleware}, +// }) +func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, config) +} + +type systemReminder struct { + Tools []string +} + +type typedMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + dynamicTools []tool.BaseTool + mapOfDynamicTools map[string]*schema.ToolInfo + dynamicToolInfos []*schema.ToolInfo + useModelToolSearch bool + sr string +} + +func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + if runCtx == nil { + return ctx, runCtx, nil + } + + nRunCtx := *runCtx + nRunCtx.Tools = make([]tool.BaseTool, len(runCtx.Tools), len(runCtx.Tools)+1+len(m.dynamicTools)) + copy(nRunCtx.Tools, runCtx.Tools) + nRunCtx.Tools = append(nRunCtx.Tools, newToolSearchTool(m.mapOfDynamicTools, m.useModelToolSearch)) + nRunCtx.Tools = append(nRunCtx.Tools, m.dynamicTools...) + if m.useModelToolSearch { + nRunCtx.ToolSearchTool = getToolSearchToolInfo() + } + return ctx, &nRunCtx, nil +} + +const toolSearchInitializedKey = "__toolsearch_initialized__" +const toolSearchReminderExtraKey = "__toolsearch_reminder__" + +func (m *typedMiddleware[M]) isInitialized(ctx context.Context) bool { + val, ok, err := adk.GetRunLocalValue(ctx, toolSearchInitializedKey) + if err != nil || !ok { + return false + } + b, _ := val.(bool) + return b +} + +func (m *typedMiddleware[M]) markInitialized(ctx context.Context) { + _ = adk.SetRunLocalValue(ctx, toolSearchInitializedKey, true) +} + +func (m *typedMiddleware[M]) ensureReminder(msgs []M) []M { + for _, msg := range msgs { + if hasToolSearchReminderExtra(msg) { + return msgs + } + } + + reminder := makeReminderMsg[M](m.sr) + result := make([]M, 0, len(msgs)+1) + inserted := false + for _, msg := range msgs { + if !inserted && !isSystemRoleTS(msg) { + inserted = true + result = append(result, reminder) + } + result = append(result, msg) + } + if !inserted { + result = append(result, reminder) + } + return result +} + +func isSystemRoleTS[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.System + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeSystem + } + return false +} + +func makeReminderMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.UserMessage(content) + msg.Extra = map[string]any{toolSearchReminderExtraKey: true} + return any(msg).(M) + case *schema.AgenticMessage: + msg := schema.UserAgenticMessage(content) + msg.Extra = map[string]any{toolSearchReminderExtraKey: true} + return any(msg).(M) + } + panic("unreachable") +} + +func hasToolSearchReminderExtra[M adk.MessageType](msg M) bool { + switch v := any(msg).(type) { + case *schema.Message: + if v.Extra != nil { + if b, ok := v.Extra[toolSearchReminderExtraKey]; ok { + if bVal, _ := b.(bool); bVal { + return true + } + } + } + case *schema.AgenticMessage: + if v.Extra != nil { + if b, ok := v.Extra[toolSearchReminderExtraKey]; ok { + if bVal, _ := b.(bool); bVal { + return true + } + } + } + } + return false +} + +func (m *typedMiddleware[M]) extractDynamicTools(tools []*schema.ToolInfo) []*schema.ToolInfo { + var result []*schema.ToolInfo + for _, t := range tools { + if _, ok := m.mapOfDynamicTools[t.Name]; ok { + result = append(result, t) + } + } + return result +} + +func (m *typedMiddleware[M]) stripDynamicTools(tools []*schema.ToolInfo) []*schema.ToolInfo { + var result []*schema.ToolInfo + for _, t := range tools { + if _, ok := m.mapOfDynamicTools[t.Name]; !ok { + result = append(result, t) + } + } + return result +} + +func removeTool(tools []*schema.ToolInfo, name string) []*schema.ToolInfo { + var result []*schema.ToolInfo + for _, t := range tools { + if t.Name != name { + result = append(result, t) + } + } + return result +} + +func toolNameSet(tools []*schema.ToolInfo) map[string]bool { + m := make(map[string]bool, len(tools)) + for _, t := range tools { + m[t.Name] = true + } + return m +} + +func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { + state.Messages = m.ensureReminder(state.Messages) + + if !m.isInitialized(ctx) { + m.markInitialized(ctx) + + if m.useModelToolSearch { + // Model-native search: move dynamic tools to DeferredToolInfos for server-side retrieval, + // keep only static tools in ToolInfos, and remove the tool_search tool (the model handles search itself). + state.DeferredToolInfos = m.extractDynamicTools(state.ToolInfos) + state.ToolInfos = m.stripDynamicTools(state.ToolInfos) + state.ToolInfos = removeTool(state.ToolInfos, toolSearchToolName) + } else { + // Client-side search: hide dynamic tools initially; they become visible + // only after the model calls tool_search and forward selection adds them back. + state.ToolInfos = m.stripDynamicTools(state.ToolInfos) + } + } + + // Forward selection (client-side search only): scan tool_search results in the + // conversation history and add the selected dynamic tools back to ToolInfos. + if !m.useModelToolSearch { + existing := toolNameSet(state.ToolInfos) + for _, msg := range state.Messages { + content, ok := extractToolSearchResult(msg, toolSearchToolName) + if !ok { + continue + } + var result toolSearchResult + if err := json.Unmarshal([]byte(content), &result); err != nil { + continue + } + for _, name := range result.Matches { + if existing[name] { + continue + } + if info, ok := m.mapOfDynamicTools[name]; ok { + state.ToolInfos = append(state.ToolInfos, info) + existing[name] = true + } + } + } + } + + return ctx, state, nil +} + +// extractToolSearchResult checks if the given message is a tool result from the tool_search tool, +// and if so returns the content string. Returns ("", false) if not a matching tool result. +func extractToolSearchResult[M adk.MessageType](msg M, toolName string) (string, bool) { + switch v := any(msg).(type) { + case *schema.Message: + if v.Role == schema.Tool && v.ToolName == toolName { + return v.Content, true + } + case *schema.AgenticMessage: + for _, block := range v.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult && + block.FunctionToolResult != nil && block.FunctionToolResult.Name == toolName { + for _, b := range block.FunctionToolResult.Content { + if b != nil && b.Text != nil { + return b.Text.Text, true + } + } + } + } + } + return "", false +} + +func newToolSearchTool(tools map[string]*schema.ToolInfo, useModelToolSearch bool) tool.BaseTool { + if useModelToolSearch { + return &modelToolSearchTool{tools: tools} + } + return &toolSearchTool{tools: tools} +} + +type toolSearchArgs struct { + Query string `json:"query"` + MaxResults *int `json:"max_results,omitempty"` +} + +type toolSearchResult struct { + Matches []string `json:"matches"` +} + +type toolSearchTool struct { + tools map[string]*schema.ToolInfo +} + +func (t *toolSearchTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return getToolSearchToolInfo(), nil +} + +func (t *toolSearchTool) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + matches, err := search(argumentsInJSON, t.tools) + if err != nil { + return "", err + } + result := &toolSearchResult{} + for _, m := range matches { + result.Matches = append(result.Matches, m.Name) + } + b, err := json.Marshal(result) + if err != nil { + return "", fmt.Errorf("failed to marshal tool search result: %w", err) + } + return string(b), nil +} + +type modelToolSearchTool struct { + tools map[string]*schema.ToolInfo +} + +func (t *modelToolSearchTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return getToolSearchToolInfo(), nil +} + +func (t *modelToolSearchTool) InvokableRun(_ context.Context, argumentsInJSON *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + ret, err := search(argumentsInJSON.Text, t.tools) + if err != nil { + return nil, err + } + + return &schema.ToolResult{Parts: []schema.ToolOutputPart{ + { + Type: schema.ToolPartTypeToolSearchResult, + ToolSearchResult: &schema.ToolSearchResult{ + Tools: ret, + }, + }, + }}, nil +} + +const ( + toolSearchToolName = "tool_search" + defaultMaxResults = 5 +) + +func getToolSearchToolInfo() *schema.ToolInfo { + return &schema.ToolInfo{ + Name: toolSearchToolName, + Desc: internal.SelectPrompt(internal.I18nPrompts{ + English: toolDescription, + Chinese: toolDescriptionChinese, + }), + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": { + Type: schema.String, + Desc: "Query to find deferred tools. Use \"select:\" for direct selection, or keywords to search.", + Required: true, + }, + "max_results": { + Type: schema.Integer, + Desc: "Maximum number of results to return (default: 5)", + Required: false, + }, + }), + } +} + +func search(argumentsInJSON string, tools map[string]*schema.ToolInfo) ([]*schema.ToolInfo, error) { + var args toolSearchArgs + if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil { + return nil, fmt.Errorf("failed to unmarshal tool search arguments: %w", err) + } + + query := strings.TrimSpace(args.Query) + if query == "" { + return nil, fmt.Errorf("query is required") + } + + maxResults := defaultMaxResults + if args.MaxResults != nil && *args.MaxResults > 0 { + maxResults = *args.MaxResults + } + + var matches []string + + // Direct selection mode: select:tool1,tool2 + // max_results is intentionally not applied here because the model has + // already specified the exact tools it wants by name. + if strings.HasPrefix(query, "select:") { + names := strings.Split(strings.TrimPrefix(query, "select:"), ",") + toolSet := make(map[string]bool, len(tools)) + for name := range tools { + toolSet[name] = true + } + for _, name := range names { + name = strings.TrimSpace(name) + if name != "" && toolSet[name] { + matches = append(matches, name) + } + } + } else { + matches = keywordSearch(query, maxResults, tools) + } + + ret := make([]*schema.ToolInfo, 0, len(matches)) + for _, name := range matches { + ti, ok := tools[name] + if !ok { + continue + } + ret = append(ret, ti) + } + return ret, nil +} + +func intMax(a, b int) int { + if a > b { + return a + } + return b +} + +func intMin(a, b int) int { + if a < b { + return a + } + return b +} + +// scoredTool pairs a tool name with its search score. +type scoredTool struct { + name string + score int +} + +// keywordSearch scores all tools against the query keywords and returns the top N. +func keywordSearch(query string, maxResults int, tools map[string]*schema.ToolInfo) []string { + keywords := parseKeywords(query) + if len(keywords) == 0 { + return nil + } + + var scored []scoredTool + + for name, tm := range tools { + nameParts := splitToolName(name) + nameLower := strings.ToLower(name) + descLower := strings.ToLower(tm.Desc) + + totalScore := 0 + allRequiredFound := true + + for _, kw := range keywords { + kwLower := strings.ToLower(kw.word) + kwScore := 0 + + // Score against name parts + for _, part := range nameParts { + partLower := strings.ToLower(part) + if partLower == kwLower { + kwScore = intMax(kwScore, 10) + } else if strings.Contains(partLower, kwLower) { + kwScore = intMax(kwScore, 5) + } + } + + // Score against full name + if strings.Contains(nameLower, kwLower) { + kwScore = intMax(kwScore, 3) + } + + // Score against description (substring match) + if descLower != "" && strings.Contains(descLower, kwLower) { + kwScore = intMax(kwScore, 2) + } + + if kw.required && kwScore == 0 { + allRequiredFound = false + break + } + + totalScore += kwScore + } + + if !allRequiredFound { + continue + } + + if totalScore > 0 { + scored = append(scored, scoredTool{name: name, score: totalScore}) + } + } + + // Sort by score descending, then by name for stability + sort.Slice(scored, func(i, j int) bool { + if scored[i].score != scored[j].score { + return scored[i].score > scored[j].score + } + return scored[i].name < scored[j].name + }) + + results := make([]string, 0, intMin(maxResults, len(scored))) + for i := 0; i < len(scored) && i < maxResults; i++ { + results = append(results, scored[i].name) + } + return results +} + +// keyword represents a parsed search keyword. +type keyword struct { + word string + required bool +} + +// parseKeywords splits a query string into keywords, handling the '+' required prefix. +func parseKeywords(query string) (keywords []keyword) { + parts := strings.Fields(query) + for _, p := range parts { + if strings.HasPrefix(p, "+") { + word := strings.TrimPrefix(p, "+") + if word != "" { + keywords = append(keywords, keyword{word: word, required: true}) + } + } else if p != "" { + keywords = append(keywords, keyword{word: p, required: false}) + } + } + return +} + +// splitToolName splits a tool name into parts by underscores, double underscores (MCP separator), +// and camelCase boundaries. +func splitToolName(name string) []string { + // First split by double underscore (MCP server__tool separator) + segments := strings.Split(name, "__") + + var parts []string + for _, seg := range segments { + // Split each segment by single underscore + underscoreParts := strings.Split(seg, "_") + for _, up := range underscoreParts { + if up == "" { + continue + } + // Further split by camelCase + camelParts := splitCamelCase(up) + parts = append(parts, camelParts...) + } + } + return parts +} + +// splitCamelCase splits a camelCase or PascalCase string into its constituent words. +func splitCamelCase(s string) []string { + if s == "" { + return nil + } + + var parts []string + runes := []rune(s) + start := 0 + + for i := 1; i < len(runes); i++ { + if unicode.IsUpper(runes[i]) { + if unicode.IsLower(runes[i-1]) { + parts = append(parts, string(runes[start:i])) + start = i + } else if i+1 < len(runes) && unicode.IsLower(runes[i+1]) { + parts = append(parts, string(runes[start:i])) + start = i + } + } + } + parts = append(parts, string(runes[start:])) + + return parts +} diff --git a/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go b/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go new file mode 100644 index 0000000..a659f07 --- /dev/null +++ b/adk/middlewares/dynamictool/toolsearch/toolsearch_generic_test.go @@ -0,0 +1,404 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package toolsearch + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// --------------------------------------------------------------------------- +// Generic table-driven tests covering both *schema.Message and *schema.AgenticMessage +// --------------------------------------------------------------------------- + +// --- Generic message construction helpers --- + +func makeUserMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(content)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(content)).(M) + default: + panic("unreachable") + } +} + +func makeSystemMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{Role: schema.System, Content: content}).(M) + case *schema.AgenticMessage: + return any(schema.SystemAgenticMessage(content)).(M) + default: + panic("unreachable") + } +} + +type testToolCall struct { + ID string + Name string + Arguments string +} + +func makeAssistantMsgWithToolCalls[M adk.MessageType](toolCalls []testToolCall) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + tcs := make([]schema.ToolCall, len(toolCalls)) + for i, tc := range toolCalls { + tcs[i] = schema.ToolCall{ + ID: tc.ID, + Function: schema.FunctionCall{Name: tc.Name, Arguments: tc.Arguments}, + } + } + return any(schema.AssistantMessage("", tcs)).(M) + case *schema.AgenticMessage: + blocks := make([]*schema.ContentBlock, len(toolCalls)) + for i, tc := range toolCalls { + blocks[i] = schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: tc.ID, + Name: tc.Name, + Arguments: tc.Arguments, + }) + } + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: blocks, + }).(M) + default: + panic("unreachable") + } +} + +func makeToolResultMsg[M adk.MessageType](content string, callID string, toolName string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{ + Role: schema.Tool, + ToolName: toolName, + ToolCallID: callID, + Content: content, + }).(M) + case *schema.AgenticMessage: + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: callID, + Name: toolName, + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: content}}, + }, + }), + }, + }).(M) + default: + panic("unreachable") + } +} + +func getMsgRole[M adk.MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return string(v.Role) + case *schema.AgenticMessage: + return string(v.Role) + default: + panic("unreachable") + } +} + +func getMsgContent[M adk.MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return v.Content + case *schema.AgenticMessage: + for _, block := range v.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeUserInputText && block.UserInputText != nil { + return block.UserInputText.Text + } + } + return "" + default: + panic("unreachable") + } +} + +func getMsgExtra[M adk.MessageType](msg M) map[string]any { + switch v := any(msg).(type) { + case *schema.Message: + return v.Extra + case *schema.AgenticMessage: + return v.Extra + default: + panic("unreachable") + } +} + +func setMsgExtra[M adk.MessageType](msg M, key string, val any) { + switch v := any(msg).(type) { + case *schema.Message: + if v.Extra == nil { + v.Extra = make(map[string]any) + } + v.Extra[key] = val + case *schema.AgenticMessage: + if v.Extra == nil { + v.Extra = make(map[string]any) + } + v.Extra[key] = val + default: + panic("unreachable") + } +} + +func newTestMiddlewareTyped[M adk.MessageType](t *testing.T, tools []tool.BaseTool) *typedMiddleware[M] { + t.Helper() + ctx := context.Background() + mw, err := NewTyped[M](ctx, &Config{ + DynamicTools: tools, + UseModelToolSearch: false, + }) + require.NoError(t, err) + return mw.(*typedMiddleware[M]) +} + +func countRemindersGeneric[M adk.MessageType](msgs []M) int { + count := 0 + for _, msg := range msgs { + extra := getMsgExtra(msg) + if extra != nil { + if v, _ := extra[toolSearchReminderExtraKey].(bool); v { + count++ + } + } + } + return count +} + +// --- Generic test functions --- + +func testEnsureReminderGeneric[M adk.MessageType](t *testing.T) { + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA}) + + t.Run("normal: system then user", func(t *testing.T) { + input := []M{ + makeSystemMsg[M]("sys"), + makeUserMsg[M]("hi"), + } + got := m.ensureReminder(input) + require.Len(t, got, 3) + assert.Equal(t, "system", getMsgRole(got[0])) + // Reminder inserted after system + extra := getMsgExtra(got[1]) + require.NotNil(t, extra) + assert.Equal(t, true, extra[toolSearchReminderExtraKey]) + assert.Equal(t, "hi", getMsgContent(got[2])) + }) + + t.Run("all system messages", func(t *testing.T) { + input := []M{ + makeSystemMsg[M]("sys1"), + makeSystemMsg[M]("sys2"), + } + got := m.ensureReminder(input) + require.Len(t, got, 3) + assert.Equal(t, "system", getMsgRole(got[0])) + assert.Equal(t, "system", getMsgRole(got[1])) + // Reminder appended at end + extra := getMsgExtra(got[2]) + require.NotNil(t, extra) + assert.Equal(t, true, extra[toolSearchReminderExtraKey]) + }) + + t.Run("empty input", func(t *testing.T) { + got := m.ensureReminder(nil) + require.Len(t, got, 1) + extra := getMsgExtra(got[0]) + require.NotNil(t, extra) + assert.Equal(t, true, extra[toolSearchReminderExtraKey]) + }) + + t.Run("no system messages", func(t *testing.T) { + input := []M{ + makeUserMsg[M]("hi"), + } + got := m.ensureReminder(input) + require.Len(t, got, 2) + // Reminder inserted at position 0 + extra := getMsgExtra(got[0]) + require.NotNil(t, extra) + assert.Equal(t, true, extra[toolSearchReminderExtraKey]) + assert.Equal(t, "hi", getMsgContent(got[1])) + }) + + t.Run("idempotent: does not insert twice", func(t *testing.T) { + reminder := makeUserMsg[M]("") + setMsgExtra(reminder, toolSearchReminderExtraKey, true) + input := []M{ + reminder, + makeUserMsg[M]("hi"), + } + got := m.ensureReminder(input) + require.Len(t, got, 2) + assert.Equal(t, "hi", getMsgContent(got[1])) + }) +} + +func testMode1InitializationGeneric[M adk.MessageType](t *testing.T) { + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + + m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA, dynamicB}) + + ctx := context.Background() + + state := &adk.TypedChatModelAgentState[M]{ + Messages: []M{ + makeSystemMsg[M]("sys"), + makeUserMsg[M]("hello"), + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + ti("dynamic_tool_a", "Dynamic tool A"), + ti("dynamic_tool_b", "Dynamic tool B"), + }, + } + + // Initialization strips dynamic tools, keeps tool_search and static tools. + _, state, err := m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names := toolNames(state.ToolInfos) + assert.Equal(t, []string{"static_tool", "tool_search"}, names) + assert.Nil(t, state.DeferredToolInfos, "Mode 1 should not populate DeferredToolInfos") + + // Verify reminder was inserted. + assert.Equal(t, 1, countRemindersGeneric(state.Messages), "reminder should be inserted") +} + +func testMode1ForwardSelectionGeneric[M adk.MessageType](t *testing.T) { + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + + m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA, dynamicB}) + + ctx := context.Background() + + // Simulate state AFTER initialization (dynamic tools already stripped). + // Include a tool_search result message that selected dynamic_tool_a. + toolSearchResultJSON, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_a"}}) + + // Build the reminder message with the extra marker + reminderMsg := makeUserMsg[M]("hello") + setMsgExtra(reminderMsg, toolSearchReminderExtraKey, true) + + state := &adk.TypedChatModelAgentState[M]{ + Messages: []M{ + makeSystemMsg[M]("sys"), + reminderMsg, + makeAssistantMsgWithToolCalls[M]([]testToolCall{ + {ID: "tc1", Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}, + }), + makeToolResultMsg[M](string(toolSearchResultJSON), "tc1", toolSearchToolName), + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + }, + } + + // Forward selection should add dynamic_tool_a from the tool_search result. + _, state, err := m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names := toolNames(state.ToolInfos) + assert.Equal(t, []string{"dynamic_tool_a", "static_tool", "tool_search"}, names) + + // Call again: forward selection should be idempotent (dynamic_tool_a already present). + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names = toolNames(state.ToolInfos) + assert.Equal(t, []string{"dynamic_tool_a", "static_tool", "tool_search"}, names) +} + +func testMalformedJSONGeneric[M adk.MessageType](t *testing.T) { + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + + m := newTestMiddlewareTyped[M](t, []tool.BaseTool{dynamicA}) + + ctx := context.Background() + + // Build the reminder message with the extra marker + reminderMsg := makeUserMsg[M]("reminder") + setMsgExtra(reminderMsg, toolSearchReminderExtraKey, true) + + state := &adk.TypedChatModelAgentState[M]{ + Messages: []M{ + makeSystemMsg[M]("sys"), + reminderMsg, + makeAssistantMsgWithToolCalls[M]([]testToolCall{ + {ID: "tc1", Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}, + }), + makeToolResultMsg[M](`{invalid json!!!`, "tc1", toolSearchToolName), + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + }, + } + + _, state, err := m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err, "malformed JSON in tool_search result should not cause an error") + + names := toolNames(state.ToolInfos) + assert.NotContains(t, names, "dynamic_tool_a", "malformed JSON result should be skipped") + assert.Contains(t, names, "static_tool") + assert.Contains(t, names, "tool_search") +} + +// --- Top-level generic test runner --- + +func TestToolSearchGeneric(t *testing.T) { + t.Run("Message", func(t *testing.T) { + t.Run("EnsureReminder", testEnsureReminderGeneric[*schema.Message]) + t.Run("Mode1Init", testMode1InitializationGeneric[*schema.Message]) + t.Run("Mode1ForwardSelection", testMode1ForwardSelectionGeneric[*schema.Message]) + t.Run("MalformedJSON", testMalformedJSONGeneric[*schema.Message]) + }) + t.Run("AgenticMessage", func(t *testing.T) { + t.Run("EnsureReminder", testEnsureReminderGeneric[*schema.AgenticMessage]) + t.Run("Mode1Init", testMode1InitializationGeneric[*schema.AgenticMessage]) + t.Run("Mode1ForwardSelection", testMode1ForwardSelectionGeneric[*schema.AgenticMessage]) + t.Run("MalformedJSON", testMalformedJSONGeneric[*schema.AgenticMessage]) + }) +} diff --git a/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go b/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go new file mode 100644 index 0000000..4bd1410 --- /dev/null +++ b/adk/middlewares/dynamictool/toolsearch/toolsearch_test.go @@ -0,0 +1,1042 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package toolsearch + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +func makeToolMap(tools ...*schema.ToolInfo) map[string]*schema.ToolInfo { + m := make(map[string]*schema.ToolInfo, len(tools)) + for _, t := range tools { + m[t.Name] = t + } + return m +} + +func ti(name, desc string) *schema.ToolInfo { + return &schema.ToolInfo{Name: name, Desc: desc} +} + +func toolNames(infos []*schema.ToolInfo) []string { + names := make([]string, len(infos)) + for i, info := range infos { + names[i] = info.Name + } + sort.Strings(names) + return names +} + +func searchJSON(query string, maxResults *int) string { + args := toolSearchArgs{Query: query, MaxResults: maxResults} + b, _ := json.Marshal(args) + return string(b) +} + +func intPtr(v int) *int { return &v } + +// --------------------------------------------------------------------------- +// TestSearch — unit tests for the search() function +// --------------------------------------------------------------------------- + +func TestSearch(t *testing.T) { + tools := makeToolMap( + ti("get_weather", "Get current weather for a city"), + ti("search_flights", "Search available flights"), + ti("mcp__slack__send_message", "Send a message to Slack channel"), + ti("mcp__slack__read_channel", "Read messages from Slack channel"), + ti("create_calendar_event", "Create a new calendar event"), + ti("NotebookEdit", "Edit Jupyter notebook cells"), + ) + + tests := []struct { + name string + json string + wantNames []string // sorted; nil means expect empty + wantErr bool + }{ + { + name: "keyword exact name part match", + json: searchJSON("weather", nil), + wantNames: []string{"get_weather"}, + }, + { + name: "keyword matches multiple tools", + json: searchJSON("slack", nil), + wantNames: []string{"mcp__slack__read_channel", "mcp__slack__send_message"}, + }, + { + name: "multi-word ranking - send_message ranked first", + json: searchJSON("send message", nil), + wantNames: []string{"mcp__slack__send_message"}, // check first element only + }, + { + name: "required keyword filters to slack only", + json: searchJSON("+slack send", nil), + wantNames: []string{"mcp__slack__read_channel", "mcp__slack__send_message"}, + }, + { + name: "required keyword no match", + json: searchJSON("+github send", nil), + wantNames: nil, + }, + { + name: "direct select single", + json: searchJSON("select:get_weather", nil), + wantNames: []string{"get_weather"}, + }, + { + name: "direct select multiple", + json: searchJSON("select:get_weather,NotebookEdit", nil), + wantNames: []string{"NotebookEdit", "get_weather"}, + }, + { + name: "direct select nonexistent", + json: searchJSON("select:nonexistent", nil), + wantNames: nil, + }, + { + name: "max_results limits output", + json: searchJSON("slack", intPtr(1)), + wantNames: []string{"mcp__slack__read_channel"}, // just check length below + }, + { + name: "camelCase split matches notebook", + json: searchJSON("notebook", nil), + wantNames: []string{"NotebookEdit"}, + }, + { + name: "empty query returns error", + json: searchJSON("", nil), + wantErr: true, + }, + { + name: "description match - jupyter", + json: searchJSON("jupyter", nil), + wantNames: []string{"NotebookEdit"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := search(tt.json, tools) + if tt.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + + // special case: max_results limit + if tt.name == "max_results limits output" { + assert.Len(t, got, 1) + return + } + + // special case: ranking — just check first element + if tt.name == "multi-word ranking - send_message ranked first" { + require.NotEmpty(t, got) + assert.Equal(t, "mcp__slack__send_message", got[0].Name) + return + } + + gotNames := toolNames(got) + if tt.wantNames == nil { + assert.Empty(t, gotNames) + } else { + assert.Equal(t, tt.wantNames, gotNames) + } + }) + } +} + +// --------------------------------------------------------------------------- +// TestMiddlewareFlow — integration test for UseModelToolSearch=false +// --------------------------------------------------------------------------- + +// simpleTool is a minimal InvokableTool for testing. +type simpleTool struct { + name string + desc string + called bool + mu sync.Mutex +} + +func (s *simpleTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: s.name, + Desc: s.desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String, Desc: "input", Required: true}, + }), + }, nil +} + +func (s *simpleTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + s.mu.Lock() + s.called = true + s.mu.Unlock() + return `{"result":"ok"}`, nil +} + +func (s *simpleTool) wasCalled() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.called +} + +// mockChatModel implements model.ToolCallingChatModel. +// It drives a 3-turn conversation: +// +// Turn 1: call tool_search with select:dynamic_tool_a +// Turn 2: call dynamic_tool_a +// Turn 3: return final text +type mockChatModel struct { + mu sync.Mutex + generateCall int + // toolsPerCall records the tool names passed via model.WithTools for each Generate call. + toolsPerCall [][]string +} + +func (m *mockChatModel) Generate(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) { + options := model.GetCommonOptions(nil, opts...) + var names []string + for _, t := range options.Tools { + names = append(names, t.Name) + } + sort.Strings(names) + + m.mu.Lock() + m.generateCall++ + call := m.generateCall + m.toolsPerCall = append(m.toolsPerCall, names) + m.mu.Unlock() + + switch call { + case 1: + // Ask tool_search to select dynamic_tool_a + return schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "tc1", + Function: schema.FunctionCall{ + Name: toolSearchToolName, + Arguments: `{"query":"select:dynamic_tool_a","max_results":5}`, + }, + }, + }), nil + case 2: + // Call dynamic_tool_a + return schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "tc2", + Function: schema.FunctionCall{ + Name: "dynamic_tool_a", + Arguments: `{"input":"hello"}`, + }, + }, + }), nil + default: + // Final response + return schema.AssistantMessage("done", nil), nil + } +} + +func (m *mockChatModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, fmt.Errorf("not implemented") +} + +func (m *mockChatModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +func (m *mockChatModel) getToolsPerCall() [][]string { + m.mu.Lock() + defer m.mu.Unlock() + ret := make([][]string, len(m.toolsPerCall)) + copy(ret, m.toolsPerCall) + return ret +} + +func TestMiddlewareFlow(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + staticTool := &simpleTool{name: "static_tool", desc: "Static tool"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA, dynamicB}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + cm := &mockChatModel{} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "test_agent", + Description: "test", + Instruction: "you are a test agent", + Model: cm, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{staticTool}, + }, + }, + Handlers: []adk.ChatModelAgentMiddleware{mw}, + }) + require.NoError(t, err) + + input := &adk.AgentInput{ + Messages: []adk.Message{schema.UserMessage("test")}, + } + iter := agent.Run(ctx, input) + + var events []*adk.AgentEvent + for { + ev, ok := iter.Next() + if !ok { + break + } + events = append(events, ev) + } + + // Verify no error event. + for _, ev := range events { + if ev.Err != nil { + t.Fatalf("unexpected error event: %v", ev.Err) + } + } + + // Verify final output is "done". + lastEvent := events[len(events)-1] + require.NotNil(t, lastEvent.Output) + require.NotNil(t, lastEvent.Output.MessageOutput) + assert.Equal(t, "done", lastEvent.Output.MessageOutput.Message.Content) + + // Verify dynamic_tool_a was actually called. + assert.True(t, dynamicA.wasCalled(), "dynamic_tool_a should have been called") + assert.False(t, dynamicB.wasCalled(), "dynamic_tool_b should not have been called") + + // Verify tool lists per Generate call. + toolsPerCall := cm.getToolsPerCall() + require.Len(t, toolsPerCall, 3, "expected 3 Generate calls") + + // Call 1: static_tool visible; dynamic tools are hidden. + assert.Contains(t, toolsPerCall[0], "static_tool") + assert.NotContains(t, toolsPerCall[0], "dynamic_tool_a") + assert.NotContains(t, toolsPerCall[0], "dynamic_tool_b") + + // Call 2: after selecting dynamic_tool_a, it becomes visible. + assert.Contains(t, toolsPerCall[1], "static_tool") + assert.Contains(t, toolsPerCall[1], "dynamic_tool_a") + assert.NotContains(t, toolsPerCall[1], "dynamic_tool_b") + + // Call 3: same as call 2. + assert.Contains(t, toolsPerCall[2], "static_tool") + assert.Contains(t, toolsPerCall[2], "dynamic_tool_a") + assert.NotContains(t, toolsPerCall[2], "dynamic_tool_b") + + // Verify reminder is present in messages (checked via tool list — the wrapper inserts it). + // The model received messages, and the reminder contains "". + // We indirectly verify this by checking that the middleware ran without error and the + // 3-turn flow completed successfully, which requires the tool_search tool to work. + + // Additional: verify that the reminder contains the dynamic tool names. + mwImpl := mw.(*typedMiddleware[*schema.Message]) + assert.True(t, strings.Contains(mwImpl.sr, "dynamic_tool_a")) + assert.True(t, strings.Contains(mwImpl.sr, "dynamic_tool_b")) + assert.True(t, strings.Contains(mwImpl.sr, "")) +} + +// --------------------------------------------------------------------------- +// TestNew — error paths for New() +// --------------------------------------------------------------------------- + +func TestNew(t *testing.T) { + ctx := context.Background() + + t.Run("nil config", func(t *testing.T) { + _, err := New(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config is required") + }) + + t.Run("empty DynamicTools", func(t *testing.T) { + _, err := New(ctx, &Config{}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "tools is required") + }) + + t.Run("success", func(t *testing.T) { + st := &simpleTool{name: "t1", desc: "tool 1"} + mw, err := New(ctx, &Config{DynamicTools: []tool.BaseTool{st}}) + require.NoError(t, err) + assert.NotNil(t, mw) + }) +} + +// --------------------------------------------------------------------------- +// TestSplitCamelCase +// --------------------------------------------------------------------------- + +func TestSplitCamelCase(t *testing.T) { + tests := []struct { + input string + want []string + }{ + {"", nil}, + {"hello", []string{"hello"}}, + {"NotebookEdit", []string{"Notebook", "Edit"}}, + {"camelCase", []string{"camel", "Case"}}, + {"HTMLParser", []string{"HTML", "Parser"}}, + {"getURL", []string{"get", "URL"}}, + {"A", []string{"A"}}, + {"AB", []string{"AB"}}, + {"HTTP", []string{"HTTP"}}, + } + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got := splitCamelCase(tt.input) + assert.Equal(t, tt.want, got) + }) + } +} + +// --------------------------------------------------------------------------- +// TestEnsureReminder +// --------------------------------------------------------------------------- + +func TestEnsureReminder(t *testing.T) { + m := &typedMiddleware[*schema.Message]{sr: ""} + + t.Run("normal: system then user", func(t *testing.T) { + input := []*schema.Message{ + {Role: schema.System, Content: "sys"}, + {Role: schema.User, Content: "hi"}, + } + got := m.ensureReminder(input) + require.Len(t, got, 3) + assert.Equal(t, schema.System, got[0].Role) + assert.Equal(t, schema.User, got[1].Role) + assert.Equal(t, "", got[1].Content) + assert.Equal(t, true, got[1].Extra[toolSearchReminderExtraKey]) + assert.Equal(t, schema.User, got[2].Role) + assert.Equal(t, "hi", got[2].Content) + }) + + t.Run("all system messages", func(t *testing.T) { + input := []*schema.Message{ + {Role: schema.System, Content: "sys1"}, + {Role: schema.System, Content: "sys2"}, + } + got := m.ensureReminder(input) + require.Len(t, got, 3) + assert.Equal(t, schema.System, got[0].Role) + assert.Equal(t, schema.System, got[1].Role) + assert.Equal(t, "", got[2].Content) + }) + + t.Run("empty input", func(t *testing.T) { + got := m.ensureReminder(nil) + require.Len(t, got, 1) + assert.Equal(t, "", got[0].Content) + }) + + t.Run("no system messages", func(t *testing.T) { + input := []*schema.Message{ + {Role: schema.User, Content: "hi"}, + {Role: schema.Assistant, Content: "hello"}, + } + got := m.ensureReminder(input) + require.Len(t, got, 3) + assert.Equal(t, "", got[0].Content) + assert.Equal(t, "hi", got[1].Content) + assert.Equal(t, "hello", got[2].Content) + }) + + t.Run("idempotent: does not insert twice", func(t *testing.T) { + input := []*schema.Message{ + {Role: schema.User, Content: "", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + {Role: schema.User, Content: "hi"}, + } + got := m.ensureReminder(input) + require.Len(t, got, 2) + assert.Equal(t, "", got[0].Content) + assert.Equal(t, "hi", got[1].Content) + }) +} + +// --------------------------------------------------------------------------- +// TestHelperFunctions +// --------------------------------------------------------------------------- + +func TestHelperFunctions(t *testing.T) { + t.Run("extractDynamicTools", func(t *testing.T) { + m := &typedMiddleware[*schema.Message]{ + mapOfDynamicTools: map[string]*schema.ToolInfo{ + "dyn_a": ti("dyn_a", "A"), + "dyn_b": ti("dyn_b", "B"), + }, + } + tools := []*schema.ToolInfo{ti("static", "S"), ti("dyn_a", "A"), ti("dyn_b", "B")} + got := m.extractDynamicTools(tools) + assert.Len(t, got, 2) + names := toolNames(got) + assert.Equal(t, []string{"dyn_a", "dyn_b"}, names) + }) + + t.Run("stripDynamicTools", func(t *testing.T) { + m := &typedMiddleware[*schema.Message]{ + mapOfDynamicTools: map[string]*schema.ToolInfo{ + "dyn_a": ti("dyn_a", "A"), + "dyn_b": ti("dyn_b", "B"), + }, + } + tools := []*schema.ToolInfo{ti("static", "S"), ti("dyn_a", "A"), ti("tool_search", "TS")} + got := m.stripDynamicTools(tools) + names := toolNames(got) + assert.Equal(t, []string{"static", "tool_search"}, names) + }) + + t.Run("removeTool", func(t *testing.T) { + tools := []*schema.ToolInfo{ti("a", "A"), ti("b", "B"), ti("c", "C")} + got := removeTool(tools, "b") + names := toolNames(got) + assert.Equal(t, []string{"a", "c"}, names) + }) + + t.Run("toolNameSet", func(t *testing.T) { + tools := []*schema.ToolInfo{ti("x", "X"), ti("y", "Y")} + got := toolNameSet(tools) + assert.True(t, got["x"]) + assert.True(t, got["y"]) + assert.False(t, got["z"]) + }) +} + +// --------------------------------------------------------------------------- +// TestBeforeModelRewriteState — direct unit tests for BeforeModelRewriteState +// --------------------------------------------------------------------------- + +// Note: these tests call BeforeModelRewriteState without a full compose context, +// so RunLocalValue (used by isInitialized/markInitialized) always returns error. +// This means every call re-runs the initialization block. Tests are designed +// accordingly: they test single-call behavior or provide pre-initialized state. + +func TestBeforeModelRewriteState_Mode1_Initialization(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA, dynamicB}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + // Simulate state: static_tool + tool_search + dynamic tools (as would come from backfill). + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.System, Content: "sys"}, + {Role: schema.User, Content: "hello"}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + ti("dynamic_tool_a", "Dynamic tool A"), + ti("dynamic_tool_b", "Dynamic tool B"), + }, + } + + // Initialization strips dynamic tools, keeps tool_search and static tools. + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names := toolNames(state.ToolInfos) + assert.Equal(t, []string{"static_tool", "tool_search"}, names) + assert.Nil(t, state.DeferredToolInfos, "Mode 1 should not populate DeferredToolInfos") + + // Verify reminder was inserted. + assert.Equal(t, 1, countReminders(state.Messages), "reminder should be inserted") +} + +func TestBeforeModelRewriteState_Mode1_ForwardSelection(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA, dynamicB}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + // Simulate state AFTER initialization (dynamic tools already stripped). + // Include a tool_search result message that selected dynamic_tool_a. + toolSearchResultJSON, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_a"}}) + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.System, Content: "sys"}, + {Role: schema.User, Content: "hello", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: string(toolSearchResultJSON)}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + }, + } + + // Forward selection should add dynamic_tool_a from the tool_search result. + // Note: init block runs (no compose ctx) but ToolInfos has no dynamic tools to strip. + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names := toolNames(state.ToolInfos) + assert.Equal(t, []string{"dynamic_tool_a", "static_tool", "tool_search"}, names) + + // Call again: forward selection should be idempotent (dynamic_tool_a already present). + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names = toolNames(state.ToolInfos) + assert.Equal(t, []string{"dynamic_tool_a", "static_tool", "tool_search"}, names) +} + +func TestBeforeModelRewriteState_Mode2_DeferredToolInfos(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA, dynamicB}, + UseModelToolSearch: true, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.User, Content: "hello"}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + ti("dynamic_tool_a", "Dynamic tool A"), + ti("dynamic_tool_b", "Dynamic tool B"), + }, + } + + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + // Mode 2: static tools in ToolInfos (tool_search removed), dynamic in DeferredToolInfos. + names := toolNames(state.ToolInfos) + assert.Equal(t, []string{"static_tool"}, names, "ToolInfos should only have static tools") + + deferredNames := toolNames(state.DeferredToolInfos) + assert.Equal(t, []string{"dynamic_tool_a", "dynamic_tool_b"}, deferredNames, "DeferredToolInfos should have all dynamic tools") +} + +func TestBeforeModelRewriteState_ReminderReinsertAfterRemoval(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.User, Content: "hello"}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + ti("dynamic_tool_a", "Dynamic tool A"), + }, + } + + // First call: reminder inserted. + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + reminderCount := countReminders(state.Messages) + assert.Equal(t, 1, reminderCount) + + // Simulate summarization removing the reminder message. + var msgsWithoutReminder []*schema.Message + for _, msg := range state.Messages { + isReminder := false + if msg.Extra != nil { + if v, ok := msg.Extra[toolSearchReminderExtraKey].(bool); ok && v { + isReminder = true + } + } + if !isReminder { + msgsWithoutReminder = append(msgsWithoutReminder, msg) + } + } + state.Messages = msgsWithoutReminder + assert.Equal(t, 0, countReminders(state.Messages), "reminder should be gone") + + // Next call: reminder should be re-inserted. + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + reminderCount = countReminders(state.Messages) + assert.Equal(t, 1, reminderCount, "reminder should be re-inserted after removal") +} + +func countReminders(msgs []*schema.Message) int { + count := 0 + for _, msg := range msgs { + if msg.Extra != nil { + if v, _ := msg.Extra[toolSearchReminderExtraKey].(bool); v { + count++ + } + } + } + return count +} + +// --------------------------------------------------------------------------- +// Edge-case tests for BeforeModelRewriteState +// --------------------------------------------------------------------------- + +func TestBeforeModelRewriteState_Mode1_MultipleToolSearchResultsAcrossTurns(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + dynamicC := &simpleTool{name: "dynamic_tool_c", desc: "Dynamic tool C"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA, dynamicB, dynamicC}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + // Build two separate tool_search result messages, each selecting a different tool. + resultA, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_a"}}) + resultB, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_b"}}) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.System, Content: "sys"}, + {Role: schema.User, Content: "reminder", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultA)}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc2", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_b"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultB)}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + }, + } + + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names := toolNames(state.ToolInfos) + assert.Contains(t, names, "dynamic_tool_a", "dynamic_tool_a should be added from first tool_search result") + assert.Contains(t, names, "dynamic_tool_b", "dynamic_tool_b should be added from second tool_search result") + assert.NotContains(t, names, "dynamic_tool_c", "dynamic_tool_c was never selected") + assert.Contains(t, names, "static_tool", "static_tool should remain") + assert.Contains(t, names, "tool_search", "tool_search should remain") +} + +func TestBeforeModelRewriteState_Mode1_MalformedJSONInToolSearchResult(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.System, Content: "sys"}, + {Role: schema.User, Content: "reminder", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: `{invalid json!!!`}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + }, + } + + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err, "malformed JSON in tool_search result should not cause an error") + + names := toolNames(state.ToolInfos) + assert.NotContains(t, names, "dynamic_tool_a", "malformed JSON result should be skipped") + assert.Contains(t, names, "static_tool") + assert.Contains(t, names, "tool_search") +} + +func TestBeforeModelRewriteState_Mode1_NonExistentToolInForwardSelection(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + resultJSON, _ := json.Marshal(toolSearchResult{Matches: []string{"nonexistent_tool", "dynamic_tool_a"}}) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.User, Content: "reminder", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:nonexistent_tool,dynamic_tool_a"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + }, + } + + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err, "nonexistent tool in forward selection should not cause an error") + + names := toolNames(state.ToolInfos) + assert.Contains(t, names, "dynamic_tool_a", "valid tool should be added") + assert.NotContains(t, names, "nonexistent_tool", "nonexistent tool should be silently ignored") +} + +func TestBeforeModelRewriteState_Mode2_EmptyToolInfos(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA}, + UseModelToolSearch: true, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.User, Content: "hello"}, + }, + ToolInfos: []*schema.ToolInfo{}, // empty, not nil + } + + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err, "empty ToolInfos should not cause an error") + + assert.Empty(t, state.ToolInfos, "ToolInfos should be empty") + assert.Empty(t, state.DeferredToolInfos, "DeferredToolInfos should be empty when no dynamic tools found in ToolInfos") +} + +func TestBeforeModelRewriteState_Mode1_DoubleInitWithoutComposeContext(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + dynamicB := &simpleTool{name: "dynamic_tool_b", desc: "Dynamic tool B"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA, dynamicB}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + resultJSON, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_a"}}) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.User, Content: "reminder", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)}, + }, + ToolInfos: []*schema.ToolInfo{ + ti("static_tool", "Static tool"), + getToolSearchToolInfo(), + ti("dynamic_tool_a", "Dynamic tool A"), + }, + } + + // First call: init runs (strips dynamic_tool_a), then forward selection re-adds it. + _, state, err = m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names := toolNames(state.ToolInfos) + assert.Contains(t, names, "dynamic_tool_a", + "forward selection should re-add dynamic_tool_a even after init re-strips it") + assert.Contains(t, names, "static_tool") + assert.Contains(t, names, "tool_search") + + // Second call: init runs AGAIN (no compose ctx), verify behavior is stable. + _, state2, err := m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + names2 := toolNames(state2.ToolInfos) + assert.Contains(t, names2, "dynamic_tool_a", + "second call should also have dynamic_tool_a re-added by forward selection") +} + +func TestBeforeModelRewriteState_ToolInfosSliceMutation(t *testing.T) { + ctx := context.Background() + + dynamicA := &simpleTool{name: "dynamic_tool_a", desc: "Dynamic tool A"} + + mw, err := New(ctx, &Config{ + DynamicTools: []tool.BaseTool{dynamicA}, + UseModelToolSearch: false, + }) + require.NoError(t, err) + + m := mw.(*typedMiddleware[*schema.Message]) + + // Create ToolInfos with excess capacity so append could mutate in place. + originalToolInfos := make([]*schema.ToolInfo, 2, 10) + originalToolInfos[0] = ti("static_tool", "Static tool") + originalToolInfos[1] = getToolSearchToolInfo() + + originalLen := len(originalToolInfos) + + resultJSON, _ := json.Marshal(toolSearchResult{Matches: []string{"dynamic_tool_a"}}) + + state := &adk.ChatModelAgentState{ + Messages: []*schema.Message{ + {Role: schema.User, Content: "reminder", Extra: map[string]any{toolSearchReminderExtraKey: true}}, + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "tc1", Function: schema.FunctionCall{Name: toolSearchToolName, Arguments: `{"query":"select:dynamic_tool_a"}`}}, + }), + {Role: schema.Tool, ToolName: toolSearchToolName, Content: string(resultJSON)}, + }, + ToolInfos: originalToolInfos, + } + + _, newState, err := m.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + + newNames := toolNames(newState.ToolInfos) + assert.Contains(t, newNames, "dynamic_tool_a") + assert.Equal(t, originalLen, len(originalToolInfos), + "original ToolInfos slice length should not be mutated by the middleware") +} + +// --------------------------------------------------------------------------- +// modelToolSearchTool (Mode 2) tests +// --------------------------------------------------------------------------- + +func TestModelToolSearchTool(t *testing.T) { + ctx := context.Background() + + tools := makeToolMap( + ti("alpha", "Alpha tool description"), + ti("beta", "Beta tool description"), + ) + mts := &modelToolSearchTool{tools: tools} + + // Info should return the standard tool_search tool info. + info, err := mts.Info(ctx) + require.NoError(t, err) + assert.Equal(t, toolSearchToolName, info.Name) + + // InvokableRun with a valid query selecting "alpha". + arg := &schema.ToolArgument{Text: searchJSON("select:alpha", nil)} + result, err := mts.InvokableRun(ctx, arg) + require.NoError(t, err) + require.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeToolSearchResult, result.Parts[0].Type) + require.NotNil(t, result.Parts[0].ToolSearchResult) + assert.Len(t, result.Parts[0].ToolSearchResult.Tools, 1) + assert.Equal(t, "alpha", result.Parts[0].ToolSearchResult.Tools[0].Name) + + // InvokableRun with an empty query should return error. + argEmpty := &schema.ToolArgument{Text: `{"query":""}`} + _, err = mts.InvokableRun(ctx, argEmpty) + assert.Error(t, err) +} + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + + // Verify that NewTyped compiles with *schema.AgenticMessage. + // DynamicTools is required, so we expect an error with an empty config. + mw, err := NewTyped[*schema.AgenticMessage](ctx, &Config{ + DynamicTools: []tool.BaseTool{&simpleTool{name: "t1", desc: "desc1"}}, + }) + assert.NoError(t, err) + assert.NotNil(t, mw) + + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} diff --git a/adk/middlewares/filesystem/backend.go b/adk/middlewares/filesystem/backend.go new file mode 100644 index 0000000..eec62f1 --- /dev/null +++ b/adk/middlewares/filesystem/backend.go @@ -0,0 +1,33 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package filesystem provides middlewares. +package filesystem + +import ( + "github.com/cloudwego/eino/adk/filesystem" +) + +type FileInfo = filesystem.FileInfo +type GrepMatch = filesystem.GrepMatch +type LsInfoRequest = filesystem.LsInfoRequest +type ReadRequest = filesystem.ReadRequest +type MultiModalReadRequest = filesystem.MultiModalReadRequest +type GrepRequest = filesystem.GrepRequest +type GlobInfoRequest = filesystem.GlobInfoRequest +type WriteRequest = filesystem.WriteRequest +type EditRequest = filesystem.EditRequest +type FileContent = filesystem.FileContent diff --git a/adk/middlewares/filesystem/filesystem.go b/adk/middlewares/filesystem/filesystem.go new file mode 100644 index 0000000..b9d64ab --- /dev/null +++ b/adk/middlewares/filesystem/filesystem.go @@ -0,0 +1,1245 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "io" + "path/filepath" + "runtime/debug" + "sort" + "strconv" + "strings" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +const ( + ToolNameLs = "ls" + ToolNameReadFile = "read_file" + ToolNameWriteFile = "write_file" + ToolNameEditFile = "edit_file" + ToolNameGlob = "glob" + ToolNameGrep = "grep" + ToolNameExecute = "execute" + + noFilesFound = "No files found" + noMatchesFound = "No matches found" +) + +// ToolConfig configures a filesystem tool +type ToolConfig struct { + // Name overrides the tool name used in tool registration + // optional, default tool name will be used if not set (empty string) + Name string + + // Desc overrides the tool description used in tool registration + // optional, default tool description will be used if not set (nil pointer) + Desc *string + + // CustomTool provides a custom implementation for this tool. + // If set, this custom tool will be used instead of the default implementation associated with Backend. + // If not set, the default tool implementation associated with Backend will be created automatically. + // optional + CustomTool tool.BaseTool + + // Disable disables this tool + // If true, the tool will not be registered + // optional, false by default + Disable bool +} + +// Config is the configuration for the filesystem middleware +type Config struct { + // Backend provides filesystem operations used by tools and offloading. + // If set, filesystem tools (read_file, write_file, edit_file, glob, grep) will be registered. + // At least one of Backend, Shell, or StreamingShell must be set. + Backend filesystem.Backend + + // Shell provides shell command execution capability. + // If set, an execute tool will be registered to support shell command execution. + // At least one of Backend, Shell, or StreamingShell must be set. + // Mutually exclusive with StreamingShell. + Shell filesystem.Shell + // StreamingShell provides streaming shell command execution capability. + // If set, a streaming execute tool will be registered to support streaming shell command execution. + // At least one of Backend, Shell, or StreamingShell must be set. + // Mutually exclusive with Shell. + StreamingShell filesystem.StreamingShell + + // LsToolConfig configures the ls tool + // optional + LsToolConfig *ToolConfig + // ReadFileToolConfig configures the read_file tool. + // This config applies to both the standard read_file tool (InvokableTool) and + // the multimodal read_file tool (EnhancedInvokableTool) when UseMultiModalRead is true. + // optional + ReadFileToolConfig *ToolConfig + // WriteFileToolConfig configures the write_file tool + // optional + WriteFileToolConfig *ToolConfig + // EditFileToolConfig configures the edit_file tool + // optional + EditFileToolConfig *ToolConfig + // GlobToolConfig configures the glob tool + // optional + GlobToolConfig *ToolConfig + // GrepToolConfig configures the grep tool + // optional + GrepToolConfig *ToolConfig + + // WithoutLargeToolResultOffloading disables automatic offloading of large tool result to Backend + // optional, false(enabled) by default + WithoutLargeToolResultOffloading bool + // LargeToolResultOffloadingTokenLimit sets the token threshold to trigger offloading + // optional, 20000 by default + LargeToolResultOffloadingTokenLimit int + // LargeToolResultOffloadingPathGen generates the write path for offloaded results based on context and ToolInput + // optional, "/large_tool_result/{ToolCallID}" by default + LargeToolResultOffloadingPathGen func(ctx context.Context, input *compose.ToolInput) (string, error) + + // CustomSystemPrompt overrides the default ToolsSystemPrompt appended to agent instruction + // optional, ToolsSystemPrompt by default + CustomSystemPrompt *string + + // CustomLsToolDesc overrides the ls tool description used in tool registration + // optional, ListFilesToolDesc by default + // Deprecated: Use LsToolConfig.Desc instead + CustomLsToolDesc *string + // CustomReadFileToolDesc overrides the read_file tool description + // optional, ReadFileToolDesc by default + // Deprecated: Use ReadFileToolConfig.Desc instead + CustomReadFileToolDesc *string + // CustomGrepToolDesc overrides the grep tool description + // optional, GrepToolDesc by default + // Deprecated: Use GrepToolConfig.Desc instead + CustomGrepToolDesc *string + // CustomGlobToolDesc overrides the glob tool description + // optional, GlobToolDesc by default + // Deprecated: Use GlobToolConfig.Desc instead + CustomGlobToolDesc *string + // CustomWriteFileToolDesc overrides the write_file tool description + // optional, WriteFileToolDesc by default + // Deprecated: Use WriteFileToolConfig.Desc instead + CustomWriteFileToolDesc *string + // CustomEditToolDesc overrides the edit_file tool description + // optional, EditFileToolDesc by default + // Deprecated: Use EditFileToolConfig.Desc instead + CustomEditToolDesc *string +} + +func (c *Config) Validate() error { + if c == nil { + return errors.New("config should not be nil") + } + if c.Backend == nil { + return errors.New("backend should not be nil") + } + if c.StreamingShell != nil && c.Shell != nil { + return errors.New("shell and streaming shell should not be both set") + } + return nil +} + +// NewMiddleware constructs and returns the filesystem middleware. +// +// Deprecated: Use New instead. New returns +// a ChatModelAgentMiddleware which provides better context propagation through wrapper methods +// and is the recommended approach for new code. See ChatModelAgentMiddleware documentation +// for details on the benefits over AgentMiddleware. +func NewMiddleware(ctx context.Context, config *Config) (adk.AgentMiddleware, error) { + err := config.Validate() + if err != nil { + return adk.AgentMiddleware{}, err + } + ts, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Backend: config.Backend, + Shell: config.Shell, + StreamingShell: config.StreamingShell, + LsToolConfig: config.LsToolConfig, + ReadFileToolConfig: config.ReadFileToolConfig, + WriteFileToolConfig: config.WriteFileToolConfig, + EditFileToolConfig: config.EditFileToolConfig, + GlobToolConfig: config.GlobToolConfig, + GrepToolConfig: config.GrepToolConfig, + CustomSystemPrompt: config.CustomSystemPrompt, + CustomLsToolDesc: config.CustomLsToolDesc, + CustomReadFileToolDesc: config.CustomReadFileToolDesc, + CustomGrepToolDesc: config.CustomGrepToolDesc, + CustomGlobToolDesc: config.CustomGlobToolDesc, + CustomWriteFileToolDesc: config.CustomWriteFileToolDesc, + CustomEditToolDesc: config.CustomEditToolDesc, + }) + if err != nil { + return adk.AgentMiddleware{}, err + } + + var systemPrompt string + if config.CustomSystemPrompt != nil { + systemPrompt = *config.CustomSystemPrompt + } + + m := adk.AgentMiddleware{ + AdditionalInstruction: systemPrompt, + AdditionalTools: ts, + } + + if !config.WithoutLargeToolResultOffloading { + m.WrapToolCall = newToolResultOffloading(ctx, &toolResultOffloadingConfig{ + Backend: config.Backend, + TokenLimit: config.LargeToolResultOffloadingTokenLimit, + PathGenerator: config.LargeToolResultOffloadingPathGen, + }) + } + + return m, nil +} + +// MiddlewareConfig is the configuration for the filesystem middleware +type MiddlewareConfig struct { + // Backend provides filesystem operations used by tools and offloading. + // required + Backend filesystem.Backend + + // Shell provides shell command execution capability. + // If set, an execute tool will be registered to support shell command execution. + // optional, mutually exclusive with StreamingShell + Shell filesystem.Shell + // StreamingShell provides streaming shell command execution capability. + // If set, a streaming execute tool will be registered for real-time output. + // optional, mutually exclusive with Shell + StreamingShell filesystem.StreamingShell + + // LsToolConfig configures the ls tool + // optional + LsToolConfig *ToolConfig + // ReadFileToolConfig configures the read_file tool. + // This config applies to both the standard read_file tool (InvokableTool) and + // the multimodal read_file tool (EnhancedInvokableTool) when UseMultiModalRead is true. + // optional + ReadFileToolConfig *ToolConfig + // WriteFileToolConfig configures the write_file tool + // optional + WriteFileToolConfig *ToolConfig + // EditFileToolConfig configures the edit_file tool + // optional + EditFileToolConfig *ToolConfig + // GlobToolConfig configures the glob tool + // optional + GlobToolConfig *ToolConfig + // GrepToolConfig configures the grep tool + // optional + GrepToolConfig *ToolConfig + + // UseMultiModalRead enables multimodal read_file tool (EnhancedInvokableTool). + // When true, read_file returns results via schema.ToolResult.Parts instead of plain text string. + // + // Requires Backend to implement filesystem.MultiModalReader interface. + // The default implementation supports reading image files (PNG, JPG, etc.) + // and PDF files with page range selection. + // + // If you provide a custom MultiModalReader, you may need to override + // ReadFileToolConfig.Desc to accurately describe your implementation's capabilities. + // The default description is composed of ReadFileToolDesc + EnhancedReadFileDescSuffix. + // + // Note: When enabled, the read_file tool becomes an EnhancedInvokableTool. + // If you use ChatModelAgentMiddleware, you must implement ChatModelAgentMiddleware.WrapEnhancedInvokableToolCall + // for the middleware to take effect on the read_file tool. + // + // Default false, preserving backward compatibility. + UseMultiModalRead bool + + // CustomSystemPrompt overrides the default ToolsSystemPrompt appended to agent instruction + // optional, ToolsSystemPrompt by default + CustomSystemPrompt *string + + // CustomLsToolDesc overrides the ls tool description used in tool registration + // optional, ListFilesToolDesc by default + // Deprecated: Use LsToolConfig.Desc instead + CustomLsToolDesc *string + // CustomReadFileToolDesc overrides the read_file tool description + // optional, ReadFileToolDesc by default + // Deprecated: Use ReadFileToolConfig.Desc instead + CustomReadFileToolDesc *string + // CustomGrepToolDesc overrides the grep tool description + // optional, GrepToolDesc by default + // Deprecated: Use GrepToolConfig.Desc instead + CustomGrepToolDesc *string + // CustomGlobToolDesc overrides the glob tool description + // optional, GlobToolDesc by default + // Deprecated: Use GlobToolConfig.Desc instead + CustomGlobToolDesc *string + // CustomWriteFileToolDesc overrides the write_file tool description + // optional, WriteFileToolDesc by default + // Deprecated: Use WriteFileToolConfig.Desc instead + CustomWriteFileToolDesc *string + // CustomEditToolDesc overrides the edit_file tool description + // optional, EditFileToolDesc by default + // Deprecated: Use EditFileToolConfig.Desc instead + CustomEditToolDesc *string +} + +func (c *MiddlewareConfig) Validate() error { + if c == nil { + return errors.New("config should not be nil") + } + if c.Backend == nil { + return errors.New("backend should not be nil") + } + if c.StreamingShell != nil && c.Shell != nil { + return errors.New("shell and streaming shell should not be both set") + } + return nil +} + +// mergeToolConfigWithDesc merges ToolConfig with legacy Desc field +// Priority: ToolConfig.Desc > legacy Desc +// Returns an empty ToolConfig if both are nil (to allow backend default implementation) +func (c *MiddlewareConfig) mergeToolConfigWithDesc( + toolConfig *ToolConfig, + legacyDesc *string, +) *ToolConfig { + if toolConfig == nil && legacyDesc == nil { + return &ToolConfig{} + } + + if toolConfig == nil { + return &ToolConfig{ + Desc: legacyDesc, + } + } + + if toolConfig.Desc == nil && legacyDesc != nil { + merged := *toolConfig + merged.Desc = legacyDesc + return &merged + } + + return toolConfig +} + +// NewTyped constructs and returns the filesystem middleware as a TypedChatModelAgentMiddleware[M]. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +// It returns a TypedChatModelAgentMiddleware[M] which provides: +// - Better context propagation through WrapInvokableToolCall and WrapStreamableToolCall methods +// - BeforeAgent hook for modifying agent instruction and tools at runtime +// - More flexible extension points compared to the struct-based AgentMiddleware +// +// The middleware provides filesystem tools (ls, read_file, write_file, edit_file, glob, grep) +// and optionally an execute tool if the Backend implements ShellBackend or StreamingShellBackend. +func NewTyped[M adk.MessageType](ctx context.Context, config *MiddlewareConfig) (adk.TypedChatModelAgentMiddleware[M], error) { + err := config.Validate() + if err != nil { + return nil, err + } + ts, err := getFilesystemTools(ctx, config) + if err != nil { + return nil, err + } + var systemPrompt string + if config.CustomSystemPrompt != nil { + systemPrompt = *config.CustomSystemPrompt + } + + m := &typedFilesystemMiddleware[M]{ + additionalInstruction: systemPrompt, + additionalTools: ts, + } + + return m, nil +} + +// New constructs and returns the filesystem middleware as a ChatModelAgentMiddleware. +// +// This is the recommended constructor for new code. It returns a ChatModelAgentMiddleware which provides: +// - Better context propagation through WrapInvokableToolCall and WrapStreamableToolCall methods +// - BeforeAgent hook for modifying agent instruction and tools at runtime +// - More flexible extension points compared to the struct-based AgentMiddleware +// +// The middleware provides filesystem tools (ls, read_file, write_file, edit_file, glob, grep) +// and optionally an execute tool if the Backend implements ShellBackend or StreamingShellBackend. +// +// Example usage: +// +// middleware, err := filesystem.New(ctx, &filesystem.Config{ +// Backend: myBackend, +// }) +// agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ +// // ... +// Handlers: []adk.ChatModelAgentMiddleware{middleware}, +// }) +func New(ctx context.Context, config *MiddlewareConfig) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, config) +} + +type typedFilesystemMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + additionalInstruction string + additionalTools []tool.BaseTool +} + +func (m *typedFilesystemMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + if runCtx == nil { + return ctx, runCtx, nil + } + + nRunCtx := *runCtx + if m.additionalInstruction != "" { + nRunCtx.Instruction = nRunCtx.Instruction + "\n" + m.additionalInstruction + } + nRunCtx.Tools = append(nRunCtx.Tools, m.additionalTools...) + return ctx, &nRunCtx, nil +} + +// toolSpec defines a specification for creating a filesystem tool. +// It unifies the tool creation process by encapsulating the tool configuration, +// legacy descriptor, and the creation function. +type toolSpec struct { + config *ToolConfig + legacyDesc *string + createFunc func(name, desc string) (tool.BaseTool, error) +} + +func getFilesystemTools(_ context.Context, middlewareConfig *MiddlewareConfig) ([]tool.BaseTool, error) { + var tools []tool.BaseTool + + toolSpecs := []toolSpec{ + { + config: middlewareConfig.LsToolConfig, + legacyDesc: middlewareConfig.CustomLsToolDesc, + createFunc: func(name, desc string) (tool.BaseTool, error) { + if middlewareConfig.Backend != nil { + return newLsTool(middlewareConfig.Backend, name, desc) + } + return nil, nil + }, + }, + { + config: middlewareConfig.ReadFileToolConfig, + legacyDesc: middlewareConfig.CustomReadFileToolDesc, + createFunc: func(name, desc string) (tool.BaseTool, error) { + if middlewareConfig.Backend != nil { + if middlewareConfig.UseMultiModalRead { + return newMultiModalReadFileTool(middlewareConfig.Backend, name, desc) + } + return newReadFileTool(middlewareConfig.Backend, name, desc) + } + return nil, nil + }, + }, + { + config: middlewareConfig.WriteFileToolConfig, + legacyDesc: middlewareConfig.CustomWriteFileToolDesc, + createFunc: func(name, desc string) (tool.BaseTool, error) { + if middlewareConfig.Backend != nil { + return newWriteFileTool(middlewareConfig.Backend, name, desc) + } + return nil, nil + }, + }, + { + config: middlewareConfig.EditFileToolConfig, + legacyDesc: middlewareConfig.CustomEditToolDesc, + createFunc: func(name, desc string) (tool.BaseTool, error) { + if middlewareConfig.Backend != nil { + return newEditFileTool(middlewareConfig.Backend, name, desc) + } + return nil, nil + }, + }, + { + config: middlewareConfig.GlobToolConfig, + legacyDesc: middlewareConfig.CustomGlobToolDesc, + createFunc: func(name, desc string) (tool.BaseTool, error) { + if middlewareConfig.Backend != nil { + return newGlobTool(middlewareConfig.Backend, name, desc) + } + return nil, nil + }, + }, + { + config: middlewareConfig.GrepToolConfig, + legacyDesc: middlewareConfig.CustomGrepToolDesc, + createFunc: func(name, desc string) (tool.BaseTool, error) { + if middlewareConfig.Backend != nil { + return newGrepTool(middlewareConfig.Backend, name, desc) + } + return nil, nil + }, + }, + } + + for _, spec := range toolSpecs { + t, err := createToolFromSpec(middlewareConfig, spec) + if err != nil { + return nil, err + } + if t != nil { + tools = append(tools, t) + } + } + + // Create execute tool if Shell or StreamingShell is available + if middlewareConfig.StreamingShell != nil { + executeDesc, err := selectToolDesc("", ExecuteToolDesc, ExecuteToolDescChinese) + if err != nil { + return nil, err + } + + executeTool, err := newStreamingExecuteTool(middlewareConfig.StreamingShell, ToolNameExecute, executeDesc) + if err != nil { + return nil, err + } + tools = append(tools, executeTool) + } else if middlewareConfig.Shell != nil { + executeDesc, err := selectToolDesc("", ExecuteToolDesc, ExecuteToolDescChinese) + if err != nil { + return nil, err + } + + executeTool, err := newExecuteTool(middlewareConfig.Shell, ToolNameExecute, executeDesc) + if err != nil { + return nil, err + } + tools = append(tools, executeTool) + } + + return tools, nil +} + +// createToolFromSpec creates a tool instance based on the provided toolSpec. +// It handles configuration merging (ToolConfig + legacy Desc), checks if the tool +// is disabled, and prioritizes CustomTool over the default implementation. +func createToolFromSpec(middlewareConfig *MiddlewareConfig, spec toolSpec) (tool.BaseTool, error) { + mergedConfig := middlewareConfig.mergeToolConfigWithDesc(spec.config, spec.legacyDesc) + + if mergedConfig.Disable { + return nil, nil + } + + return getOrCreateTool(mergedConfig.CustomTool, func() (tool.BaseTool, error) { + desc := "" + if mergedConfig.Desc != nil { + desc = *mergedConfig.Desc + } + return spec.createFunc(mergedConfig.Name, desc) + }) +} + +func getOrCreateTool(customTool tool.BaseTool, createFunc func() (tool.BaseTool, error)) (tool.BaseTool, error) { + if customTool != nil { + return customTool, nil + } + return createFunc() +} + +type lsArgs struct { + Path string `json:"path"` +} + +func newLsTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameLs) + d, err := selectToolDesc(desc, ListFilesToolDesc, ListFilesToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input lsArgs) (string, error) { + infos, err := fs.LsInfo(ctx, &filesystem.LsInfoRequest{Path: input.Path}) + if err != nil { + return "", err + } + if len(infos) == 0 { + return noFilesFound, nil + } + paths := make([]string, 0, len(infos)) + for _, fi := range infos { + paths = append(paths, fi.Path) + } + return strings.Join(paths, "\n"), nil + }) +} + +type readFileArgs struct { + // FilePath is the path to the file to read. + FilePath string `json:"file_path" jsonschema:"description=The path to the file to read"` + + // Offset is the line number to start reading from. + Offset int `json:"offset" jsonschema:"description=The line number to start reading from. Only provide if the file is too large to read at once"` + + // Limit is the number of lines to read. + Limit int `json:"limit" jsonschema:"description=The number of lines to read. Only provide if the file is too large to read at once."` +} + +// multiModalReadFileArgs extends readFileArgs with PDF-specific parameters for MultiModalReadFileTool. +type multiModalReadFileArgs struct { + readFileArgs + + // Pages is the page range for PDF files. + Pages string `json:"pages,omitempty" jsonschema:"description=Page range for PDF files (e.g.\\, \"1-5\"\\, \"3\"\\, \"10-20\"). Only applicable to PDF files. Maximum 20 pages per request."` +} + +func newReadFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameReadFile) + d, err := selectToolDesc(desc, ReadFileToolDesc, ReadFileToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input readFileArgs) (string, error) { + if input.Offset <= 0 { + input.Offset = 1 + } + if input.Limit <= 0 { + input.Limit = 2000 + } + + fileCt, err := fs.Read(ctx, &filesystem.ReadRequest{ + FilePath: input.FilePath, + Offset: input.Offset, + Limit: input.Limit, + }) + if err != nil { + return "", err + } + if fileCt == nil { + return fmt.Sprintf("No content found at path: %s", input.FilePath), nil + } + + return formatLineNumbers(fileCt.Content, input.Offset), nil + }) +} + +// formatLineNumbers prefixes each line of content with a 1-based line number +// starting at startLine (e.g. " 1\tfoo"). startLine corresponds to the +// line number of the first line in content (usually ReadRequest.Offset). +func formatLineNumbers(content string, startLine int) string { + lines := strings.Split(content, "\n") + var b strings.Builder + for i, line := range lines { + if i < len(lines)-1 { + fmt.Fprintf(&b, "%6d\t%s\n", startLine+i, line) + } else { + fmt.Fprintf(&b, "%6d\t%s", startLine+i, line) + } + } + return b.String() +} + +const maxPagesPerRequest = 20 + +func validatePages(pages string) error { + parts := strings.SplitN(pages, "-", 2) + start, err := strconv.Atoi(parts[0]) + if err != nil || start < 1 { + return fmt.Errorf("invalid pages parameter %q: expected format like \"3\" or \"1-10\"", pages) + } + if len(parts) == 1 { + return nil + } + if parts[1] == "" { + return fmt.Errorf("invalid pages parameter %q: expected format like \"3\" or \"1-10\"", pages) + } + end, err := strconv.Atoi(parts[1]) + if err != nil || end < 1 { + return fmt.Errorf("invalid pages parameter %q: expected format like \"3\" or \"1-10\"", pages) + } + if end < start { + return fmt.Errorf("invalid pages parameter %q: end page must be >= start page", pages) + } + if end-start+1 > maxPagesPerRequest { + return fmt.Errorf("invalid pages parameter %q: range exceeds maximum of %d pages per request", pages, maxPagesPerRequest) + } + return nil +} + +func newMultiModalReadFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + er, ok := fs.(filesystem.MultiModalReader) + if !ok { + return nil, fmt.Errorf("UseMultiModalRead is enabled, but backend (type %T) does not implement filesystem.MultiModalReader interface. "+ + "Either implement the MultiModalReader interface on your backend, or set UseMultiModalRead to false", fs) + } + toolName := selectToolName(name, ToolNameReadFile) + d, err := selectToolDesc(desc, ReadFileToolDesc, ReadFileToolDescChinese) + if err != nil { + return nil, err + } + // Only append the multimodal suffix when falling back to the built-in desc. + // A custom desc is expected to describe its own capabilities, so appending + // would produce duplicated or contradictory descriptions. + if desc == "" { + d += internal.SelectPrompt(internal.I18nPrompts{ + English: EnhancedReadFileDescSuffix, + Chinese: EnhancedReadFileDescSuffixChinese, + }) + } + + return utils.InferEnhancedTool(toolName, d, func(ctx context.Context, input multiModalReadFileArgs) (*schema.ToolResult, error) { + if input.Offset <= 0 { + input.Offset = 1 + } + if input.Limit <= 0 { + input.Limit = 2000 + } + + if input.Pages != "" { + if err := validatePages(input.Pages); err != nil { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: err.Error()}}, + }, nil + } + } + + fileCt, err := er.MultiModalRead(ctx, &filesystem.MultiModalReadRequest{ + ReadRequest: filesystem.ReadRequest{ + FilePath: input.FilePath, + Offset: input.Offset, + Limit: input.Limit, + }, + Pages: input.Pages, + }) + if err != nil { + return nil, err + } + + if fileCt == nil { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: fmt.Sprintf("No content found at path: %s", input.FilePath)}}, + }, nil + } + + // Multimodal result: convert FileContentPart to ToolOutputPart + if len(fileCt.Parts) > 0 { + parts := make([]schema.ToolOutputPart, 0, len(fileCt.Parts)) + enc := base64Encoder{} + for _, p := range fileCt.Parts { + if len(p.Data) == 0 { + return nil, fmt.Errorf("FileContentPart.Data is empty for type %s", p.Type) + } + if p.MIMEType == "" { + return nil, fmt.Errorf("FileContentPart.MIMEType is empty for type %s", p.Type) + } + b64 := enc.encode(p.Data) + switch p.Type { + case filesystem.FileContentPartTypeImage: + parts = append(parts, schema.ToolOutputPart{ + Type: schema.ToolPartTypeImage, + Image: &schema.ToolOutputImage{ + MessagePartCommon: schema.MessagePartCommon{ + MIMEType: p.MIMEType, + Base64Data: &b64, + }, + }, + }) + case filesystem.FileContentPartTypePDF: + parts = append(parts, schema.ToolOutputPart{ + Type: schema.ToolPartTypeFile, + File: &schema.ToolOutputFile{ + MessagePartCommon: schema.MessagePartCommon{ + MIMEType: p.MIMEType, + Base64Data: &b64, + }, + }, + }) + default: + // FileContentPartType is defined by Backend implementations. + // Unrecognized types are unlikely but should fail explicitly rather than silently. + return nil, fmt.Errorf("unsupported FileContentPartType: %s", p.Type) + } + } + return &schema.ToolResult{Parts: parts}, nil + } + + if fileCt.FileContent == nil { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: fmt.Sprintf("No content found at path: %s", input.FilePath)}}, + }, nil + } + + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: formatLineNumbers(fileCt.Content, input.Offset)}}, + }, nil + }) +} + +type writeFileArgs struct { + // FilePath is the path to the file to write. + FilePath string `json:"file_path" jsonschema:"description=The path to the file to write"` + + // Content is the content to write to the file. + Content string `json:"content" jsonschema:"description=The content to write to the file"` +} + +func newWriteFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameWriteFile) + d, err := selectToolDesc(desc, WriteFileToolDesc, WriteFileToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input writeFileArgs) (string, error) { + err := fs.Write(ctx, &filesystem.WriteRequest{ + FilePath: input.FilePath, + Content: input.Content, + }) + if err != nil { + return "", err + } + return fmt.Sprintf("Updated file %s", input.FilePath), nil + }) +} + +type editFileArgs struct { + // FilePath is the path to the file to modify. + FilePath string `json:"file_path" jsonschema:"description=The path to the file to modify"` + + // OldString is the text to replace. + OldString string `json:"old_string" jsonschema:"description=The text to replace"` + + // NewString is the text to replace it with. + NewString string `json:"new_string" jsonschema:"description=The text to replace it with (must be different from old_string)"` + + // ReplaceAll indicates whether to replace all occurrences of old_string. + ReplaceAll bool `json:"replace_all" jsonschema:"description=Replace all occurrences of old_string (default false),default=false"` +} + +func newEditFileTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameEditFile) + d, err := selectToolDesc(desc, EditFileToolDesc, EditFileToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input editFileArgs) (string, error) { + err := fs.Edit(ctx, &filesystem.EditRequest{ + FilePath: input.FilePath, + OldString: input.OldString, + NewString: input.NewString, + ReplaceAll: input.ReplaceAll, + }) + if err != nil { + return "", err + } + return fmt.Sprintf("Successfully replaced the string in '%s'", input.FilePath), nil + }) +} + +type globArgs struct { + // Pattern is the glob pattern to match files against. + Pattern string `json:"pattern" jsonschema:"description=The glob pattern to match files against"` + + // Path is the directory to search in. + Path string `json:"path" jsonschema:"description=The directory to search in. If not specified\\, the current working directory will be used. IMPORTANT: Omit this field to use the default directory. DO NOT enter 'undefined' or 'null' - simply omit it for the default behavior. Must be a valid directory path if provided."` +} + +func newGlobTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameGlob) + d, err := selectToolDesc(desc, GlobToolDesc, GlobToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input globArgs) (string, error) { + infos, err := fs.GlobInfo(ctx, &filesystem.GlobInfoRequest{ + Pattern: input.Pattern, + Path: input.Path, + }) + if err != nil { + return "", err + } + if len(infos) == 0 { + return noFilesFound, nil + } + paths := make([]string, 0, len(infos)) + for _, fi := range infos { + paths = append(paths, fi.Path) + } + return strings.Join(paths, "\n"), nil + }) +} + +type grepArgs struct { + // Pattern is the regular expression pattern to search for in file contents. + Pattern string `json:"pattern" jsonschema:"description=The regular expression pattern to search for in file contents"` + + // Path is the file or directory to search in. Defaults to current working directory. + Path *string `json:"path,omitempty" jsonschema:"description=File or directory to search in (rg PATH). Defaults to current working directory."` + + // Glob is the glob pattern to filter files (e.g. "*.js", "*.{ts,tsx}"). + Glob *string `json:"glob,omitempty" jsonschema:"description=Glob pattern to filter files (e.g. '*.js'\\, '*.{ts\\,tsx}') - maps to rg --glob"` + + // OutputMode specifies the output format. + // "content" shows matching lines (supports context, line numbers, head_limit). + // "files_with_matches" shows file paths (supports head_limit). + // "count" shows match counts (supports head_limit). + // Defaults to "files_with_matches". + OutputMode string `json:"output_mode,omitempty" jsonschema:"description=Output mode: 'content' shows matching lines (supports -A/-B/-C context\\, -n line numbers\\, head_limit)\\, 'files_with_matches' shows file paths (supports head_limit)\\, 'count' shows match counts (supports head_limit). Defaults to 'files_with_matches'.,enum=content,enum=files_with_matches,enum=count"` + + // Context is the number of lines to show before and after each match. + // Only applicable when output_mode is "content". + Context *int `json:"-C,omitempty" jsonschema:"description=Number of lines to show before and after each match (rg -C). Requires output_mode: 'content'\\, ignored otherwise."` + + // BeforeLines is the number of lines to show before each match. + // Only applicable when output_mode is "content". + BeforeLines *int `json:"-B,omitempty" jsonschema:"description=Number of lines to show before each match (rg -B). Requires output_mode: 'content'\\, ignored otherwise."` + + // AfterLines is the number of lines to show after each match. + // Only applicable when output_mode is "content". + AfterLines *int `json:"-A,omitempty" jsonschema:"description=Number of lines to show after each match (rg -A). Requires output_mode: 'content'\\, ignored otherwise."` + + // ShowLineNumbers enables showing line numbers in output. + // Only applicable when output_mode is "content". Defaults to true. + ShowLineNumbers *bool `json:"-n,omitempty" jsonschema:"description=Show line numbers in output (rg -n). Requires output_mode: 'content'\\, ignored otherwise. Defaults to true."` + + // CaseInsensitive enables case insensitive search. + CaseInsensitive *bool `json:"-i,omitempty" jsonschema:"description=Case insensitive search (rg -i)"` + + // FileType is the file type to search (e.g., js, py, rust, go, java). + // More efficient than Glob for standard file types. + FileType *string `json:"type,omitempty" jsonschema:"description=File type to search (rg --type). Common types: js\\, py\\, rust\\, go\\, java\\, etc. More efficient than include for standard file types."` + + // HeadLimit limits output to first N lines/entries. + // Works across all output modes. Defaults to 0 (unlimited). + HeadLimit *int `json:"head_limit,omitempty" jsonschema:"description=Limit output to first N lines/entries\\, equivalent to '| head -N'. Works across all output modes: content (limits output lines)\\, files_with_matches (limits file paths)\\, count (limits count entries). Defaults to 0 (unlimited)."` + + // Offset skips first N lines/entries before applying HeadLimit. + // Works across all output modes. Defaults to 0. + Offset *int `json:"offset,omitempty" jsonschema:"description=Skip first N lines/entries before applying head_limit\\, equivalent to '| tail -n +N | head -N'. Works across all output modes. Defaults to 0."` + + // Multiline enables multiline mode where patterns can span lines. + // - true: Allows patterns to match across lines, "." matches newlines + // - false: Default, matches only within single lines + Multiline *bool `json:"multiline,omitempty" jsonschema:"description=Enable multiline mode where . matches newlines and patterns can span lines (rg -U --multiline-dotall). Default: false."` +} + +func newGrepTool(fs filesystem.Backend, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameGrep) + d, err := selectToolDesc(desc, GrepToolDesc, GrepToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input grepArgs) (string, error) { + // Extract string parameters + path := valueOrDefault(input.Path, "") + glob := valueOrDefault(input.Glob, "") + fileType := valueOrDefault(input.FileType, "") + var beforeLines, afterLines int + + if input.Context != nil { + beforeLines = valueOrDefault(input.Context, 0) + afterLines = valueOrDefault(input.Context, 0) + } else { + // Extract context parameters + beforeLines = valueOrDefault(input.BeforeLines, 0) + afterLines = valueOrDefault(input.AfterLines, 0) + } + + // Extract boolean flags + caseInsensitive := valueOrDefault(input.CaseInsensitive, false) + enableMultiline := valueOrDefault(input.Multiline, false) + + // Extract pagination parameters + headLimit := valueOrDefault(input.HeadLimit, 0) + offset := valueOrDefault(input.Offset, 0) + + matches, err := fs.GrepRaw(ctx, &filesystem.GrepRequest{ + Pattern: input.Pattern, + Path: path, + Glob: glob, + FileType: fileType, + CaseInsensitive: caseInsensitive, + AfterLines: afterLines, + BeforeLines: beforeLines, + EnableMultiline: enableMultiline, + }) + if err != nil { + return "", err + } + + sort.SliceStable(matches, func(i, j int) bool { + return filepath.Base(matches[i].Path) < filepath.Base(matches[j].Path) + }) + + switch input.OutputMode { + case "content": + matches = applyPagination(matches, offset, headLimit) + return formatContentMatches(matches, valueOrDefault(input.ShowLineNumbers, true)), nil + + case "count": + return formatCountMatches(matches, offset, headLimit), nil + + case "files_with_matches": + return formatFileMatches(matches, offset, headLimit), nil + + default: + return formatFileMatches(matches, offset, headLimit), nil + } + }) +} + +type executeArgs struct { + Command string `json:"command"` +} + +func newExecuteTool(sb filesystem.Shell, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameExecute) + d, err := selectToolDesc(desc, ExecuteToolDesc, ExecuteToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferTool(toolName, d, func(ctx context.Context, input executeArgs) (string, error) { + result, err := sb.Execute(ctx, &filesystem.ExecuteRequest{ + Command: input.Command, + }) + if err != nil { + return "", err + } + + return convExecuteResponse(result), nil + }) +} + +func newStreamingExecuteTool(sb filesystem.StreamingShell, name string, desc string) (tool.BaseTool, error) { + toolName := selectToolName(name, ToolNameExecute) + d, err := selectToolDesc(desc, ExecuteToolDesc, ExecuteToolDescChinese) + if err != nil { + return nil, err + } + return utils.InferStreamTool(toolName, d, func(ctx context.Context, input executeArgs) (*schema.StreamReader[string], error) { + result, err := sb.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{ + Command: input.Command, + }) + if err != nil { + return nil, err + } + sr, sw := schema.Pipe[string](10) + go func() { + defer func() { + e := recover() + if e != nil { + sw.Send("", fmt.Errorf("panic: %v,\n stack: %s", e, string(debug.Stack()))) + } + sw.Close() + }() + + var hasSentContent bool + var exitCode *int + + for { + chunk, recvErr := result.Recv() + if recvErr == io.EOF { + break + } + if recvErr != nil { + sw.Send("", recvErr) + return + } + + if chunk == nil { + continue + } + if chunk.ExitCode != nil { + exitCode = chunk.ExitCode + } + + parts := make([]string, 0, 2) + if chunk.Output != "" { + parts = append(parts, chunk.Output) + } + if chunk.Truncated { + parts = append(parts, "[Output was truncated due to size limits]") + } + if len(parts) > 0 { + sw.Send(strings.Join(parts, "\n"), nil) + hasSentContent = true + } + } + + if exitCode != nil && *exitCode != 0 { + sw.Send(fmt.Sprintf("\n[Command failed with exit code %d]", *exitCode), nil) + } else if !hasSentContent { + sw.Send("[Command executed successfully with no output]", nil) + } + }() + + return sr, nil + }) +} + +func convExecuteResponse(response *filesystem.ExecuteResponse) string { + if response == nil { + return "" + } + parts := []string{response.Output} + if response.ExitCode != nil && *response.ExitCode != 0 { + parts = append(parts, fmt.Sprintf("[Command failed with exit code %d]", *response.ExitCode)) + } + if response.Truncated { + parts = append(parts, "[Output was truncated due to size limits]") + } + + result := strings.Join(parts, "\n") + if result == "" && (response.ExitCode == nil || *response.ExitCode == 0) { + return "[Command executed successfully with no output]" + } + return result +} + +// valueOrDefault returns the value pointed to by ptr, or defaultValue if ptr is nil. +func valueOrDefault[T any](ptr *T, defaultValue T) T { + if ptr != nil { + return *ptr + } + return defaultValue +} + +// base64Encoder reuses a buffer across multiple base64 encoding calls to reduce allocations. +type base64Encoder struct { + buf []byte +} + +func (e *base64Encoder) encode(data []byte) string { + n := base64.StdEncoding.EncodedLen(len(data)) + if cap(e.buf) < n { + e.buf = make([]byte, n) + } else { + e.buf = e.buf[:n] + } + base64.StdEncoding.Encode(e.buf, data) + return string(e.buf) +} + +func applyPagination[T any](items []T, offset, headLimit int) []T { + if offset < 0 { + offset = 0 + } + if offset >= len(items) { + return []T{} + } + items = items[offset:] + + if headLimit > 0 && headLimit < len(items) { + items = items[:headLimit] + } + return items +} + +func formatFileMatches(matches []filesystem.GrepMatch, offset, headLimit int) string { + if len(matches) == 0 { + return noFilesFound + } + seen := make(map[string]bool) + var uniquePaths []string + for _, match := range matches { + if !seen[match.Path] { + seen[match.Path] = true + uniquePaths = append(uniquePaths, match.Path) + } + } + totalFiles := len(uniquePaths) + uniquePaths = applyPagination(uniquePaths, offset, headLimit) + + fileWord := "files" + if totalFiles == 1 { + fileWord = "file" + } + return fmt.Sprintf("Found %d %s\n%s", totalFiles, fileWord, strings.Join(uniquePaths, "\n")) +} + +func formatContentMatches(matches []filesystem.GrepMatch, showLineNum bool) string { + if len(matches) == 0 { + return noMatchesFound + } + var b strings.Builder + for _, match := range matches { + b.WriteString(match.Path) + if showLineNum { + b.WriteString(":") + b.WriteString(strconv.Itoa(match.Line)) + } + b.WriteString(":") + b.WriteString(match.Content) + b.WriteString("\n") + } + return strings.TrimSuffix(b.String(), "\n") +} + +func formatCountMatches(matches []filesystem.GrepMatch, offset, headLimit int) string { + countMap := make(map[string]int) + for _, match := range matches { + countMap[match.Path]++ + } + + var paths []string + for path := range countMap { + paths = append(paths, path) + } + sort.Strings(paths) + + totalOccurrences := len(matches) + totalFiles := len(paths) + + occurrenceWord := "occurrences" + if totalOccurrences == 1 { + occurrenceWord = "occurrence" + } + fileWord := "files" + if totalFiles == 1 { + fileWord = "file" + } + + if totalOccurrences == 0 { + return fmt.Sprintf("%s\n\nFound %d total %s across %d %s.", noMatchesFound, totalOccurrences, occurrenceWord, totalFiles, fileWord) + } + + paths = applyPagination(paths, offset, headLimit) + + var b strings.Builder + for _, path := range paths { + b.WriteString(path) + b.WriteString(":") + b.WriteString(strconv.Itoa(countMap[path])) + b.WriteString("\n") + } + result := strings.TrimSuffix(b.String(), "\n") + return fmt.Sprintf("%s\n\nFound %d total %s across %d %s.", result, totalOccurrences, occurrenceWord, totalFiles, fileWord) +} + +// selectToolDesc returns the custom description if provided, otherwise selects the appropriate +// i18n description based on the current language setting. +func selectToolDesc(customDesc string, defaultEnglish, defaultChinese string) (string, error) { + if customDesc != "" { + return customDesc, nil + } + return internal.SelectPrompt(internal.I18nPrompts{ + English: defaultEnglish, + Chinese: defaultChinese, + }), nil +} + +// selectToolName returns the custom tool name if provided, otherwise returns the default name. +func selectToolName(customName string, defaultName string) string { + if customName != "" { + return customName + } + return defaultName +} diff --git a/adk/middlewares/filesystem/filesystem_test.go b/adk/middlewares/filesystem/filesystem_test.go new file mode 100644 index 0000000..cb59353 --- /dev/null +++ b/adk/middlewares/filesystem/filesystem_test.go @@ -0,0 +1,2647 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// setupTestBackend creates a test backend with some initial files +func setupTestBackend() *filesystem.InMemoryBackend { + backend := filesystem.NewInMemoryBackend() + ctx := context.Background() + + // Create test files + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/file1.txt", + Content: "line1\nline2\nline3\nline4\nline5", + }) + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/file2.go", + Content: "package main\n\nfunc main() {\n\tprintln(\"hello\")\n}", + }) + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/dir1/file3.txt", + Content: "hello world\nfoo bar\nhello again", + }) + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/dir1/file4.py", + Content: "print('hello')\nprint('world')", + }) + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/dir2/file5.go", + Content: "package test\n\nfunc test() {}", + }) + + return backend +} + +// invokeTool is a helper to invoke a tool with JSON input +func invokeTool(_ *testing.T, bt tool.BaseTool, input string) (string, error) { + ctx := context.Background() + result, err := bt.(tool.InvokableTool).InvokableRun(ctx, input) + if err != nil { + return "", err + } + return result, nil +} + +func TestLsTool(t *testing.T) { + backend := setupTestBackend() + lsTool, err := newLsTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create ls tool: %v", err) + } + + tests := []struct { + name string + input string + expected []string // expected paths in output + }{ + { + name: "list root", + input: `{"path": "/"}`, + expected: []string{"file1.txt", "file2.go", "dir1", "dir2"}, + }, + { + name: "list empty path (defaults to root)", + input: `{"path": ""}`, + expected: []string{"file1.txt", "file2.go", "dir1", "dir2"}, + }, + { + name: "list dir1", + input: `{"path": "/dir1"}`, + expected: []string{"file3.txt", "file4.py"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := invokeTool(t, lsTool, tt.input) + if err != nil { + t.Fatalf("ls tool failed: %v", err) + } + + for _, expectedPath := range tt.expected { + if !strings.Contains(result, expectedPath) { + t.Errorf("Expected output to contain %q, got: %s", expectedPath, result) + } + } + }) + } +} + +func TestReadFileTool(t *testing.T) { + backend := setupTestBackend() + readTool, err := newReadFileTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create read_file tool: %v", err) + } + + tests := []struct { + name string + input string + expected string + shouldError bool + }{ + { + name: "read full file", + input: `{"file_path": "/file1.txt", "offset": 0, "limit": 100}`, + expected: " 1\tline1\n 2\tline2\n 3\tline3\n 4\tline4\n 5\tline5", + }, + { + name: "read with offset", + input: `{"file_path": "/file1.txt", "offset": 2, "limit": 2}`, + expected: " 2\tline2\n 3\tline3", + }, + { + name: "read with default limit", + input: `{"file_path": "/file1.txt", "offset": 0, "limit": 0}`, + expected: " 1\tline1\n 2\tline2\n 3\tline3\n 4\tline4\n 5\tline5", + }, + { + name: "read with negative offset (treated as 0)", + input: `{"file_path": "/file1.txt", "offset": -1, "limit": 2}`, + expected: " 1\tline1\n 2\tline2", + }, + { + name: "read non-existent file", + input: `{"file_path": "/nonexistent.txt", "offset": 0, "limit": 10}`, + shouldError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := invokeTool(t, readTool, tt.input) + if tt.shouldError { + if err == nil { + t.Error("Expected error but got none") + } + return + } + if err != nil { + t.Fatalf("read_file tool failed: %v", err) + } + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestReadFileTool_DefaultLimit(t *testing.T) { + // Build a file with more than 2000 lines to verify the tool layer applies the default limit + backend := filesystem.NewInMemoryBackend() + var b strings.Builder + totalLines := 2500 + for i := 1; i <= totalLines; i++ { + if i > 1 { + b.WriteString("\n") + } + fmt.Fprintf(&b, "line%d", i) + } + backend.Write(context.Background(), &filesystem.WriteRequest{ + FilePath: "/big.txt", + Content: b.String(), + }) + + readTool, err := newReadFileTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create read_file tool: %v", err) + } + + t.Run("limit=0 defaults to 2000 lines in tool layer", func(t *testing.T) { + result, err := invokeTool(t, readTool, `{"file_path": "/big.txt", "offset": 0, "limit": 0}`) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + lines := strings.Split(result, "\n") + if len(lines) != 2000 { + t.Errorf("expected 2000 lines, got %d", len(lines)) + } + }) + + t.Run("explicit limit is respected", func(t *testing.T) { + result, err := invokeTool(t, readTool, `{"file_path": "/big.txt", "offset": 0, "limit": 10}`) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + lines := strings.Split(result, "\n") + if len(lines) != 10 { + t.Errorf("expected 10 lines, got %d", len(lines)) + } + }) + + t.Run("backend read with limit=0 returns all lines", func(t *testing.T) { + content, err := backend.Read(context.Background(), &filesystem.ReadRequest{ + FilePath: "/big.txt", + Limit: 0, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + lines := strings.Split(content.Content, "\n") + if len(lines) != totalLines { + t.Errorf("expected %d lines from backend, got %d", totalLines, len(lines)) + } + }) +} + +func TestWriteFileTool(t *testing.T) { + backend := setupTestBackend() + writeTool, err := newWriteFileTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create write_file tool: %v", err) + } + + tests := []struct { + name string + input string + expected string + isError bool + }{ + { + name: "write new file", + input: `{"file_path": "/newfile.txt", "content": "new content"}`, + expected: "Updated file /newfile.txt", + }, + { + name: "overwrite existing file", + input: `{"file_path": "/file1.txt", "content": "overwritten"}`, + isError: false, + expected: "Updated file /file1.txt", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := invokeTool(t, writeTool, tt.input) + if tt.isError { + if err == nil { + t.Errorf("Expected an error, but got none") + } + return + } + if err != nil { + t.Fatalf("write_file tool failed: %v", err) + } + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } + + // Verify the file was actually written + ctx := context.Background() + content, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/newfile.txt", + Offset: 0, + Limit: 100, + }) + if err != nil { + t.Fatalf("Failed to read written file: %v", err) + } + if content.Content != "new content" { + t.Errorf("Expected written content to be 'new content', got %q", content.Content) + } +} + +func TestEditFileTool(t *testing.T) { + backend := setupTestBackend() + editTool, err := newEditFileTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create edit_file tool: %v", err) + } + + tests := []struct { + name string + setupFile string + setupContent string + input string + expected string + shouldError bool + }{ + { + name: "replace first occurrence", + setupFile: "/edit1.txt", + setupContent: "hello world\nhello again\nhello world", + input: `{"file_path": "/edit1.txt", "old_string": "hello again", "new_string": "hi", "replace_all": false}`, + expected: "hello world\nhi\nhello world", + }, + { + name: "replace all occurrences", + setupFile: "/edit2.txt", + setupContent: "hello world\nhello again\nhello world", + input: `{"file_path": "/edit2.txt", "old_string": "hello", "new_string": "hi", "replace_all": true}`, + expected: "hi world\nhi again\nhi world", + }, + { + name: "non-existent file", + setupFile: "", + setupContent: "", + input: `{"file_path": "/nonexistent.txt", "old_string": "old", "new_string": "new", "replace_all": false}`, + shouldError: true, + }, + { + name: "empty old_string", + setupFile: "/edit3.txt", + setupContent: "content", + input: `{"file_path": "/edit3.txt", "old_string": "", "new_string": "new", "replace_all": false}`, + shouldError: true, + }, + } + + ctx := context.Background() + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Setup file if needed + if tt.setupFile != "" { + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: tt.setupFile, + Content: tt.setupContent, + }) + } + + _, err := invokeTool(t, editTool, tt.input) + if tt.shouldError { + if err == nil { + t.Error("Expected error but got none") + } + return + } + if err != nil { + t.Fatalf("edit_file tool failed: %v", err) + } + result, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: tt.setupFile, + Offset: 0, + Limit: 0, + }) + if err != nil { + t.Fatalf("edit_file tool failed: %v", err) + } + if result.Content != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result.Content) + } + }) + } +} + +func TestGlobTool(t *testing.T) { + backend := setupTestBackend() + globTool, err := newGlobTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create glob tool: %v", err) + } + + tests := []struct { + name string + input string + expected []string + }{ + { + name: "match all .txt files in root", + input: `{"pattern": "*.txt", "path": "/"}`, + expected: []string{"file1.txt"}, + }, + { + name: "match all .go files in root", + input: `{"pattern": "*.go", "path": "/"}`, + expected: []string{"file2.go"}, + }, + { + name: "match all .txt files in dir1", + input: `{"pattern": "*.txt", "path": "/dir1"}`, + expected: []string{"file3.txt"}, + }, + { + name: "match all .py files in dir1", + input: `{"pattern": "*.py", "path": "/dir1"}`, + expected: []string{"file4.py"}, + }, + + { + name: "empty path defaults to root", + input: `{"pattern": "*.go", "path": ""}`, + expected: []string{"file2.go"}, + }, + + { + name: "match all .txt files in dir1 in root dir", + input: `{"pattern": "/dir1/*.txt", "path": "/"}`, + expected: []string{"/dir1/file3.txt"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := invokeTool(t, globTool, tt.input) + if err != nil { + t.Fatalf("glob tool failed: %v", err) + } + + for _, expectedPath := range tt.expected { + if !strings.Contains(result, expectedPath) { + t.Errorf("Expected output to contain %q, got: %s", expectedPath, result) + } + } + }) + } +} + +func TestGrepTool(t *testing.T) { + backend := setupTestBackend() + grepTool, err := newGrepTool(backend, "", "") + if err != nil { + t.Fatalf("Failed to create grep tool: %v", err) + } + + tests := []struct { + name string + input string + expected string + contains []string + }{ + { + name: "grep with count mode", + input: `{"pattern": "hello", "output_mode": "count"}`, + expected: "/dir1/file3.txt:2\n/dir1/file4.py:1\n/file2.go:1\n\nFound 4 total occurrences across 3 files.", // 2 in file3.txt, 1 in file4.py, 1 in file2.go + }, + { + name: "grep with content mode", + input: `{"pattern": "hello", "output_mode": "content"}`, + contains: []string{"/dir1/file3.txt:1:hello world", "/dir1/file3.txt:3:hello again", "/dir1/file4.py:1:print('hello')"}, + }, + { + name: "grep with files_with_matches mode (default)", + input: `{"pattern": "hello", "output_mode": "files_with_matches"}`, + contains: []string{"/dir1/file3.txt", "/dir1/file4.py"}, + }, + { + name: "grep with glob filter", + input: `{"pattern": "hello", "glob": "*.txt", "output_mode": "count"}`, + expected: "/dir1/file3.txt:2\n\nFound 2 total occurrences across 1 file.", // only in file3.txt + }, + { + name: "grep withpath filter", + input: `{"pattern": "package", "path": "/dir2", "output_mode": "count"}`, + expected: "/dir2/file5.go:1\n\nFound 1 total occurrence across 1 file.", // only in dir2/file5.go + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := invokeTool(t, grepTool, tt.input) + if err != nil { + t.Fatalf("grep tool failed: %v", err) + } + + if tt.expected != "" { + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + } + + for _, expectedStr := range tt.contains { + if !strings.Contains(result, expectedStr) { + t.Errorf("Expected output to contain %q, got: %s", expectedStr, result) + } + } + }) + } +} + +func TestExecuteTool(t *testing.T) { + backend := setupTestBackend() + + tests := []struct { + name string + resp *filesystem.ExecuteResponse + input string + expected string + shouldError bool + }{ + { + name: "successful command execution", + resp: &filesystem.ExecuteResponse{ + Output: "hello world", + ExitCode: ptrOf(0), + }, + input: `{"command": "echo hello world"}`, + expected: "hello world", + }, + { + name: "command with non-zero exit code", + resp: &filesystem.ExecuteResponse{ + Output: "error: file not found", + ExitCode: ptrOf(1), + }, + input: `{"command": "cat nonexistent.txt"}`, + expected: "error: file not found\n[Command failed with exit code 1]", + }, + { + name: "command with truncated output", + resp: &filesystem.ExecuteResponse{ + Output: "partial output...", + ExitCode: ptrOf(0), + Truncated: true, + }, + input: `{"command": "cat largefile.txt"}`, + expected: "partial output...\n[Output was truncated due to size limits]", + }, + { + name: "command with both non-zero exit code and truncated output", + resp: &filesystem.ExecuteResponse{ + Output: "error output...", + ExitCode: ptrOf(2), + Truncated: true, + }, + input: `{"command": "failing command"}`, + expected: "error output...\n[Command failed with exit code 2]\n[Output was truncated due to size limits]", + }, + { + name: "successful command with no output", + resp: &filesystem.ExecuteResponse{ + Output: "", + ExitCode: ptrOf(0), + }, + input: `{"command": "mkdir /tmp/test"}`, + expected: "[Command executed successfully with no output]", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + executeTool, err := newExecuteTool(&mockShellBackend{ + Backend: backend, + resp: tt.resp, + }, "", "") + assert.NoError(t, err) + + result, err := invokeTool(t, executeTool, tt.input) + if tt.shouldError { + assert.Error(t, err) + return + } + assert.NoError(t, err) + assert.Equal(t, tt.expected, result) + }) + } +} + +func ptrOf[T any](t T) *T { + return &t +} + +type mockShellBackend struct { + filesystem.Backend + resp *filesystem.ExecuteResponse +} + +func (m *mockShellBackend) Execute(ctx context.Context, req *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) { + return m.resp, nil +} + +func TestGetFilesystemTools(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("returns 6 tools for regular Backend", func(t *testing.T) { + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{Backend: backend}) + assert.NoError(t, err) + assert.Len(t, tools, 6) + + // Verify tool names + toolNames := make([]string, 0, len(tools)) + for _, to := range tools { + info, _ := to.Info(ctx) + toolNames = append(toolNames, info.Name) + } + assert.Contains(t, toolNames, "ls") + assert.Contains(t, toolNames, "read_file") + assert.Contains(t, toolNames, "write_file") + assert.Contains(t, toolNames, "edit_file") + assert.Contains(t, toolNames, "glob") + assert.Contains(t, toolNames, "grep") + }) + + t.Run("returns 7 tools for Shell", func(t *testing.T) { + shellBackend := &mockShellBackend{ + Backend: backend, + resp: &filesystem.ExecuteResponse{Output: "ok"}, + } + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{Backend: shellBackend, Shell: shellBackend}) + assert.NoError(t, err) + assert.Len(t, tools, 7) + + // Verify execute tool is included + toolNames := make([]string, 0, len(tools)) + for _, to := range tools { + info, _ := to.Info(ctx) + toolNames = append(toolNames, info.Name) + } + assert.Contains(t, toolNames, "execute") + }) + + t.Run("custom tool descriptions", func(t *testing.T) { + customLsDesc := "Custom ls description" + customReadDesc := "Custom read description" + + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Backend: backend, + CustomLsToolDesc: &customLsDesc, + CustomReadFileToolDesc: &customReadDesc, + }) + assert.NoError(t, err) + assert.Len(t, tools, 6) + + // Verify custom descriptions are applied + for _, to := range tools { + info, _ := to.Info(ctx) + if info.Name == "ls" { + assert.Equal(t, customLsDesc, info.Desc) + } + if info.Name == "read_file" { + assert.Equal(t, customReadDesc, info.Desc) + } + } + }) +} + +func TestNew(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("nil config returns error", func(t *testing.T) { + _, err := New(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config should not be nil") + }) + + t.Run("nil backend returns error", func(t *testing.T) { + _, err := New(ctx, &MiddlewareConfig{Backend: nil}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend should not be nil") + }) + + t.Run("valid config with default settings", func(t *testing.T) { + m, err := New(ctx, &MiddlewareConfig{Backend: backend}) + assert.NoError(t, err) + assert.NotNil(t, m) + + fm, ok := m.(*typedFilesystemMiddleware[*schema.Message]) + assert.True(t, ok) + assert.Len(t, fm.additionalTools, 6) + }) + + t.Run("custom system prompt", func(t *testing.T) { + customPrompt := "Custom system prompt" + m, err := New(ctx, &MiddlewareConfig{ + Backend: backend, + CustomSystemPrompt: &customPrompt, + }) + assert.NoError(t, err) + + fm, ok := m.(*typedFilesystemMiddleware[*schema.Message]) + assert.True(t, ok) + assert.Equal(t, customPrompt, fm.additionalInstruction) + }) + + t.Run("ShellBackend adds execute tool", func(t *testing.T) { + shellBackend := &mockShellBackend{ + Backend: backend, + resp: &filesystem.ExecuteResponse{Output: "ok"}, + } + m, err := New(ctx, &MiddlewareConfig{Backend: shellBackend, Shell: shellBackend}) + assert.NoError(t, err) + + fm, ok := m.(*typedFilesystemMiddleware[*schema.Message]) + assert.True(t, ok) + assert.Len(t, fm.additionalTools, 7) + }) +} + +func TestFilesystemMiddleware_BeforeAgent(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("adds instruction and tools to context", func(t *testing.T) { + m, err := New(ctx, &MiddlewareConfig{Backend: backend}) + assert.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext{ + Instruction: "Original instruction", + Tools: nil, + } + + newCtx, newRunCtx, err := m.BeforeAgent(ctx, runCtx) + assert.NoError(t, err) + assert.NotNil(t, newCtx) + assert.NotNil(t, newRunCtx) + assert.Contains(t, newRunCtx.Instruction, "Original instruction") + assert.Len(t, newRunCtx.Tools, 6) + }) + + t.Run("nil runCtx returns nil", func(t *testing.T) { + m, err := New(ctx, &MiddlewareConfig{Backend: backend}) + assert.NoError(t, err) + + newCtx, newRunCtx, err := m.BeforeAgent(ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, newCtx) + assert.Nil(t, newRunCtx) + }) +} + +func TestFilesystemMiddleware_WrapInvokableToolCall(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("small result passes through unchanged", func(t *testing.T) { + m, err := New(ctx, &MiddlewareConfig{Backend: backend}) + assert.NoError(t, err) + + endpoint := func(ctx context.Context, args string, opts ...tool.Option) (string, error) { + return "small result", nil + } + + tCtx := &adk.ToolContext{Name: "test_tool", CallID: "call-1"} + wrapped, err := m.WrapInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + + result, err := wrapped(ctx, "{}") + assert.NoError(t, err) + assert.Equal(t, "small result", result) + }) + +} + +func TestGrepToolWithSortingAndPagination(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + ctx := context.Background() + + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/zebra.txt", + Content: "match1\nmatch2\nmatch3", + }) + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/apple.txt", + Content: "match4\nmatch5", + }) + backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/banana.txt", + Content: "match6\nmatch7\nmatch8", + }) + + grepTool, err := newGrepTool(backend, "", "") + assert.NoError(t, err) + + t.Run("files sorted by basename", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "files_with_matches"}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 4, len(lines)) // 1 summary + 3 files + assert.Contains(t, lines[0], "Found 3 files") + assert.Contains(t, lines[1], "apple.txt") + assert.Contains(t, lines[2], "banana.txt") + assert.Contains(t, lines[3], "zebra.txt") + }) + + t.Run("files_with_matches with offset", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "files_with_matches", "offset": 1}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 3, len(lines)) // 1 summary + 2 files (pagination applied) + assert.Contains(t, lines[0], "Found 3 files") // total count before pagination + assert.Contains(t, lines[1], "banana.txt") + assert.Contains(t, lines[2], "zebra.txt") + }) + + t.Run("files_with_matches with head_limit", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "files_with_matches", "head_limit": 2}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 3, len(lines)) // 1 summary + 2 files (pagination applied) + assert.Contains(t, lines[0], "Found 3 files") // total count before pagination + assert.Contains(t, lines[1], "apple.txt") + assert.Contains(t, lines[2], "banana.txt") + }) + + t.Run("files_with_matches with offset and head_limit", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "files_with_matches", "offset": 1, "head_limit": 1}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 2, len(lines)) // 1 summary + 1 file (pagination applied) + assert.Contains(t, lines[0], "Found 3 files") // total count before pagination + assert.Contains(t, lines[1], "banana.txt") + }) + + t.Run("content mode sorted and paginated", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "content", "head_limit": 3}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 3, len(lines)) + assert.Contains(t, lines[0], "apple.txt") + }) + + t.Run("content mode with offset", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "content", "offset": 2, "head_limit": 2}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 2, len(lines)) + }) + + t.Run("count mode sorted", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "count"}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 5, len(lines)) // 3 file counts + 1 empty line + 1 summary line + assert.Contains(t, lines[0], "apple.txt:2") + assert.Contains(t, lines[1], "banana.txt:3") + assert.Contains(t, lines[2], "zebra.txt:3") + assert.Contains(t, lines[4], "Found 8 total occurrences across 3 files.") + }) + + t.Run("count mode with pagination", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "count", "offset": 1, "head_limit": 1}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 3, len(lines)) // 1 file count + 1 empty line + 1 summary line + assert.Contains(t, lines[0], "banana.txt:3") + assert.Contains(t, lines[2], "Found 8 total occurrences across 3 files.") // summary shows total before pagination + }) + + t.Run("offset exceeds result count", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "files_with_matches", "offset": 100}`) + assert.NoError(t, err) + assert.Contains(t, result, "Found 3 files") // still shows total count + }) + + t.Run("negative offset treated as zero", func(t *testing.T) { + result, err := invokeTool(t, grepTool, `{"pattern": "match", "output_mode": "files_with_matches", "offset": -5}`) + assert.NoError(t, err) + lines := strings.Split(strings.TrimSpace(result), "\n") + assert.Equal(t, 4, len(lines)) // 1 summary + 3 files + }) +} + +func TestApplyPagination(t *testing.T) { + t.Run("basic pagination", func(t *testing.T) { + items := []string{"a", "b", "c", "d", "e"} + result := applyPagination(items, 0, 3) + assert.Equal(t, []string{"a", "b", "c"}, result) + }) + + t.Run("with offset", func(t *testing.T) { + items := []string{"a", "b", "c", "d", "e"} + result := applyPagination(items, 2, 2) + assert.Equal(t, []string{"c", "d"}, result) + }) + + t.Run("offset exceeds length", func(t *testing.T) { + items := []string{"a", "b", "c"} + result := applyPagination(items, 10, 5) + assert.Equal(t, []string{}, result) + }) + + t.Run("negative offset", func(t *testing.T) { + items := []string{"a", "b", "c"} + result := applyPagination(items, -1, 2) + assert.Equal(t, []string{"a", "b"}, result) + }) + + t.Run("zero head limit means no limit", func(t *testing.T) { + items := []string{"a", "b", "c", "d", "e"} + result := applyPagination(items, 1, 0) + assert.Equal(t, []string{"b", "c", "d", "e"}, result) + }) +} + +func TestCustomToolNames(t *testing.T) { + backend := setupTestBackend() + ctx := context.Background() + + t.Run("custom tool names applied to individual tools", func(t *testing.T) { + customLsName := "list_files" + customReadName := "read" + customWriteName := "write" + customEditName := "edit" + customGlobName := "find_files" + customGrepName := "search" + + lsTool, err := newLsTool(backend, customLsName, "") + assert.NoError(t, err) + info, _ := lsTool.Info(ctx) + assert.Equal(t, "list_files", info.Name) + + readTool, err := newReadFileTool(backend, customReadName, "") + assert.NoError(t, err) + info, _ = readTool.Info(ctx) + assert.Equal(t, "read", info.Name) + + writeTool, err := newWriteFileTool(backend, customWriteName, "") + assert.NoError(t, err) + info, _ = writeTool.Info(ctx) + assert.Equal(t, "write", info.Name) + + editTool, err := newEditFileTool(backend, customEditName, "") + assert.NoError(t, err) + info, _ = editTool.Info(ctx) + assert.Equal(t, "edit", info.Name) + + globTool, err := newGlobTool(backend, customGlobName, "") + assert.NoError(t, err) + info, _ = globTool.Info(ctx) + assert.Equal(t, "find_files", info.Name) + + grepTool, err := newGrepTool(backend, customGrepName, "") + assert.NoError(t, err) + info, _ = grepTool.Info(ctx) + assert.Equal(t, "search", info.Name) + }) + + t.Run("default tool names when custom names not provided", func(t *testing.T) { + lsTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + info, _ := lsTool.Info(ctx) + assert.Equal(t, ToolNameLs, info.Name) + + readTool, err := newReadFileTool(backend, "", "") + assert.NoError(t, err) + info, _ = readTool.Info(ctx) + assert.Equal(t, ToolNameReadFile, info.Name) + + writeTool, err := newWriteFileTool(backend, "", "") + assert.NoError(t, err) + info, _ = writeTool.Info(ctx) + assert.Equal(t, ToolNameWriteFile, info.Name) + + editTool, err := newEditFileTool(backend, "", "") + assert.NoError(t, err) + info, _ = editTool.Info(ctx) + assert.Equal(t, ToolNameEditFile, info.Name) + + globTool, err := newGlobTool(backend, "", "") + assert.NoError(t, err) + info, _ = globTool.Info(ctx) + assert.Equal(t, ToolNameGlob, info.Name) + + grepTool, err := newGrepTool(backend, "", "") + assert.NoError(t, err) + info, _ = grepTool.Info(ctx) + assert.Equal(t, ToolNameGrep, info.Name) + }) + + t.Run("custom execute tool name", func(t *testing.T) { + customExecuteName := "run_command" + shellBackend := &mockShellBackend{ + Backend: backend, + resp: &filesystem.ExecuteResponse{Output: "ok"}, + } + + executeTool, err := newExecuteTool(shellBackend, customExecuteName, "") + assert.NoError(t, err) + info, _ := executeTool.Info(ctx) + assert.Equal(t, "run_command", info.Name) + }) + + t.Run("custom tool names via ToolConfig in getFilesystemTools", func(t *testing.T) { + customLsName := "list_files" + customReadName := "read" + + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Name: customLsName, + }, + ReadFileToolConfig: &ToolConfig{ + Name: customReadName, + }, + }) + assert.NoError(t, err) + + toolNames := make(map[string]bool) + for _, to := range tools { + info, _ := to.Info(ctx) + toolNames[info.Name] = true + } + + assert.True(t, toolNames["list_files"]) + assert.True(t, toolNames["read"]) + }) + + t.Run("custom tool names via ToolConfig in middleware", func(t *testing.T) { + customLsName := "list_files" + customReadName := "read" + + m, err := New(ctx, &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Name: customLsName, + }, + ReadFileToolConfig: &ToolConfig{ + Name: customReadName, + }, + }) + assert.NoError(t, err) + + fm, ok := m.(*typedFilesystemMiddleware[*schema.Message]) + assert.True(t, ok) + + toolNames := make(map[string]bool) + for _, to := range fm.additionalTools { + info, _ := to.Info(ctx) + toolNames[info.Name] = true + } + + assert.True(t, toolNames["list_files"]) + assert.True(t, toolNames["read"]) + }) +} + +func TestSelectToolName(t *testing.T) { + t.Run("returns custom name when provided", func(t *testing.T) { + customName := "custom_tool" + result := selectToolName(customName, "default_tool") + assert.Equal(t, "custom_tool", result) + }) + + t.Run("returns default name when custom name is nil", func(t *testing.T) { + result := selectToolName("", "default_tool") + assert.Equal(t, "default_tool", result) + }) +} + +func TestGetOrCreateTool(t *testing.T) { + backend := setupTestBackend() + + t.Run("returns custom tool when provided", func(t *testing.T) { + customTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + + result, err := getOrCreateTool(customTool, func() (tool.BaseTool, error) { + t.Fatal("createFunc should not be called when custom tool is provided") + return nil, nil + }) + + assert.NoError(t, err) + assert.Equal(t, customTool, result) + }) + + t.Run("calls createFunc when custom tool is nil", func(t *testing.T) { + expectedTool, err := newReadFileTool(backend, "", "") + assert.NoError(t, err) + + createFuncCalled := false + result, err := getOrCreateTool(nil, func() (tool.BaseTool, error) { + createFuncCalled = true + return expectedTool, nil + }) + + assert.NoError(t, err) + assert.True(t, createFuncCalled, "createFunc should be called when custom tool is nil") + assert.Equal(t, expectedTool, result) + }) + + t.Run("returns nil when custom tool is nil and createFunc returns nil", func(t *testing.T) { + result, err := getOrCreateTool(nil, func() (tool.BaseTool, error) { + return nil, nil + }) + + assert.NoError(t, err) + assert.Nil(t, result) + }) + + t.Run("propagates error from createFunc", func(t *testing.T) { + expectedErr := assert.AnError + + result, err := getOrCreateTool(nil, func() (tool.BaseTool, error) { + return nil, expectedErr + }) + + assert.Error(t, err) + assert.Equal(t, expectedErr, err) + assert.Nil(t, result) + }) +} + +func TestCustomTools(t *testing.T) { + backend := setupTestBackend() + ctx := context.Background() + + t.Run("custom ls tool is used via ToolConfig", func(t *testing.T) { + customLsTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + LsToolConfig: &ToolConfig{ + CustomTool: customLsTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customLsTool, tools[0]) + }) + + t.Run("custom read file tool is used via ToolConfig", func(t *testing.T) { + customReadTool, err := newReadFileTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + ReadFileToolConfig: &ToolConfig{ + CustomTool: customReadTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customReadTool, tools[0]) + }) + + t.Run("custom write file tool is used via ToolConfig", func(t *testing.T) { + customWriteTool, err := newWriteFileTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + WriteFileToolConfig: &ToolConfig{ + CustomTool: customWriteTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customWriteTool, tools[0]) + }) + + t.Run("custom edit file tool is used via ToolConfig", func(t *testing.T) { + customEditTool, err := newEditFileTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + EditFileToolConfig: &ToolConfig{ + CustomTool: customEditTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customEditTool, tools[0]) + }) + + t.Run("custom glob tool is used via ToolConfig", func(t *testing.T) { + customGlobTool, err := newGlobTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + GlobToolConfig: &ToolConfig{ + CustomTool: customGlobTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customGlobTool, tools[0]) + }) + + t.Run("custom grep tool is used via ToolConfig", func(t *testing.T) { + customGrepTool, err := newGrepTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + GrepToolConfig: &ToolConfig{ + CustomTool: customGrepTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, customGrepTool, tools[0]) + }) + + t.Run("multiple custom tools can be used together", func(t *testing.T) { + customLsTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + customReadTool, err := newReadFileTool(backend, "", "") + assert.NoError(t, err) + customGlobTool, err := newGlobTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + LsToolConfig: &ToolConfig{ + CustomTool: customLsTool, + }, + ReadFileToolConfig: &ToolConfig{ + CustomTool: customReadTool, + }, + GlobToolConfig: &ToolConfig{ + CustomTool: customGlobTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 3) + }) + + t.Run("custom tools take precedence over backend", func(t *testing.T) { + customLsTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + CustomTool: customLsTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + lsToolFound := false + for _, t := range tools { + if t == customLsTool { + lsToolFound = true + break + } + } + assert.True(t, lsToolFound, "custom ls tool should be in the tools list") + }) + + t.Run("backend tools are created when custom tools not provided", func(t *testing.T) { + config := &MiddlewareConfig{ + Backend: backend, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Greater(t, len(tools), 0, "should create backend tools when custom tools not provided") + }) +} + +func TestToolConfig(t *testing.T) { + backend := setupTestBackend() + ctx := context.Background() + + t.Run("use new ToolConfig", func(t *testing.T) { + customName := "my_ls" + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Name: customName, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 6) + + var lsToolFound bool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == "my_ls" { + lsToolFound = true + break + } + } + assert.True(t, lsToolFound) + }) + + t.Run("ToolConfig disabled", func(t *testing.T) { + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Disable: true, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 5) + + for _, tool := range tools { + info, _ := tool.Info(ctx) + assert.NotEqual(t, ToolNameLs, info.Name) + } + }) + + t.Run("ToolConfig with custom tool", func(t *testing.T) { + customLsTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + CustomTool: customLsTool, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var lsToolFound bool + for _, tool := range tools { + if tool == customLsTool { + lsToolFound = true + break + } + } + assert.True(t, lsToolFound) + }) + + t.Run("ToolConfig Desc takes precedence over legacy Desc", func(t *testing.T) { + customDesc := "new description" + legacyDesc := "old description" + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Desc: &customDesc, + }, + CustomLsToolDesc: &legacyDesc, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var found bool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == ToolNameLs && info.Desc == "new description" { + found = true + break + } + } + assert.True(t, found) + }) + + t.Run("legacy Desc field still works", func(t *testing.T) { + legacyDesc := "legacy description" + config := &MiddlewareConfig{ + Backend: backend, + CustomLsToolDesc: &legacyDesc, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var found bool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == ToolNameLs && info.Desc == "legacy description" { + found = true + break + } + } + assert.True(t, found) + }) + + t.Run("multiple ToolConfig", func(t *testing.T) { + lsName := "my_ls" + readName := "my_read" + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Name: lsName, + }, + ReadFileToolConfig: &ToolConfig{ + Name: readName, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + toolNames := make(map[string]bool) + for _, tool := range tools { + info, _ := tool.Info(ctx) + toolNames[info.Name] = true + } + + assert.True(t, toolNames["my_ls"]) + assert.True(t, toolNames["my_read"]) + }) +} + +func TestToolConfigEdgeCases(t *testing.T) { + backend := setupTestBackend() + ctx := context.Background() + + t.Run("nil ToolConfig.Desc with nil legacyDesc", func(t *testing.T) { + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Desc: nil, + }, + CustomLsToolDesc: nil, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var lsTool tool.BaseTool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == ToolNameLs { + lsTool = tool + break + } + } + assert.NotNil(t, lsTool, "ls tool should be created even with nil Desc") + }) + + t.Run("nil ToolConfig.Desc falls back to legacyDesc", func(t *testing.T) { + legacyDesc := "legacy description from pointer" + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Desc: nil, + }, + CustomLsToolDesc: &legacyDesc, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var found bool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == ToolNameLs && info.Desc == "legacy description from pointer" { + found = true + break + } + } + assert.True(t, found, "nil ToolConfig.Desc should fall back to legacyDesc") + }) + + t.Run("CustomTool with Disable flag should not create tool", func(t *testing.T) { + customLsTool, err := newLsTool(backend, "", "") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + CustomTool: customLsTool, + Disable: true, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + for _, tool := range tools { + info, _ := tool.Info(ctx) + assert.NotEqual(t, ToolNameLs, info.Name, "disabled tool should not be created even if CustomTool is set") + } + }) + + t.Run("multiple ToolConfig with conflicting settings", func(t *testing.T) { + legacyDesc := "legacy ls desc" + customDesc := "custom desc" + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Name: "custom_ls", + Desc: &customDesc, + Disable: false, + }, + CustomLsToolDesc: &legacyDesc, + ReadFileToolConfig: &ToolConfig{ + Disable: true, + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + hasLsTool := false + hasReadTool := false + + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == "custom_ls" { + hasLsTool = true + assert.Equal(t, "custom desc", info.Desc, "ToolConfig.Desc should take precedence over legacy") + } + if info.Name == ToolNameReadFile { + hasReadTool = true + } + } + + assert.True(t, hasLsTool, "ls tool should be created") + assert.False(t, hasReadTool, "read_file tool should be disabled") + }) + + t.Run("nil ToolConfig with nil legacyDesc creates default tool", func(t *testing.T) { + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: nil, + CustomLsToolDesc: nil, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var lsTool tool.BaseTool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == ToolNameLs { + lsTool = tool + break + } + } + assert.NotNil(t, lsTool, "tool should be created with backend even when config is nil") + }) + + t.Run("empty Name in ToolConfig uses default name", func(t *testing.T) { + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + Name: "", + }, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var lsTool tool.BaseTool + for _, tool := range tools { + info, _ := tool.Info(ctx) + if info.Name == ToolNameLs { + lsTool = tool + break + } + } + assert.NotNil(t, lsTool, "tool should use default name when Name is empty") + }) +} + +func TestGetFilesystemTools_DisableAllTools(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{Disable: true}, + ReadFileToolConfig: &ToolConfig{Disable: true}, + WriteFileToolConfig: &ToolConfig{Disable: true}, + EditFileToolConfig: &ToolConfig{Disable: true}, + GlobToolConfig: &ToolConfig{Disable: true}, + GrepToolConfig: &ToolConfig{Disable: true}, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 0) +} + +func TestGetFilesystemTools_StreamingShell(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("returns 7 tools with StreamingShell", func(t *testing.T) { + mockSS := &mockStreamingShell{} + tools, err := getFilesystemTools(ctx, &MiddlewareConfig{ + Backend: backend, + StreamingShell: mockSS, + }) + assert.NoError(t, err) + assert.Len(t, tools, 7) + + toolNames := make([]string, 0, len(tools)) + for _, to := range tools { + info, _ := to.Info(ctx) + toolNames = append(toolNames, info.Name) + } + assert.Contains(t, toolNames, ToolNameExecute) + }) + + t.Run("StreamingShell takes precedence over Shell", func(t *testing.T) { + mockSS := &mockStreamingShell{} + shellBackend := &mockShellBackend{ + Backend: backend, + resp: &filesystem.ExecuteResponse{Output: "ok"}, + } + + // When both are set, Validate should fail + config := &MiddlewareConfig{ + Backend: backend, + Shell: shellBackend, + StreamingShell: mockSS, + } + err := config.Validate() + assert.Error(t, err) + assert.Contains(t, err.Error(), "shell and streaming shell should not be both set") + }) +} + +func TestGetFilesystemTools_NilBackend(t *testing.T) { + ctx := context.Background() + + t.Run("nil backend with shell only returns execute tool", func(t *testing.T) { + mockSS := &mockStreamingShell{} + config := &MiddlewareConfig{ + Backend: nil, + StreamingShell: mockSS, + } + // Validate should fail, but getFilesystemTools itself handles nil backend gracefully + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + // Only execute tool should be returned since backend is nil + assert.Len(t, tools, 1) + + info, _ := tools[0].Info(ctx) + assert.Equal(t, ToolNameExecute, info.Name) + }) + + t.Run("nil backend with regular Shell returns execute tool", func(t *testing.T) { + mockShell := &mockShellBackend{ + resp: &filesystem.ExecuteResponse{Output: "ok"}, + } + config := &MiddlewareConfig{ + Backend: nil, + Shell: mockShell, + } + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 1) + + info, _ := tools[0].Info(ctx) + assert.Equal(t, ToolNameExecute, info.Name) + }) + + t.Run("nil backend and nil shell returns empty tools", func(t *testing.T) { + config := &MiddlewareConfig{ + Backend: nil, + } + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 0) + }) +} + +func TestGetFilesystemTools_PartialDisable(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{Disable: true}, + ReadFileToolConfig: &ToolConfig{Disable: true}, + } + + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + assert.Len(t, tools, 4) + + toolNames := make([]string, 0, len(tools)) + for _, to := range tools { + info, _ := to.Info(ctx) + toolNames = append(toolNames, info.Name) + } + assert.NotContains(t, toolNames, ToolNameLs) + assert.NotContains(t, toolNames, ToolNameReadFile) + assert.Contains(t, toolNames, ToolNameWriteFile) + assert.Contains(t, toolNames, ToolNameEditFile) + assert.Contains(t, toolNames, ToolNameGlob) + assert.Contains(t, toolNames, ToolNameGrep) +} + +type mockStreamingShell struct{} + +func (m *mockStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{ + Output: "streaming output", + ExitCode: ptrOf(0), + }, nil) + }() + return sr, nil +} + +type mockStreamingShellWithError struct{} + +func (m *mockStreamingShellWithError) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + return nil, fmt.Errorf("streaming shell error") +} + +type mockStreamingShellWithRecvError struct{} + +func (m *mockStreamingShellWithRecvError) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(nil, fmt.Errorf("recv error during streaming")) + }() + return sr, nil +} + +type mockStreamingShellWithExitCode struct { + exitCode int +} + +func (m *mockStreamingShellWithExitCode) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{ + Output: "some output", + ExitCode: ptrOf(m.exitCode), + }, nil) + }() + return sr, nil +} + +type mockStreamingShellNoOutput struct{} + +func (m *mockStreamingShellNoOutput) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{ + ExitCode: ptrOf(0), + }, nil) + }() + return sr, nil +} + +type mockStreamingShellTruncated struct{} + +func (m *mockStreamingShellTruncated) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{ + Output: "partial", + Truncated: true, + ExitCode: ptrOf(0), + }, nil) + }() + return sr, nil +} + +type mockStreamingShellNilChunk struct{} + +func (m *mockStreamingShellNilChunk) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(nil, nil) + sw.Send(&filesystem.ExecuteResponse{ + Output: "after nil", + ExitCode: ptrOf(0), + }, nil) + }() + return sr, nil +} + +func TestNewStreamingExecuteTool(t *testing.T) { + t.Run("successful streaming execution", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShell{}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "echo hello"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + assert.True(t, len(chunks) > 0) + result := strings.Join(chunks, "") + assert.Contains(t, result, "streaming output") + }) + + t.Run("streaming execution with ExecuteStreaming error", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShellWithError{}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + _, err = st.StreamableRun(context.Background(), `{"command": "fail"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "streaming shell error") + }) + + t.Run("streaming execution with recv error", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShellWithRecvError{}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "echo hello"}`) + assert.NoError(t, err) + defer sr.Close() + + var gotError bool + for { + _, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + if recvErr != nil { + gotError = true + assert.Contains(t, recvErr.Error(), "recv error during streaming") + break + } + } + assert.True(t, gotError) + }) + + t.Run("streaming execution with non-zero exit code", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShellWithExitCode{exitCode: 1}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "false"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + result := strings.Join(chunks, "") + assert.Contains(t, result, "[Command failed with exit code 1]") + }) + + t.Run("streaming execution with zero exit code and no output", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShellNoOutput{}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "true"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + result := strings.Join(chunks, "") + assert.Contains(t, result, "[Command executed successfully with no output]") + }) + + t.Run("streaming execution with truncated output", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShellTruncated{}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "cat largefile"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + result := strings.Join(chunks, "") + assert.Contains(t, result, "partial") + assert.Contains(t, result, "[Output was truncated due to size limits]") + }) + + t.Run("streaming execution with nil chunk skipped", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShellNilChunk{}, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "echo test"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + result := strings.Join(chunks, "") + assert.Contains(t, result, "after nil") + }) + + t.Run("streaming execution with custom name and desc", func(t *testing.T) { + executeTool, err := newStreamingExecuteTool(&mockStreamingShell{}, "custom_execute", "custom desc") + assert.NoError(t, err) + + info, err := executeTool.Info(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "custom_execute", info.Name) + assert.Equal(t, "custom desc", info.Desc) + }) +} + +func TestNew_StreamingShell(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("StreamingShell adds streaming execute tool", func(t *testing.T) { + m, err := New(ctx, &MiddlewareConfig{ + Backend: backend, + StreamingShell: &mockStreamingShell{}, + }) + assert.NoError(t, err) + + fm, ok := m.(*typedFilesystemMiddleware[*schema.Message]) + assert.True(t, ok) + assert.Len(t, fm.additionalTools, 7) + }) + + t.Run("both Shell and StreamingShell returns error", func(t *testing.T) { + _, err := New(ctx, &MiddlewareConfig{ + Backend: backend, + Shell: &mockShellBackend{Backend: backend, resp: &filesystem.ExecuteResponse{Output: "ok"}}, + StreamingShell: &mockStreamingShell{}, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "shell and streaming shell should not be both set") + }) +} + +func TestNewMiddleware_Validation(t *testing.T) { + ctx := context.Background() + + t.Run("nil config returns error", func(t *testing.T) { + _, err := NewMiddleware(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config should not be nil") + }) + + t.Run("nil backend returns error", func(t *testing.T) { + _, err := NewMiddleware(ctx, &Config{Backend: nil}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend should not be nil") + }) + + t.Run("both Shell and StreamingShell returns error", func(t *testing.T) { + backend := setupTestBackend() + _, err := NewMiddleware(ctx, &Config{ + Backend: backend, + Shell: &mockShellBackend{Backend: backend, resp: &filesystem.ExecuteResponse{Output: "ok"}}, + StreamingShell: &mockStreamingShell{}, + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "shell and streaming shell should not be both set") + }) +} + +func TestMiddlewareConfig_Validate(t *testing.T) { + t.Run("nil config returns error", func(t *testing.T) { + var c *MiddlewareConfig + err := c.Validate() + assert.Error(t, err) + assert.Contains(t, err.Error(), "config should not be nil") + }) + + t.Run("nil backend returns error", func(t *testing.T) { + c := &MiddlewareConfig{} + err := c.Validate() + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend should not be nil") + }) + + t.Run("both shells returns error", func(t *testing.T) { + c := &MiddlewareConfig{ + Backend: setupTestBackend(), + Shell: &mockShellBackend{}, + StreamingShell: &mockStreamingShell{}, + } + err := c.Validate() + assert.Error(t, err) + assert.Contains(t, err.Error(), "shell and streaming shell should not be both set") + }) + + t.Run("valid config passes", func(t *testing.T) { + c := &MiddlewareConfig{ + Backend: setupTestBackend(), + } + err := c.Validate() + assert.NoError(t, err) + }) +} + +func TestNewStreamingExecuteTool_MultipleChunks(t *testing.T) { + mockSS := &mockStreamingShellMultiChunk{} + executeTool, err := newStreamingExecuteTool(mockSS, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "long-running"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + // Should have received multiple chunks + assert.True(t, len(chunks) >= 3) + result := strings.Join(chunks, "") + assert.Contains(t, result, "chunk1") + assert.Contains(t, result, "chunk2") + assert.Contains(t, result, "chunk3") +} + +type mockStreamingShellMultiChunk struct{} + +func (m *mockStreamingShellMultiChunk) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{Output: "chunk1\n"}, nil) + sw.Send(&filesystem.ExecuteResponse{Output: "chunk2\n"}, nil) + sw.Send(&filesystem.ExecuteResponse{Output: "chunk3\n", ExitCode: ptrOf(0)}, nil) + }() + return sr, nil +} + +func TestNewStreamingExecuteTool_ExitCodeOnlyInLastChunk(t *testing.T) { + mockSS := &mockStreamingShellExitCodeLast{exitCode: 2} + executeTool, err := newStreamingExecuteTool(mockSS, "", "") + assert.NoError(t, err) + + st := executeTool.(tool.StreamableTool) + sr, err := st.StreamableRun(context.Background(), `{"command": "fail-at-end"}`) + assert.NoError(t, err) + defer sr.Close() + + var chunks []string + for { + chunk, recvErr := sr.Recv() + if recvErr == io.EOF { + break + } + assert.NoError(t, recvErr) + chunks = append(chunks, chunk) + } + result := strings.Join(chunks, "") + assert.Contains(t, result, "output line") + assert.Contains(t, result, "[Command failed with exit code 2]") +} + +type mockStreamingShellExitCodeLast struct { + exitCode int +} + +func (m *mockStreamingShellExitCodeLast) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + sr, sw := schema.Pipe[*filesystem.ExecuteResponse](10) + go func() { + defer sw.Close() + sw.Send(&filesystem.ExecuteResponse{Output: "output line"}, nil) + sw.Send(&filesystem.ExecuteResponse{ExitCode: ptrOf(m.exitCode)}, nil) + }() + return sr, nil +} + +func TestConvExecuteResponse_NilResponse(t *testing.T) { + result := convExecuteResponse(nil) + assert.Equal(t, "", result) +} + +func TestConvExecuteResponse_NilExitCode(t *testing.T) { + result := convExecuteResponse(&filesystem.ExecuteResponse{ + Output: "some output", + }) + assert.Equal(t, "some output", result) +} + +func TestConfig_Validate(t *testing.T) { + t.Run("nil config returns error", func(t *testing.T) { + var c *Config + err := c.Validate() + assert.Error(t, err) + }) + + t.Run("nil backend returns error", func(t *testing.T) { + c := &Config{} + err := c.Validate() + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend should not be nil") + }) + + t.Run("both shells returns error", func(t *testing.T) { + c := &Config{ + Backend: setupTestBackend(), + Shell: &mockShellBackend{}, + StreamingShell: &mockStreamingShell{}, + } + err := c.Validate() + assert.Error(t, err) + }) + + t.Run("valid config passes", func(t *testing.T) { + c := &Config{ + Backend: setupTestBackend(), + } + err := c.Validate() + assert.NoError(t, err) + }) +} + +func TestGetFilesystemTools_CustomToolWithShell(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("custom tool replaces default for all disabled except custom", func(t *testing.T) { + customLs, err := newLsTool(backend, "my_ls", "my ls desc") + assert.NoError(t, err) + + config := &MiddlewareConfig{ + Backend: backend, + LsToolConfig: &ToolConfig{ + CustomTool: customLs, + }, + } + tools, err := getFilesystemTools(ctx, config) + assert.NoError(t, err) + + var found bool + for _, to := range tools { + info, _ := to.Info(ctx) + if info.Name == "my_ls" { + found = true + break + } + } + assert.True(t, found) + }) +} + +func TestMergeToolConfigWithDesc(t *testing.T) { + config := &MiddlewareConfig{Backend: setupTestBackend()} + + t.Run("both nil returns empty ToolConfig", func(t *testing.T) { + result := config.mergeToolConfigWithDesc(nil, nil) + assert.NotNil(t, result) + assert.Equal(t, "", result.Name) + assert.Nil(t, result.Desc) + assert.False(t, result.Disable) + }) + + t.Run("nil toolConfig with legacyDesc", func(t *testing.T) { + desc := "legacy" + result := config.mergeToolConfigWithDesc(nil, &desc) + assert.NotNil(t, result) + assert.Equal(t, "legacy", *result.Desc) + }) + + t.Run("toolConfig with Desc overrides legacyDesc", func(t *testing.T) { + tcDesc := "tc desc" + legacyDesc := "legacy" + tc := &ToolConfig{Desc: &tcDesc} + result := config.mergeToolConfigWithDesc(tc, &legacyDesc) + assert.Equal(t, "tc desc", *result.Desc) + }) + + t.Run("toolConfig with nil Desc falls back to legacyDesc", func(t *testing.T) { + legacyDesc := "legacy" + tc := &ToolConfig{Name: "custom"} + result := config.mergeToolConfigWithDesc(tc, &legacyDesc) + assert.Equal(t, "legacy", *result.Desc) + assert.Equal(t, "custom", result.Name) + }) + + t.Run("toolConfig with nil Desc and nil legacyDesc", func(t *testing.T) { + tc := &ToolConfig{Name: "custom"} + result := config.mergeToolConfigWithDesc(tc, nil) + assert.Nil(t, result.Desc) + assert.Equal(t, "custom", result.Name) + }) +} + +func TestNewMiddleware_WithShell(t *testing.T) { + ctx := context.Background() + backend := setupTestBackend() + + t.Run("Shell backend creates execute tool", func(t *testing.T) { + shellBackend := &mockShellBackend{ + Backend: backend, + resp: &filesystem.ExecuteResponse{Output: "ok"}, + } + m, err := NewMiddleware(ctx, &Config{ + Backend: backend, + Shell: shellBackend, + }) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 7) + }) + + t.Run("StreamingShell backend creates streaming execute tool", func(t *testing.T) { + m, err := NewMiddleware(ctx, &Config{ + Backend: backend, + StreamingShell: &mockStreamingShell{}, + }) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 7) + }) +} + +func TestNewExecuteTool_ShellError(t *testing.T) { + mockShell := &mockShellBackendWithError{} + executeTool, err := newExecuteTool(mockShell, "", "") + assert.NoError(t, err) + + result, err := invokeTool(t, executeTool, `{"command": "fail"}`) + assert.Error(t, err) + assert.Equal(t, "", result) + assert.Contains(t, err.Error(), "shell execution error") +} + +type mockShellBackendWithError struct{} + +func (m *mockShellBackendWithError) Execute(ctx context.Context, req *filesystem.ExecuteRequest) (*filesystem.ExecuteResponse, error) { + return nil, errors.New("shell execution error") +} + +// multiModalBackend wraps InMemoryBackend and implements MultiModalReader for testing. +type multiModalBackend struct { + *filesystem.InMemoryBackend + multiModalReadFunc func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) +} + +func (b *multiModalBackend) MultiModalRead(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + return b.multiModalReadFunc(ctx, req) +} + +func TestMultiModalReadFileTool_TextOnly(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + ct, err := base.Read(ctx, &req.ReadRequest) + if err != nil { + return nil, err + } + return &filesystem.MultiFileContent{ + FileContent: ct, + }, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + result, err := mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/file1.txt", "offset": 0, "limit": 100}`}) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Contains(t, result.Parts[0].Text, "line1") + assert.Contains(t, result.Parts[0].Text, "line5") +} + +func TestMultiModalReadFileTool_Multimodal(t *testing.T) { + base := setupTestBackend() + imgData := []byte("rawimagedata") + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + return &filesystem.MultiFileContent{ + Parts: []filesystem.FileContentPart{ + { + Type: filesystem.FileContentPartTypeImage, + MIMEType: "image/png", + Data: imgData, + }, + }, + }, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + result, err := mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/image.png", "offset": 0, "limit": 100}`}) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeImage, result.Parts[0].Type) + + // Verify base64 encoding correctness + assert.NotNil(t, result.Parts[0].Image) + assert.Equal(t, "image/png", result.Parts[0].Image.MIMEType) + assert.Equal(t, base64.StdEncoding.EncodeToString(imgData), *result.Parts[0].Image.Base64Data) +} + +func TestMultiModalReadFileTool_FileType(t *testing.T) { + base := setupTestBackend() + pdfData := []byte("fakepdfcontent") + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + return &filesystem.MultiFileContent{ + Parts: []filesystem.FileContentPart{ + { + Type: filesystem.FileContentPartTypePDF, + MIMEType: "application/pdf", + Data: pdfData, + }, + }, + }, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + result, err := mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/doc.pdf", "offset": 0, "limit": 100}`}) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeFile, result.Parts[0].Type) + assert.NotNil(t, result.Parts[0].File) + assert.Equal(t, "application/pdf", result.Parts[0].File.MIMEType) + assert.Equal(t, base64.StdEncoding.EncodeToString(pdfData), *result.Parts[0].File.Base64Data) +} + +func TestMultiModalReadFileTool_UnsupportedPartType(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + return &filesystem.MultiFileContent{ + Parts: []filesystem.FileContentPart{ + { + Type: filesystem.FileContentPartType("unknown"), + MIMEType: "application/octet-stream", + Data: []byte("data"), + }, + }, + }, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + _, err = mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/file.bin", "offset": 0, "limit": 100}`}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unsupported FileContentPartType") +} + +func TestMultiModalReadFileTool_PagesPassThrough(t *testing.T) { + base := setupTestBackend() + var capturedPages string + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + capturedPages = req.Pages + return &filesystem.MultiFileContent{FileContent: &filesystem.FileContent{Content: "page content"}}, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + _, err = mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/doc.pdf", "pages": "1-5"}`}) + assert.NoError(t, err) + assert.Equal(t, "1-5", capturedPages) +} + +func TestMultiModalReadFileTool_BackendNotMultiModalReader(t *testing.T) { + base := setupTestBackend() + _, err := newMultiModalReadFileTool(base, "", "") + assert.Error(t, err) + assert.Contains(t, err.Error(), "MultiModalReader") +} + +func TestUseMultiModalRead_Routing(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + ct, err := base.Read(ctx, &req.ReadRequest) + if err != nil { + return nil, err + } + return &filesystem.MultiFileContent{FileContent: ct}, nil + }, + } + + // UseMultiModalRead=false should create standard tool + tools, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Backend: base, + UseMultiModalRead: false, + }) + assert.NoError(t, err) + for _, tl := range tools { + info, _ := tl.Info(context.Background()) + if info != nil && info.Name == ToolNameReadFile { + _, isEnhanced := tl.(tool.EnhancedInvokableTool) + assert.False(t, isEnhanced, "should be standard InvokableTool when UseMultiModalRead=false") + } + } + + // UseMultiModalRead=true with enhanced backend should create enhanced tool + tools2, err := getFilesystemTools(context.Background(), &MiddlewareConfig{ + Backend: eb, + UseMultiModalRead: true, + }) + assert.NoError(t, err) + for _, tl := range tools2 { + info, _ := tl.Info(context.Background()) + if info != nil && info.Name == ToolNameReadFile { + _, isEnhanced := tl.(tool.EnhancedInvokableTool) + assert.True(t, isEnhanced, "should be EnhancedInvokableTool when UseMultiModalRead=true") + } + } +} + +// TestMultiModalReadFileTool_SchemaContainsAllFields verifies that the JSON schema +// exposed to the LLM includes both the embedded readFileArgs fields (file_path, +// offset, limit) and the enhanced-only "pages" field. Guards against the +// jsonschema library failing to flatten an unexported anonymous embedded struct. +func TestMultiModalReadFileTool_SchemaContainsAllFields(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + ct, err := base.Read(ctx, &req.ReadRequest) + if err != nil { + return nil, err + } + return &filesystem.MultiFileContent{FileContent: ct}, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + info, err := mmTool.Info(context.Background()) + assert.NoError(t, err) + assert.NotNil(t, info) + + js, err := info.ParamsOneOf.ToJSONSchema() + assert.NoError(t, err) + assert.NotNil(t, js) + assert.NotNil(t, js.Properties, "schema should have properties") + + for _, field := range []string{"file_path", "offset", "limit", "pages"} { + _, ok := js.Properties.Get(field) + assert.True(t, ok, "expected JSON schema to contain field %q, schema=%+v", field, js.Properties) + } +} + +// TestMultiModalReadFileTool_CustomDescNoSuffix verifies that when a custom desc is +// provided, the multimodal suffix is NOT appended (user's desc replaces default). +func TestMultiModalReadFileTool_CustomDescNoSuffix(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + ct, err := base.Read(ctx, &req.ReadRequest) + if err != nil { + return nil, err + } + return &filesystem.MultiFileContent{FileContent: ct}, nil + }, + } + + customDesc := "my custom read tool description" + mmTool, err := newMultiModalReadFileTool(eb, "", customDesc) + assert.NoError(t, err) + + info, err := mmTool.Info(context.Background()) + assert.NoError(t, err) + assert.Equal(t, customDesc, info.Desc, "custom desc should not be augmented with multimodal suffix") + + // With empty desc (fallback to default), suffix should be appended. + defaultTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + defaultInfo, err := defaultTool.Info(context.Background()) + assert.NoError(t, err) + assert.Contains(t, defaultInfo.Desc, "multimodal", "default desc should include multimodal suffix") +} + +// TestMultiModalReadFileTool_EmptyPartDataError verifies that a FileContentPart +// with empty Data fails explicitly rather than silently encoding to an empty +// base64 string. +func TestMultiModalReadFileTool_EmptyPartDataError(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + return &filesystem.MultiFileContent{ + Parts: []filesystem.FileContentPart{ + {Type: filesystem.FileContentPartTypeImage, MIMEType: "image/png", Data: nil}, + }, + }, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + _, err = mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/x"}`}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "empty") +} + +// nilReadBackend wraps InMemoryBackend but returns nil, nil from Read. +type nilReadBackend struct { + *filesystem.InMemoryBackend +} + +func (b *nilReadBackend) Read(_ context.Context, _ *filesystem.ReadRequest) (*filesystem.FileContent, error) { + return nil, nil +} + +// TestReadFileTool_NilResult verifies that newReadFileTool does not panic when +// Backend.Read returns nil, and emits a human-readable fallback message instead. +func TestReadFileTool_NilResult(t *testing.T) { + base := setupTestBackend() + backend := &nilReadBackend{InMemoryBackend: base} + + readTool, err := newReadFileTool(backend, "", "") + assert.NoError(t, err) + + out, err := invokeTool(t, readTool, `{"file_path": "/missing.txt"}`) + assert.NoError(t, err) + assert.Contains(t, out, "No content found at path") + assert.Contains(t, out, "/missing.txt") +} + +// TestMultiModalReadFileTool_NilResult verifies that newMultiModalReadFileTool +// does not panic when MultiModalRead returns nil, and returns a text part with +// a human-readable fallback message. +func TestMultiModalReadFileTool_NilResult(t *testing.T) { + base := setupTestBackend() + eb := &multiModalBackend{ + InMemoryBackend: base, + multiModalReadFunc: func(ctx context.Context, req *filesystem.MultiModalReadRequest) (*filesystem.MultiFileContent, error) { + return nil, nil + }, + } + + mmTool, err := newMultiModalReadFileTool(eb, "", "") + assert.NoError(t, err) + + result, err := mmTool.(tool.EnhancedInvokableTool).InvokableRun( + context.Background(), &schema.ToolArgument{Text: `{"file_path": "/missing.txt"}`}) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Contains(t, result.Parts[0].Text, "No content found at path") + assert.Contains(t, result.Parts[0].Text, "/missing.txt") +} + +func TestValidatePages(t *testing.T) { + tests := []struct { + name string + pages string + wantErr string + }{ + {name: "single page", pages: "3"}, + {name: "valid range", pages: "1-10"}, + {name: "same start end", pages: "1-1"}, + {name: "max 20 pages", pages: "1-20"}, + {name: "trailing dash", pages: "1-", wantErr: "expected format"}, + {name: "leading dash", pages: "-5", wantErr: "expected format"}, + {name: "non-numeric", pages: "abc", wantErr: "expected format"}, + {name: "non-numeric end", pages: "1-abc", wantErr: "expected format"}, + {name: "zero start", pages: "0-5", wantErr: "expected format"}, + {name: "zero end", pages: "1-0", wantErr: "expected format"}, + {name: "end less than start", pages: "10-5", wantErr: "end page must be >= start page"}, + {name: "exceeds max pages", pages: "1-21", wantErr: "range exceeds maximum of 20"}, + {name: "large range", pages: "1-30", wantErr: "range exceeds maximum of 20"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validatePages(tt.pages) + if tt.wantErr == "" { + assert.NoError(t, err) + } else { + assert.ErrorContains(t, err, tt.wantErr) + } + }) + } +} diff --git a/adk/middlewares/filesystem/large_tool_result.go b/adk/middlewares/filesystem/large_tool_result.go new file mode 100644 index 0000000..215b69d --- /dev/null +++ b/adk/middlewares/filesystem/large_tool_result.go @@ -0,0 +1,177 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "strings" + "unicode/utf8" + + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type toolResultOffloadingConfig struct { + Backend filesystem.Backend + TokenLimit int + PathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error) +} + +func newToolResultOffloading(ctx context.Context, config *toolResultOffloadingConfig) compose.ToolMiddleware { + offloading := &toolResultOffloading{ + backend: config.Backend, + tokenLimit: config.TokenLimit, + pathGenerator: config.PathGenerator, + } + + if offloading.tokenLimit == 0 { + offloading.tokenLimit = 20000 + } + + if offloading.pathGenerator == nil { + offloading.pathGenerator = func(ctx context.Context, input *compose.ToolInput) (string, error) { + return fmt.Sprintf("/large_tool_result/%s", input.CallID), nil + } + } + + return compose.ToolMiddleware{ + Invokable: offloading.invoke, + Streamable: offloading.stream, + } +} + +type toolResultOffloading struct { + backend filesystem.Backend + tokenLimit int + pathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error) +} + +func (t *toolResultOffloading) invoke(endpoint compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + output, err := endpoint(ctx, input) + if err != nil { + return nil, err + } + result, err := t.handleResult(ctx, output.Result, input) + if err != nil { + return nil, err + } + return &compose.ToolOutput{Result: result}, nil + } +} + +func (t *toolResultOffloading) stream(endpoint compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + output, err := endpoint(ctx, input) + if err != nil { + return nil, err + } + result, err := concatString(output.Result) + if err != nil { + return nil, err + } + result, err = t.handleResult(ctx, result, input) + if err != nil { + return nil, err + } + return &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{result})}, nil + } +} + +func (t *toolResultOffloading) handleResult(ctx context.Context, result string, input *compose.ToolInput) (string, error) { + if len(result) > t.tokenLimit*4 { + path, err := t.pathGenerator(ctx, input) + if err != nil { + return "", err + } + + nResult := formatToolMessage(result) + msgTemplate := internal.SelectPrompt(internal.I18nPrompts{ + English: tooLargeToolMessage, + Chinese: tooLargeToolMessageChinese, + }) + nResult, err = pyfmt.Fmt(msgTemplate, map[string]any{ + "tool_call_id": input.CallID, + "file_path": path, + "content_sample": nResult, + }) + if err != nil { + return "", err + } + + err = t.backend.Write(ctx, &WriteRequest{ + FilePath: path, + Content: result, + }) + if err != nil { + return "", err + } + + return nResult, nil + } + + return result, nil +} + +func concatString(sr *schema.StreamReader[string]) (string, error) { + if sr == nil { + return "", errors.New("stream is nil") + } + sb := strings.Builder{} + for { + str, err := sr.Recv() + if errors.Is(err, io.EOF) { + return sb.String(), nil + } + if err != nil { + return "", err + } + sb.WriteString(str) + } +} + +func formatToolMessage(s string) string { + reader := bufio.NewScanner(strings.NewReader(s)) + var b strings.Builder + + lineNum := 1 + for reader.Scan() { + if lineNum > 10 { + break + } + line := reader.Text() + + if utf8.RuneCountInString(line) > 1000 { + runes := []rune(line) + line = string(runes[:1000]) + } + + b.WriteString(fmt.Sprintf("%d: %s\n", lineNum, line)) + + lineNum++ + } + + return b.String() +} diff --git a/adk/middlewares/filesystem/large_tool_result_test.go b/adk/middlewares/filesystem/large_tool_result_test.go new file mode 100644 index 0000000..42cf27b --- /dev/null +++ b/adk/middlewares/filesystem/large_tool_result_test.go @@ -0,0 +1,611 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + "testing" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// mockBackend is a simple in-memory backend for testing +type mockBackend struct { + files map[string]string +} + +func newMockBackend() *mockBackend { + return &mockBackend{ + files: make(map[string]string), + } +} + +func (m *mockBackend) Write(ctx context.Context, req *WriteRequest) error { + m.files[req.FilePath] = req.Content + return nil +} + +func (m *mockBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) { + content, ok := m.files[req.FilePath] + if !ok { + return nil, errors.New("file not found") + } + return &FileContent{Content: content}, nil +} + +func (m *mockBackend) LsInfo(ctx context.Context, _ *LsInfoRequest) ([]FileInfo, error) { + return nil, nil +} + +func (m *mockBackend) GrepRaw(ctx context.Context, _ *GrepRequest) ([]GrepMatch, error) { + return nil, nil +} + +func (m *mockBackend) GlobInfo(ctx context.Context, _ *GlobInfoRequest) ([]FileInfo, error) { + return nil, nil +} + +func (m *mockBackend) Edit(ctx context.Context, _ *EditRequest) error { + return nil +} + +func TestToolResultOffloading_SmallResult(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 100, // Small limit for testing + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a mock endpoint that returns a small result + smallResult := "This is a small result" + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: smallResult}, nil + } + + // Wrap the endpoint with the middleware + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + // Execute + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_123", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Small result should pass through unchanged + if output.Result != smallResult { + t.Errorf("expected result %q, got %q", smallResult, output.Result) + } + + // No file should be written + if len(backend.files) != 0 { + t.Errorf("expected no files to be written, got %d files", len(backend.files)) + } +} + +func TestToolResultOffloading_LargeResult(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, // Very small limit to trigger offloading + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a large result (more than 10 * 4 = 40 bytes) + largeResult := strings.Repeat("This is a long line of text that will exceed the token limit.\n", 10) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_456", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Result should be replaced with a message + if !strings.Contains(output.Result, "Tool result too large") { + t.Errorf("expected result to contain 'Tool result too large', got %q", output.Result) + } + + if !strings.Contains(output.Result, "call_456") { + t.Errorf("expected result to contain call ID 'call_456', got %q", output.Result) + } + + if !strings.Contains(output.Result, "/large_tool_result/call_456") { + t.Errorf("expected result to contain file path, got %q", output.Result) + } + + // File should be written + if len(backend.files) != 1 { + t.Fatalf("expected 1 file to be written, got %d files", len(backend.files)) + } + + savedContent, ok := backend.files["/large_tool_result/call_456"] + if !ok { + t.Fatalf("expected file at /large_tool_result/call_456, got files: %v", backend.files) + } + + if savedContent != largeResult { + t.Errorf("saved content doesn't match original result") + } +} + +func TestToolResultOffloading_CustomPathGenerator(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + customPath := "/custom/path/result.txt" + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + PathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) { + return customPath, nil + }, + } + + middleware := newToolResultOffloading(ctx, config) + + largeResult := strings.Repeat("Large content ", 100) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_789", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Check custom path is used + if !strings.Contains(output.Result, customPath) { + t.Errorf("expected result to contain custom path %q, got %q", customPath, output.Result) + } + + // File should be written to custom path + savedContent, ok := backend.files[customPath] + if !ok { + t.Fatalf("expected file at %q, got files: %v", customPath, backend.files) + } + + if savedContent != largeResult { + t.Errorf("saved content doesn't match original result") + } +} + +func TestToolResultOffloading_PathGeneratorError(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + expectedErr := errors.New("path generation failed") + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + PathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) { + return "", expectedErr + }, + } + + middleware := newToolResultOffloading(ctx, config) + + largeResult := strings.Repeat("Large content ", 100) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } +} + +func TestToolResultOffloading_EndpointError(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 100, + } + + middleware := newToolResultOffloading(ctx, config) + + expectedErr := errors.New("endpoint execution failed") + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return nil, expectedErr + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_endpoint_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } +} + +func TestToolResultOffloading_DefaultTokenLimit(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 0, // Should default to 20000 + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a result smaller than 20000 * 4 = 80000 bytes + smallResult := strings.Repeat("x", 1000) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: smallResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_default", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Should pass through unchanged + if output.Result != smallResult { + t.Errorf("expected result to pass through unchanged") + } + + // No file should be written + if len(backend.files) != 0 { + t.Errorf("expected no files to be written, got %d files", len(backend.files)) + } +} + +func TestToolResultOffloading_Stream(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a streaming endpoint that returns large content + largeResult := strings.Repeat("Large streaming content ", 100) + mockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + // Split the result into chunks + chunks := []string{largeResult[:len(largeResult)/2], largeResult[len(largeResult)/2:]} + return &compose.StreamToolOutput{ + Result: schema.StreamReaderFromArray(chunks), + }, nil + } + + wrappedEndpoint := middleware.Streamable(mockStreamEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_stream", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Read the stream + var result strings.Builder + for { + chunk, err := output.Result.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("error reading stream: %v", err) + } + result.WriteString(chunk) + } + + resultStr := result.String() + + // Result should be replaced with a message + if !strings.Contains(resultStr, "Tool result too large") { + t.Errorf("expected result to contain 'Tool result too large', got %q", resultStr) + } + + if !strings.Contains(resultStr, "call_stream") { + t.Errorf("expected result to contain call ID 'call_stream', got %q", resultStr) + } + + // File should be written + if len(backend.files) != 1 { + t.Fatalf("expected 1 file to be written, got %d files", len(backend.files)) + } + + savedContent, ok := backend.files["/large_tool_result/call_stream"] + if !ok { + t.Fatalf("expected file at /large_tool_result/call_stream, got files: %v", backend.files) + } + + if savedContent != largeResult { + t.Errorf("saved content doesn't match original result") + } +} + +func TestToolResultOffloading_StreamError(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + } + + middleware := newToolResultOffloading(ctx, config) + + expectedErr := errors.New("stream endpoint failed") + mockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return nil, expectedErr + } + + wrappedEndpoint := middleware.Streamable(mockStreamEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_stream_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } +} + +func TestFormatToolMessage(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "single line", + input: "single line", + expected: "1: single line\n", + }, + { + name: "multiple lines", + input: "line1\nline2\nline3", + expected: "1: line1\n2: line2\n3: line3\n", + }, + { + name: "more than 10 lines", + input: "1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12", + expected: "1: 1\n2: 2\n3: 3\n4: 4\n5: 5\n6: 6\n7: 7\n8: 8\n9: 9\n10: 10\n", + }, + { + name: "long line truncation", + input: strings.Repeat("a", 1500), + expected: fmt.Sprintf("1: %s\n", strings.Repeat("a", 1000)), + }, + { + name: "unicode characters", + input: "你好世界\n测试", + expected: "1: 你好世界\n2: 测试\n", + }, + { + name: "long unicode line", + input: strings.Repeat("你", 1500), + expected: fmt.Sprintf("1: %s\n", strings.Repeat("你", 1000)), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := formatToolMessage(tt.input) + if result != tt.expected { + t.Errorf("formatToolMessage() = %q, want %q", result, tt.expected) + } + }) + } +} + +func TestConcatString(t *testing.T) { + tests := []struct { + name string + chunks []string + expected string + expectError bool + }{ + { + name: "single chunk", + chunks: []string{"hello"}, + expected: "hello", + }, + { + name: "multiple chunks", + chunks: []string{"hello", " ", "world"}, + expected: "hello world", + }, + { + name: "empty chunks", + chunks: []string{"", "", ""}, + expected: "", + }, + { + name: "mixed chunks", + chunks: []string{"a", "", "b", "c"}, + expected: "abc", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := schema.StreamReaderFromArray(tt.chunks) + result, err := concatString(sr) + + if tt.expectError { + if err == nil { + t.Error("expected error, got nil") + } + return + } + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if result != tt.expected { + t.Errorf("concatString() = %q, want %q", result, tt.expected) + } + }) + } + + // Test nil stream + t.Run("nil stream", func(t *testing.T) { + _, err := concatString(nil) + if err == nil { + t.Error("expected error for nil stream, got nil") + } + if !strings.Contains(err.Error(), "stream is nil") { + t.Errorf("expected 'stream is nil' error, got %v", err) + } + }) +} + +func TestToolResultOffloading_BackendWriteError(t *testing.T) { + ctx := context.Background() + + // Create a backend that fails on write + backend := &failingBackend{ + writeErr: errors.New("write failed"), + } + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + } + + middleware := newToolResultOffloading(ctx, config) + + largeResult := strings.Repeat("Large content ", 100) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_write_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !strings.Contains(err.Error(), "write failed") { + t.Errorf("expected 'write failed' error, got %v", err) + } +} + +// failingBackend is a mock backend that can be configured to fail +type failingBackend struct { + writeErr error +} + +func (f *failingBackend) Write(ctx context.Context, req *WriteRequest) error { + if f.writeErr != nil { + return f.writeErr + } + return nil +} + +func (f *failingBackend) Read(ctx context.Context, req *ReadRequest) (*FileContent, error) { + return &FileContent{}, nil +} + +func (f *failingBackend) LsInfo(ctx context.Context, _ *LsInfoRequest) ([]FileInfo, error) { + return nil, nil +} + +func (f *failingBackend) GrepRaw(ctx context.Context, _ *GrepRequest) ([]GrepMatch, error) { + return nil, nil +} + +func (f *failingBackend) GlobInfo(ctx context.Context, _ *GlobInfoRequest) ([]FileInfo, error) { + return nil, nil +} + +func (f *failingBackend) Edit(ctx context.Context, _ *EditRequest) error { + return nil +} diff --git a/adk/middlewares/filesystem/prompt.go b/adk/middlewares/filesystem/prompt.go new file mode 100644 index 0000000..a20d6d7 --- /dev/null +++ b/adk/middlewares/filesystem/prompt.go @@ -0,0 +1,270 @@ +/* + * Copyright (c) 2025 Harrison Chase + * Copyright (c) 2025 CloudWeGo Authors + * SPDX-License-Identifier: MIT + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package filesystem + +// This file contains prompt templates and tool descriptions adapted from the DeepAgents project. +// Original source: https://github.com/langchain-ai/deepagents +// +// These prompts are used under the terms of the original project's open source license. +// When using this code in your own open source project, ensure compliance with the original license requirements. + +const ( + tooLargeToolMessage = `Tool result too large, the result of this tool call {tool_call_id} was saved in the filesystem at this path: {file_path} +You can read the result from the filesystem by using the read_file tool, but make sure to only read part of the result at a time. +You can do this by specifying an offset and limit in the read_file tool call. +For example, to read the first 100 lines, you can use the read_file tool with offset=0 and limit=100. + +Here are the first 10 lines of the result: +{content_sample}` + + tooLargeToolMessageChinese = `工具结果过大，此工具调用 {tool_call_id} 的结果已保存到文件系统的以下路径：{file_path} +你可以使用 read_file 工具从文件系统读取结果，但请确保每次只读取部分结果。 +你可以通过在 read_file 工具调用中指定 offset 和 limit 来实现。 +例如，要读取前 100 行，你可以使用 read_file 工具，设置 offset=0 和 limit=100。 + +以下是结果的前 10 行： +{content_sample}` + + ListFilesToolDesc = `Lists all files in the filesystem, filtering by directory. + +Usage: +- The path parameter must be an absolute path, not a relative path +- The ls tool will return a list of all files in the specified directory. +- This is very useful for exploring the file system and finding the right file to read or edit. +- You should almost ALWAYS use this tool before using the read_file or edit_file tools.` + + ListFilesToolDescChinese = `列出文件系统中的所有文件，按目录过滤。 + +使用方法： +- path 参数必须是绝对路径，不能是相对路径 +- ls 工具将返回指定目录中所有文件的列表 +- 这对于探索文件系统和找到要读取或编辑的正确文件非常有用 +- 在使用 read_file 或 edit_file 工具之前，你几乎总是应该先使用此工具` + + ReadFileToolDesc = `Reads a file from the filesystem. You can access any file directly by using this tool. +Assume this tool is able to read all files on the machine. If the User provides a path to a file assume that path is valid. It is okay to read a file that does not exist; an error will be returned. + +Usage: +- The file_path parameter must be an absolute path, not a relative path +- By default, it reads up to 2000 lines starting from the beginning of the file +- **IMPORTANT for large files and codebase exploration**: Use pagination with offset and limit parameters to avoid context overflow + - First scan: read_file(path, limit=100) to see file structure + - Read more sections: read_file(path, offset=100, limit=200) for next 200 lines + - Only omit limit (read full file) when necessary for editing +- Specify offset and limit: read_file(path, offset=0, limit=100) reads first 100 lines +- Results are returned using cat -n format, with line numbers starting at 1 +- You have the capability to call multiple tools in a single response. It is always better to speculatively read multiple files as a batch that are potentially useful. +- If you read a file that exists but has empty contents you will receive a system reminder warning in place of file contents. +- You should ALWAYS make sure a file has been read before editing it.` + + ReadFileToolDescChinese = `从文件系统读取文件。你可以使用此工具直接访问任何文件。 +假设此工具能够读取机器上的所有文件。如果用户提供了文件路径，假设该路径是有效的。读取不存在的文件是可以的；将返回错误。 + +使用方法： +- file_path 参数必须是绝对路径，不能是相对路径 +- 默认情况下，从文件开头读取最多 2000 行 +- **大文件和代码库探索的重要提示**：使用 offset 和 limit 参数进行分页，以避免上下文溢出 + - 首次扫描：read_file(path, limit=100) 查看文件结构 + - 读取更多部分：read_file(path, offset=100, limit=200) 读取接下来的 200 行 + - 仅在编辑必要时才省略 limit（读取完整文件） +- 指定 offset 和 limit：read_file(path, offset=0, limit=100) 读取前 100 行 +- 结果以 cat -n 格式返回，行号从 1 开始 +- 你可以在单个响应中调用多个工具。最好同时推测性地批量读取多个可能有用的文件 +- 如果你读取的文件存在但内容为空，你将收到系统提醒警告而不是文件内容 +- 在编辑文件之前，你应该始终确保已读取该文件` + + // EnhancedReadFileDescSuffix is appended to ReadFileToolDesc when using MultiModalReadFileTool. + EnhancedReadFileDescSuffix = ` +- This tool supports reading image files (e.g., PNG, JPG, etc.). When reading an image file, the contents are presented visually, as the underlying model is a multimodal LLM. +- This tool can read PDF files (.pdf). For large PDFs (more than 10 pages), you MUST provide the pages parameter to read specific page ranges (e.g., pages: "1-5"). Reading a large PDF without the pages parameter will fail. Maximum 20 pages per request.` + + EnhancedReadFileDescSuffixChinese = ` +- 此工具支持读取图片文件（如 PNG、JPG 等）。读取图片文件时，内容将以视觉方式呈现，因为底层模型是多模态 LLM。 +- 此工具可以读取 PDF 文件（.pdf）。对于大型 PDF（超过 10 页），你必须提供 pages 参数来指定页面范围（例如 pages: "1-5"）。不提供 pages 参数读取大型 PDF 将会失败。每次请求最多 20 页。` + + EditFileToolDesc = `Performs exact string replacements in files. + +Usage: +- You must use your 'read_file' tool at least once in the conversation before editing. This tool will error if you attempt an edit without reading the file. +- When editing text from Read tool output, ensure you preserve the exact indentation (tabs/spaces) as it appears AFTER the line number prefix. The line number prefix format is: spaces + line number + tab. Everything after that tab is the actual file content to match. Never include any part of the line number prefix in the old_string or new_string. +- ALWAYS prefer editing existing files. NEVER write new files unless explicitly required. +- Only use emojis if the user explicitly requests it. Avoid adding emojis to files unless asked. +- The edit will FAIL if 'old_string' is not unique in the file. Either provide a larger string with more surrounding context to make it unique or use 'replace_all' to change every instance of 'old_string'. +- Use 'replace_all' for replacing and renaming strings across the file. This parameter is useful if you want to rename a variable for instance.` + + EditFileToolDescChinese = `在文件中执行精确的字符串替换。 + +使用方法： +- 在编辑之前，你必须在对话中至少使用一次 'read_file' 工具。如果你在未读取文件的情况下尝试编辑，此工具将报错 +- 当从 Read 工具输出编辑文本时，请确保保留行号前缀之后的确切缩进（制表符/空格）。行号前缀格式为：空格 + 行号 + 制表符。制表符之后的所有内容都是要匹配的实际文件内容。永远不要在 old_string 或 new_string 中包含行号前缀的任何部分 +- 始终优先编辑现有文件。除非明确要求，否则不要创建新文件 +- 仅在用户明确要求时使用表情符号。除非被要求，否则避免在文件中添加表情符号 +- 如果 'old_string' 在文件中不唯一，编辑将失败。要么提供包含更多上下文的更长字符串使其唯一，要么使用 'replace_all' 更改 'old_string' 的每个实例 +- 使用 'replace_all' 在整个文件中替换和重命名字符串。例如，如果你想重命名变量，此参数很有用` + + WriteFileToolDesc = `Writes a file to the local filesystem. + +Usage: +- This tool will overwrite the existing file if there is one at the provided path. +- If this is an existing file, you MUST use the Read tool first to read the file's contents. This tool will fail if you did not read the file first. +- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required. +- NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User. +- Only use emojis if the user explicitly requests it. Avoid writing emojis to files unless asked.` + + WriteFileToolDescChinese = `将文件写入本地文件系统。 + +使用方法： +- 如果提供的路径已存在文件，此工具将覆盖现有文件 +- 如果这是一个现有文件，你必须先使用 Read 工具读取文件内容。如果你没有先读取文件，此工具将失败 +- 始终优先编辑代码库中的现有文件。除非明确要求，否则不要创建新文件 +- 不要主动创建文档文件（*.md）或 README 文件。仅在用户明确要求时才创建文档文件 +- 仅在用户明确要求时使用表情符号。除非被要求，否则避免在文件中写入表情符号` + + GlobToolDesc = `Fast file pattern matching tool that works with any codebase size +- Supports glob patterns like "**/*.js" or "src/**/*.ts" +- Returns matching file paths sorted by modification time +- Use this tool when you need to find files by name patterns +- You can call multiple tools in a single response. It is always better to speculatively perform multiple searches in parallel if they are potentially useful. + +Examples: +- '**/*.py' - Find all Python files +- '*.txt' - Find all text files in root +- '/subdir/**/*.md' - Find all markdown files under /subdir` + + GlobToolDescChinese = `适用于任何代码库大小的快速文件模式匹配工具 +- 支持 glob 模式，如 "**/*.js" 或 "src/**/*.ts" +- 返回按修改时间排序的匹配文件路径 +- 当你需要按名称模式查找文件时使用此工具 +- 你可以在单个响应中调用多个工具。最好同时并行执行多个可能有用的搜索 + +示例： +- '**/*.py' - 查找所有 Python 文件 +- '*.txt' - 查找根目录中的所有文本文件 +- '/subdir/**/*.md' - 查找 /subdir 下的所有 markdown 文件` + + GrepToolDesc = ` +A powerful search tool built on ripgrep + + Usage: + - ALWAYS use Grep for search tasks. NEVER invoke 'grep' or 'rg' as a Bash command. The Grep tool has been optimized for correct permissions and access. + - Supports full regex syntax (e.g., "log.*Error", "function\s+\w+") + - Filter files with glob parameter (e.g., "*.js", "**/*.tsx") or type parameter (e.g., "js", "py", "rust") + - Output modes: "content" shows matching lines, "files_with_matches" shows only file paths (default), "count" shows match counts + - Use Task tool for open-ended searches requiring multiple rounds + - Pattern syntax: Uses ripgrep (not grep) - literal braces need escaping (use 'interface\{\}' to find 'interface{}' in Go code) + - Multiline matching: By default patterns match within single lines only. For cross-line patterns like 'struct \{[\s\S]*?field', use 'multiline: true'` + + GrepToolDescChinese = ` +基于 ripgrep 的强大搜索工具 + + 使用方法： + - 始终使用 Grep 进行搜索任务。不要将 'grep' 或 'rg' 作为 Bash 命令调用。Grep 工具已针对正确的权限和访问进行了优化 + - 支持完整的正则表达式语法（例如，"log.*Error"，"function\s+\w+"） + - 使用 glob 参数（例如，"*.js"，"**/*.tsx"）或 type 参数（例如，"js"，"py"，"rust"）过滤文件 + - 输出模式："content" 显示匹配行，"files_with_matches" 仅显示文件路径（默认），"count" 显示匹配计数 + - 对于需要多轮的开放式搜索，使用 Task 工具 + - 模式语法：使用 ripgrep（不是 grep）- 字面大括号需要转义（使用 'interface\{\}' 在 Go 代码中查找 'interface{}'） + - 多行匹配：默认情况下，模式仅在单行内匹配。对于跨行模式如 'struct \{[\s\S]*?field'，使用 'multiline: true'` + + ExecuteToolDesc = ` +Executes a given command in the sandbox environment with proper handling and security measures. + +Before executing the command, please follow these steps: + +1. Directory Verification: +- If the command will create new directories or files, first use the ls tool to verify the parent directory exists and is the correct location +- For example, before running "mkdir foo/bar", first use ls to check that "foo" exists and is the intended parent directory + +2. Command Execution: +- Always quote file paths that contain spaces with double quotes (e.g., cd "path with spaces/file.txt") +- Examples of proper quoting: +- cd "/Users/name/My Documents" (correct) +- cd /Users/name/My Documents (incorrect - will fail) +- python "/path/with spaces/script.py" (correct) +- python /path/with spaces/script.py (incorrect - will fail) +- After ensuring proper quoting, execute the command +- Capture the output of the command + +Usage notes: +- The command parameter is required +- Commands run in an isolated sandbox environment +- Returns combined stdout/stderr output with exit code +- If the output is very large, it may be truncated +- VERY IMPORTANT: You MUST avoid using search commands like find and grep. Instead use the grep, glob tools to search. You MUST avoid read tools like cat, head, tail, and use read_file to read files. +- When issuing multiple commands, use the ';' or '&&' operator to separate them. DO NOT use newlines (newlines are ok in quoted strings) +- Use '&&' when commands depend on each other (e.g., "mkdir dir && cd dir") +- Use ';' only when you need to run commands sequentially but don't care if earlier commands fail +- Try to maintain your current working directory throughout the session by using absolute paths and avoiding usage of cd + +Examples: +Good examples: +- execute(command="pytest /foo/bar/tests") +- execute(command="python /path/to/script.py") +- execute(command="npm install && npm test") + +Bad examples (avoid these): +- execute(command="cd /foo/bar && pytest tests") # Use absolute path instead +- execute(command="cat file.txt") # Use read_file tool instead +- execute(command="find . -name '*.py'") # Use glob tool instead +- execute(command="grep -r 'pattern' .") # Use grep tool instead +` + + ExecuteToolDescChinese = ` +在沙箱环境中执行给定命令，具有适当的处理和安全措施。 + +执行命令前，请按照以下步骤操作： + +1. 目录验证： +- 如果命令将创建新目录或文件，首先使用 ls 工具验证父目录是否存在且是正确的位置 +- 例如，在运行 "mkdir foo/bar" 之前，首先使用 ls 检查 "foo" 是否存在且是预期的父目录 + +2. 命令执行： +- 始终用双引号引用包含空格的文件路径（例如，cd "path with spaces/file.txt"） +- 正确引用的示例： +- cd "/Users/name/My Documents"（正确） +- cd /Users/name/My Documents（错误 - 将失败） +- python "/path/with spaces/script.py"（正确） +- python /path/with spaces/script.py（错误 - 将失败） +- 确保正确引用后，执行命令 +- 捕获命令的输出 + +使用说明： +- command 参数是必需的 +- 命令在隔离的沙箱环境中运行 +- 返回合并的 stdout/stderr 输出和退出代码 +- 如果输出非常大，可能会被截断 +- 非常重要：你必须避免使用 find 和 grep 等搜索命令。请改用 grep、glob 工具进行搜索。你必须避免使用 cat、head、tail 等读取工具，请使用 read_file 读取文件 +- 发出多个命令时，使用 ';' 或 '&&' 运算符分隔它们。不要使用换行符（引号字符串中的换行符是可以的） +- 当命令相互依赖时使用 '&&'（例如，"mkdir dir && cd dir"） +- 仅当你需要按顺序运行命令但不关心早期命令是否失败时使用 ';' +- 尝试通过使用绝对路径并避免使用 cd 来在整个会话中保持当前工作目录 + +示例： +好的示例： +- execute(command="pytest /foo/bar/tests") +- execute(command="python /path/to/script.py") +- execute(command="npm install && npm test") + +不好的示例（避免这些）： +- execute(command="cd /foo/bar && pytest tests") # 改用绝对路径 +- execute(command="cat file.txt") # 改用 read_file 工具 +- execute(command="find . -name '*.py'") # 改用 glob 工具 +- execute(command="grep -r 'pattern' .") # 改用 grep 工具 +` +) diff --git a/adk/middlewares/patchtoolcalls/patchtoolcalls.go b/adk/middlewares/patchtoolcalls/patchtoolcalls.go new file mode 100644 index 0000000..484c881 --- /dev/null +++ b/adk/middlewares/patchtoolcalls/patchtoolcalls.go @@ -0,0 +1,263 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package patchtoolcalls provides a middleware that patches dangling tool calls in the message history. +package patchtoolcalls + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/schema" +) + +// Config defines the configuration options for the patch tool calls middleware. +type Config struct { + // PatchedContentGenerator is an optional custom function to generate the content + // of patched tool messages. If not provided, a default message will be used. + // + // Parameters: + // - ctx: the context for the operation + // - toolName: the name of the tool that was called + // - toolCallID: the id of the tool call + // + // Returns: + // - string: the content to use for the patched tool message + // - error: any error that occurred during generation + PatchedContentGenerator func(ctx context.Context, toolName, toolCallID string) (string, error) +} + +// NewTyped creates a new generic patch tool calls middleware. +// +// The middleware scans the message history before each model invocation and inserts +// placeholder tool messages for any tool calls that don't have corresponding responses. +func NewTyped[M adk.MessageType](_ context.Context, cfg *Config) (adk.TypedChatModelAgentMiddleware[M], error) { + if cfg == nil { + cfg = &Config{} + } + return &typedMiddleware[M]{ + gen: cfg.PatchedContentGenerator, + }, nil +} + +// New creates a new patch tool calls middleware with the given configuration. +// +// The middleware scans the message history before each model invocation and inserts +// placeholder tool messages for any tool calls that don't have corresponding responses. +func New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, cfg) +} + +type typedMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + gen func(ctx context.Context, toolName, toolCallID string) (string, error) +} + +func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], + mc *adk.TypedModelContext[M], +) (context.Context, *adk.TypedChatModelAgentState[M], error) { + if len(state.Messages) == 0 { + return ctx, state, nil + } + + var zero M + switch any(zero).(type) { + case *schema.Message: + return patchToolCallsForMessage(ctx, m.gen, any(state).(*adk.TypedChatModelAgentState[*schema.Message]), mc) + case *schema.AgenticMessage: + return patchToolCallsForAgenticMessage(ctx, m.gen, any(state).(*adk.TypedChatModelAgentState[*schema.AgenticMessage]), mc) + default: + panic("unreachable: unknown MessageType") + } +} + +func patchToolCallsForMessage[M adk.MessageType](ctx context.Context, + gen func(ctx context.Context, toolName, toolCallID string) (string, error), + state *adk.TypedChatModelAgentState[*schema.Message], + _ *adk.TypedModelContext[M], +) (context.Context, *adk.TypedChatModelAgentState[M], error) { + patched := make([]*schema.Message, 0, len(state.Messages)) + + for i, msg := range state.Messages { + patched = append(patched, msg) + + if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { + continue + } + + for _, tc := range msg.ToolCalls { + if hasCorrespondingToolMessage(state.Messages[i+1:], tc.ID) { + continue + } + + toolMsg, err := createPatchedToolMessage(ctx, gen, tc) + if err != nil { + return ctx, nil, err + } + patched = append(patched, toolMsg) + } + } + + nState := *state + nState.Messages = patched + return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil +} + +func patchToolCallsForAgenticMessage[M adk.MessageType](ctx context.Context, + gen func(ctx context.Context, toolName, toolCallID string) (string, error), + state *adk.TypedChatModelAgentState[*schema.AgenticMessage], + _ *adk.TypedModelContext[M], +) (context.Context, *adk.TypedChatModelAgentState[M], error) { + patched := make([]*schema.AgenticMessage, 0, len(state.Messages)) + + for i, msg := range state.Messages { + patched = append(patched, msg) + + if msg.Role != schema.AgenticRoleTypeAssistant { + continue + } + + // Collect tool call IDs from this assistant message. + var toolCalls []struct { + callID string + name string + } + for _, block := range msg.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil { + toolCalls = append(toolCalls, struct { + callID string + name string + }{callID: block.FunctionToolCall.CallID, name: block.FunctionToolCall.Name}) + } + } + if len(toolCalls) == 0 { + continue + } + + for _, tc := range toolCalls { + if hasCorrespondingAgenticToolResult(state.Messages[i+1:], tc.callID) { + continue + } + + toolMsg, err := createPatchedAgenticToolMessage(ctx, gen, tc.name, tc.callID) + if err != nil { + return ctx, nil, err + } + patched = append(patched, toolMsg) + } + } + + nState := *state + nState.Messages = patched + return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil +} + +func hasCorrespondingToolMessage(messages []*schema.Message, toolCallID string) bool { + for _, msg := range messages { + // Only consider successive tool messages after the tool call message + if msg.Role != schema.Tool { + return false + } + if msg.ToolCallID == toolCallID { + return true + } + } + return false +} + +func hasCorrespondingAgenticToolResult(messages []*schema.AgenticMessage, toolCallID string) bool { + for _, msg := range messages { + // Only consider successive tool messages after the tool call message + if msg.Role != schema.AgenticRoleTypeUser { + return false + } + hasToolResult := false + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + if block.Type == schema.ContentBlockTypeFunctionToolResult { + hasToolResult = true + if block.FunctionToolResult != nil && block.FunctionToolResult.CallID == toolCallID { + return true + } + } + if block.Type == schema.ContentBlockTypeToolSearchResult { + hasToolResult = true + if block.ToolSearchFunctionToolResult != nil && block.ToolSearchFunctionToolResult.CallID == toolCallID { + return true + } + } + } + if !hasToolResult { + return false + } + } + return false +} + +func createPatchedToolMessage(ctx context.Context, gen func(ctx context.Context, toolName, toolCallID string) (string, error), tc schema.ToolCall) (*schema.Message, error) { + if gen != nil { + content, err := gen(ctx, tc.Function.Name, tc.ID) + if err != nil { + return nil, err + } + return schema.ToolMessage(content, tc.ID, schema.WithToolName(tc.Function.Name)), nil + } + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: defaultPatchedToolMessageTemplate, + Chinese: defaultPatchedToolMessageTemplateChinese, + }) + + return schema.ToolMessage(fmt.Sprintf(tpl, tc.Function.Name, tc.ID), tc.ID, schema.WithToolName(tc.Function.Name)), nil +} + +func createPatchedAgenticToolMessage(ctx context.Context, gen func(ctx context.Context, toolName, toolCallID string) (string, error), toolName, callID string) (*schema.AgenticMessage, error) { + var content string + if gen != nil { + var err error + content, err = gen(ctx, toolName, callID) + if err != nil { + return nil, err + } + } else { + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: defaultPatchedToolMessageTemplate, + Chinese: defaultPatchedToolMessageTemplateChinese, + }) + content = fmt.Sprintf(tpl, toolName, callID) + } + + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: callID, + Name: toolName, + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: content}}, + }, + }), + }, + }, nil +} + +const ( + defaultPatchedToolMessageTemplate = "Tool call %s with id %s was canceled - another message came in before it could be completed." + defaultPatchedToolMessageTemplateChinese = "工具调用 %s（ID 为 %s）已被取消——在其完成之前收到了另一条消息。" +) diff --git a/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go b/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go new file mode 100644 index 0000000..2fdb3c1 --- /dev/null +++ b/adk/middlewares/patchtoolcalls/patchtoolcalls_test.go @@ -0,0 +1,394 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package patchtoolcalls + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.AgenticMessage](ctx, nil) + assert.NoError(t, err) + assert.NotNil(t, mw) + + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} + +type testToolCall struct { + ID string + Name string + Arguments string +} + +func makeUserMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(content)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(content)).(M) + } + panic("unreachable") +} + +func makeAssistantMsgWithToolCalls[M adk.MessageType](content string, toolCalls []testToolCall) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + tcs := make([]schema.ToolCall, len(toolCalls)) + for i, tc := range toolCalls { + tcs[i] = schema.ToolCall{ID: tc.ID, Function: schema.FunctionCall{Name: tc.Name, Arguments: tc.Arguments}} + } + return any(schema.AssistantMessage(content, tcs)).(M) + case *schema.AgenticMessage: + blocks := make([]*schema.ContentBlock, 0, len(toolCalls)+1) + if content != "" { + blocks = append(blocks, schema.NewContentBlock(&schema.AssistantGenText{Text: content})) + } + for _, tc := range toolCalls { + blocks = append(blocks, schema.NewContentBlock(&schema.FunctionToolCall{CallID: tc.ID, Name: tc.Name, Arguments: tc.Arguments})) + } + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: blocks, + }).(M) + } + panic("unreachable") +} + +func makeToolResultMsg[M adk.MessageType](content string, callID string, toolName string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.ToolMessage(content, callID, schema.WithToolName(toolName))).(M) + case *schema.AgenticMessage: + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: callID, + Name: toolName, + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: content}}, + }, + }), + }, + }).(M) + } + panic("unreachable") +} + +func assertMsgContent[M adk.MessageType](t *testing.T, msg M, expectedContent string) { + t.Helper() + switch m := any(msg).(type) { + case *schema.Message: + assert.Equal(t, expectedContent, m.Content) + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + for _, b := range block.FunctionToolResult.Content { + if b.Text != nil { + assert.Equal(t, expectedContent, b.Text.Text) + return + } + } + } + } + t.Errorf("no text content found in agentic message, expected %q", expectedContent) + } +} + +func assertToolResultID[M adk.MessageType](t *testing.T, msg M, expectedID string) { + t.Helper() + switch m := any(msg).(type) { + case *schema.Message: + assert.Equal(t, expectedID, m.ToolCallID) + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + assert.Equal(t, expectedID, block.FunctionToolResult.CallID) + return + } + } + t.Errorf("no tool result found in agentic message, expected call ID %q", expectedID) + } +} + +func assertToolResultName[M adk.MessageType](t *testing.T, msg M, expectedName string) { + t.Helper() + switch m := any(msg).(type) { + case *schema.Message: + assert.Equal(t, expectedName, m.ToolName) + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + assert.Equal(t, expectedName, block.FunctionToolResult.Name) + return + } + } + t.Errorf("no tool result found in agentic message, expected tool name %q", expectedName) + } +} + +func testPatchToolCallsGeneric[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + tests := []struct { + name string + config *Config + messages []M + wantLen int + checkPatchedAt int // index of the patched message to check (-1 if no check needed) + wantCallID string + wantToolName string + wantContent string + }{ + { + name: "empty messages", + config: nil, + messages: nil, + wantLen: 0, + checkPatchedAt: -1, + }, + { + name: "no tool calls to patch", + config: nil, + messages: []M{ + makeUserMsg[M]("hello"), + makeAssistantMsgWithToolCalls[M]("hi there", nil), + }, + wantLen: 2, + checkPatchedAt: -1, + }, + { + name: "missing tool result", + config: nil, + messages: []M{ + makeUserMsg[M]("hello"), + makeAssistantMsgWithToolCalls[M]("", []testToolCall{ + {ID: "call_1", Name: "tool_a", Arguments: "{}"}, + {ID: "call_2", Name: "tool_b", Arguments: "{}"}, + }), + makeToolResultMsg[M]("result_a", "call_1", "tool_a"), + }, + wantLen: 4, + checkPatchedAt: 2, + wantCallID: "call_2", + wantToolName: "tool_b", + wantContent: fmt.Sprintf(defaultPatchedToolMessageTemplate, "tool_b", "call_2"), + }, + { + name: "custom content generator", + config: &Config{ + PatchedContentGenerator: func(ctx context.Context, toolName, toolCallID string) (string, error) { + return fmt.Sprintf("123 %s %s", toolName, toolCallID), nil + }, + }, + messages: []M{ + makeUserMsg[M]("hello"), + makeAssistantMsgWithToolCalls[M]("", []testToolCall{ + {ID: "call_1", Name: "tool_a", Arguments: "{}"}, + {ID: "call_2", Name: "tool_b", Arguments: "{}"}, + }), + makeToolResultMsg[M]("result_a", "call_1", "tool_a"), + }, + wantLen: 4, + checkPatchedAt: 2, + wantCallID: "call_2", + wantToolName: "tool_b", + wantContent: "123 tool_b call_2", + }, + { + name: "two consecutive assistant messages with tool calls", + config: nil, + messages: []M{ + makeUserMsg[M]("hello"), + makeAssistantMsgWithToolCalls[M]("", []testToolCall{ + {ID: "call_1", Name: "tool_a", Arguments: "{}"}, + }), + makeAssistantMsgWithToolCalls[M]("continued...", nil), + }, + wantLen: 4, + checkPatchedAt: 2, + wantCallID: "call_1", + wantToolName: "tool_a", + wantContent: fmt.Sprintf(defaultPatchedToolMessageTemplate, "tool_a", "call_1"), + }, + { + name: "assistant message followed by user message without tool result", + config: nil, + messages: []M{ + makeUserMsg[M]("hello"), + makeAssistantMsgWithToolCalls[M]("", []testToolCall{ + {ID: "call_1", Name: "tool_a", Arguments: "{}"}, + }), + makeUserMsg[M]("continued..."), + }, + wantLen: 4, + checkPatchedAt: 2, + wantCallID: "call_1", + wantToolName: "tool_a", + wantContent: fmt.Sprintf(defaultPatchedToolMessageTemplate, "tool_a", "call_1"), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mw, err := NewTyped[M](ctx, tt.config) + assert.NoError(t, err) + + state := &adk.TypedChatModelAgentState[M]{ + Messages: tt.messages, + } + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + assert.NoError(t, err) + assert.Len(t, newState.Messages, tt.wantLen) + + if tt.checkPatchedAt >= 0 && tt.checkPatchedAt < len(newState.Messages) { + patched := newState.Messages[tt.checkPatchedAt] + assertToolResultID(t, patched, tt.wantCallID) + assertToolResultName(t, patched, tt.wantToolName) + assertMsgContent(t, patched, tt.wantContent) + } + }) + } +} + +func TestPatchToolCallsGeneric(t *testing.T) { + t.Run("Message", testPatchToolCallsGeneric[*schema.Message]) + t.Run("AgenticMessage", testPatchToolCallsGeneric[*schema.AgenticMessage]) +} + +func TestPatchToolCallsAgenticToolSearchResult(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.AgenticMessage](ctx, nil) + require.NoError(t, err) + + messages := []*schema.AgenticMessage{ + makeAssistantMsgWithToolCalls[*schema.AgenticMessage]("", []testToolCall{ + {ID: "call_1", Name: "tool_search", Arguments: `{"query":"dynamic"}`}, + }), + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.ToolSearchFunctionToolResult{ + CallID: "call_1", + Name: "tool_search", + Result: &schema.ToolSearchResult{Tools: []*schema.ToolInfo{ + {Name: "dynamic_tool", Desc: "dynamic tool"}, + }}, + }), + }, + }, + } + + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: messages} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + require.NoError(t, err) + assert.Len(t, newState.Messages, 2) + assert.Equal(t, schema.ContentBlockTypeToolSearchResult, newState.Messages[1].ContentBlocks[0].Type) +} + +// TestPatchToolCalls_NilFunctionToolCallInBlock verifies the middleware handles +// a ContentBlock with Type=FunctionToolCall but FunctionToolCall=nil without panicking. +func TestPatchToolCalls_NilFunctionToolCallInBlock(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.AgenticMessage](ctx, nil) + require.NoError(t, err) + + msgs := []*schema.AgenticMessage{ + schema.UserAgenticMessage("hello"), + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: nil, // nil despite type indicating tool call + }, + schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: "call_1", + Name: "real_tool", + }), + }, + }, + } + + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: msgs} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 3, "should patch call_1 but skip nil FunctionToolCall block") + + patchMsg := newState.Messages[2] + assert.Equal(t, schema.AgenticRoleTypeUser, patchMsg.Role) + foundResult := false + for _, block := range patchMsg.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult && + block.FunctionToolResult != nil && block.FunctionToolResult.CallID == "call_1" { + foundResult = true + } + } + assert.True(t, foundResult, "patched message should contain tool result for call_1") +} + +// TestPatchToolCalls_AgenticMessage_NilBlockInUserMessage verifies the middleware handles +// a User Agentic Message with nil ContentBlock without panicking. +func TestPatchToolCalls_AgenticMessage_NilBlockInUserMessage(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.AgenticMessage](ctx, nil) + require.NoError(t, err) + + msgs := []*schema.AgenticMessage{ + schema.UserAgenticMessage("hello"), + makeAssistantMsgWithToolCalls[*schema.AgenticMessage]("", []testToolCall{ + {ID: "call_1", Name: "tool_a", Arguments: "{}"}, + }), + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + nil, // nil block to test robustness + }, + }, + } + + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: msgs} + _, newState, err := mw.BeforeModelRewriteState(ctx, state, nil) + assert.NoError(t, err, "should not panic when encountering nil block in user message") + assert.Len(t, newState.Messages, 4, "should patch call_1 and insert tool response") + + // Verify the patched message is inserted at index 2 + patchMsg := newState.Messages[2] + assert.Equal(t, schema.AgenticRoleTypeUser, patchMsg.Role) + + foundResult := false + for _, block := range patchMsg.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult && + block.FunctionToolResult != nil && block.FunctionToolResult.CallID == "call_1" { + foundResult = true + break + } + } + assert.True(t, foundResult, "patched message should contain tool result for call_1") +} diff --git a/adk/middlewares/plantask/backend_test.go b/adk/middlewares/plantask/backend_test.go new file mode 100644 index 0000000..d381ff7 --- /dev/null +++ b/adk/middlewares/plantask/backend_test.go @@ -0,0 +1,80 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "errors" + "path/filepath" + "strings" + "sync" + + fspkg "github.com/cloudwego/eino/adk/filesystem" +) + +type inMemoryBackend struct { + files map[string]string + mu sync.RWMutex +} + +func newInMemoryBackend() *inMemoryBackend { + return &inMemoryBackend{ + files: make(map[string]string), + } +} + +func (b *inMemoryBackend) LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + reqPath := strings.TrimSuffix(req.Path, "/") + var result []FileInfo + for path := range b.files { + dir := filepath.Dir(path) + if dir == reqPath { + result = append(result, FileInfo{Path: path}) + } + } + return result, nil +} + +func (b *inMemoryBackend) Read(ctx context.Context, req *ReadRequest) (*fspkg.FileContent, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + content, ok := b.files[req.FilePath] + if !ok { + return nil, errors.New("file not found") + } + return &fspkg.FileContent{Content: content}, nil +} + +func (b *inMemoryBackend) Write(ctx context.Context, req *WriteRequest) error { + b.mu.Lock() + defer b.mu.Unlock() + + b.files[req.FilePath] = req.Content + return nil +} + +func (b *inMemoryBackend) Delete(ctx context.Context, req *DeleteRequest) error { + b.mu.Lock() + defer b.mu.Unlock() + + delete(b.files, req.FilePath) + return nil +} diff --git a/adk/middlewares/plantask/plantask.go b/adk/middlewares/plantask/plantask.go new file mode 100644 index 0000000..fb201bd --- /dev/null +++ b/adk/middlewares/plantask/plantask.go @@ -0,0 +1,81 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "fmt" + "sync" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// Config is the configuration for the tool search middleware. +type Config struct { + Backend Backend + BaseDir string +} + +// NewTyped creates a new plantask middleware that provides task management tools for agents. +// It adds TaskCreate, TaskGet, TaskUpdate, and TaskList tools to the agent's tool set, +// allowing agents to create and manage structured task lists during coding sessions. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +func NewTyped[M adk.MessageType](_ context.Context, config *Config) (adk.TypedChatModelAgentMiddleware[M], error) { + if config == nil { + return nil, fmt.Errorf("config is required") + } + if config.Backend == nil { + return nil, fmt.Errorf("backend is required") + } + if config.BaseDir == "" { + return nil, fmt.Errorf("baseDir is required") + } + + return &typedMiddleware[M]{backend: config.Backend, baseDir: config.BaseDir}, nil +} + +// New creates a new plantask middleware that provides task management tools for agents. +// It adds TaskCreate, TaskGet, TaskUpdate, and TaskList tools to the agent's tool set, +// allowing agents to create and manage structured task lists during coding sessions. +func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped[*schema.Message](ctx, config) +} + +type typedMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + backend Backend + baseDir string +} + +func (m *typedMiddleware[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + if runCtx == nil { + return ctx, runCtx, nil + } + + nRunCtx := *runCtx + lock := sync.Mutex{} + nRunCtx.Tools = append(nRunCtx.Tools, + newTaskCreateTool(m.backend, m.baseDir, &lock), + newTaskGetTool(m.backend, m.baseDir, &lock), + newTaskUpdateTool(m.backend, m.baseDir, &lock), + newTaskListTool(m.backend, m.baseDir, &lock), + ) + + return ctx, &nRunCtx, nil +} diff --git a/adk/middlewares/plantask/plantask_test.go b/adk/middlewares/plantask/plantask_test.go new file mode 100644 index 0000000..2354e79 --- /dev/null +++ b/adk/middlewares/plantask/plantask_test.go @@ -0,0 +1,137 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func TestNew(t *testing.T) { + ctx := context.Background() + + _, err := New(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config is required") + + _, err = New(ctx, &Config{}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend is required") + + _, err = New(ctx, &Config{Backend: newInMemoryBackend()}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "baseDir is required") + + m, err := New(ctx, &Config{Backend: newInMemoryBackend(), BaseDir: "/tmp/tasks"}) + assert.NoError(t, err) + assert.NotNil(t, m) +} + +func TestMiddlewareBeforeAgent(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + + m, err := New(ctx, &Config{Backend: backend, BaseDir: baseDir}) + assert.NoError(t, err) + + mw := m.(*typedMiddleware[*schema.Message]) + + ctx, runCtx, err := mw.BeforeAgent(ctx, nil) + assert.NoError(t, err) + assert.Nil(t, runCtx) + + runCtx = &adk.ChatModelAgentContext{ + Tools: []tool.BaseTool{}, + } + ctx, newRunCtx, err := mw.BeforeAgent(ctx, runCtx) + assert.NoError(t, err) + assert.NotNil(t, newRunCtx) + assert.Len(t, newRunCtx.Tools, 4) + + toolNames := make([]string, 0, 4) + for _, t := range newRunCtx.Tools { + info, _ := t.Info(ctx) + toolNames = append(toolNames, info.Name) + } + assert.Contains(t, toolNames, "TaskCreate") + assert.Contains(t, toolNames, "TaskGet") + assert.Contains(t, toolNames, "TaskUpdate") + assert.Contains(t, toolNames, "TaskList") +} + +func TestIntegration(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + createTool := newTaskCreateTool(backend, baseDir, lock) + getTool := newTaskGetTool(backend, baseDir, lock) + updateTool := newTaskUpdateTool(backend, baseDir, lock) + listTool := newTaskListTool(backend, baseDir, lock) + + result, err := createTool.InvokableRun(ctx, `{"subject": "Task 1", "description": "First task"}`) + assert.NoError(t, err) + assert.Contains(t, result, "Task #1") + + result, err = createTool.InvokableRun(ctx, `{"subject": "Task 2", "description": "Second task"}`) + assert.NoError(t, err) + assert.Contains(t, result, "Task #2") + + _, err = updateTool.InvokableRun(ctx, `{"taskId": "2", "addBlockedBy": ["1"]}`) + assert.NoError(t, err) + + result, err = listTool.InvokableRun(ctx, `{}`) + assert.NoError(t, err) + assert.Contains(t, result, "#1 [pending] Task 1") + assert.Contains(t, result, "#2 [pending] Task 2") + assert.Contains(t, result, "[blocked by #1]") + + _, err = updateTool.InvokableRun(ctx, `{"taskId": "1", "status": "in_progress"}`) + assert.NoError(t, err) + + result, err = getTool.InvokableRun(ctx, `{"taskId": "1"}`) + assert.NoError(t, err) + assert.Contains(t, result, "Status: in_progress") + + _, err = updateTool.InvokableRun(ctx, `{"taskId": "1", "status": "completed"}`) + assert.NoError(t, err) + + result, err = listTool.InvokableRun(ctx, `{}`) + assert.NoError(t, err) + assert.Contains(t, result, "#1 [completed] Task 1") +} + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped[*schema.AgenticMessage](ctx, &Config{ + Backend: newInMemoryBackend(), + BaseDir: "/tmp/tasks", + }) + assert.NoError(t, err) + assert.NotNil(t, mw) + + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} diff --git a/adk/middlewares/plantask/task.go b/adk/middlewares/plantask/task.go new file mode 100644 index 0000000..ff1ed28 --- /dev/null +++ b/adk/middlewares/plantask/task.go @@ -0,0 +1,123 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "regexp" + + "github.com/cloudwego/eino/adk/middlewares/filesystem" +) + +var validTaskIDRegex = regexp.MustCompile(`^\d+$`) + +const highWatermarkFileName = ".highwatermark" + +type task struct { + ID string `json:"id"` + Subject string `json:"subject"` + Description string `json:"description"` + Status string `json:"status"` + Blocks []string `json:"blocks"` + BlockedBy []string `json:"blockedBy"` + ActiveForm string `json:"activeForm,omitempty"` + Owner string `json:"owner,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +type taskOut struct { + Result string `json:"result"` +} + +const ( + taskStatusPending = "pending" + taskStatusInProgress = "in_progress" + taskStatusCompleted = "completed" + taskStatusDeleted = "deleted" +) + +type FileInfo = filesystem.FileInfo +type LsInfoRequest = filesystem.LsInfoRequest +type ReadRequest = filesystem.ReadRequest +type WriteRequest = filesystem.WriteRequest + +type DeleteRequest struct { + FilePath string +} + +// Backend defines the storage interface for task persistence. +// Implementations can use local filesystem, remote storage, or any other storage backend. +type Backend interface { + // LsInfo lists file information in the specified directory. + LsInfo(ctx context.Context, req *LsInfoRequest) ([]FileInfo, error) + // Read reads the content of a file. + Read(ctx context.Context, req *ReadRequest) (*filesystem.FileContent, error) + // Write writes content to a file, creating it if it doesn't exist. + Write(ctx context.Context, req *WriteRequest) error + // Delete removes a file from storage. + Delete(ctx context.Context, req *DeleteRequest) error +} + +func isValidTaskID(taskID string) bool { + return validTaskIDRegex.MatchString(taskID) +} + +func appendUnique(slice []string, items ...string) []string { + seen := make(map[string]struct{}, len(slice)) + for _, s := range slice { + seen[s] = struct{}{} + } + for _, item := range items { + if _, exists := seen[item]; !exists { + slice = append(slice, item) + seen[item] = struct{}{} + } + } + return slice +} + +func hasCyclicDependency(taskMap map[string]*task, blockerID, blockedID string) bool { + if blockerID == blockedID { + return true + } + + visited := make(map[string]bool) + return canReach(taskMap, blockedID, blockerID, visited) +} + +func canReach(taskMap map[string]*task, fromID, toID string, visited map[string]bool) bool { + if fromID == toID { + return true + } + if visited[fromID] { + return false + } + visited[fromID] = true + + fromTask, exists := taskMap[fromID] + if !exists { + return false + } + + for _, blockedID := range fromTask.Blocks { + if canReach(taskMap, blockedID, toID, visited) { + return true + } + } + + return false +} diff --git a/adk/middlewares/plantask/task_create.go b/adk/middlewares/plantask/task_create.go new file mode 100644 index 0000000..478b14d --- /dev/null +++ b/adk/middlewares/plantask/task_create.go @@ -0,0 +1,265 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "fmt" + "path/filepath" + "sync" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func newTaskCreateTool(backend Backend, baseDir string, lock *sync.Mutex) *taskCreateTool { + return &taskCreateTool{Backend: backend, BaseDir: baseDir, lock: lock} +} + +type taskCreateTool struct { + Backend Backend + BaseDir string + lock *sync.Mutex +} + +type taskCreateArgs struct { + Subject string `json:"subject"` + Description string `json:"description"` + ActiveForm string `json:"activeForm,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +func (t *taskCreateTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + desc := internal.SelectPrompt(internal.I18nPrompts{ + English: taskCreateToolDesc, + Chinese: taskCreateToolDescChinese, + }) + + return &schema.ToolInfo{ + Name: TaskCreateToolName, + Desc: desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "subject": { + Type: schema.String, + Desc: "A brief title for the task", + Required: true, + }, + "description": { + Type: schema.String, + Desc: "A detailed description of what needs to be done", + Required: true, + }, + "activeForm": { + Type: schema.String, + Desc: "Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")", + Required: false, + }, + "metadata": { + Type: schema.Object, + Desc: "Arbitrary metadata to attach to the task", + SubParams: map[string]*schema.ParameterInfo{ + "propertyNames": { + Type: schema.String, + }, + }, + Required: false, + }, + }), + }, nil +} + +func (t *taskCreateTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + t.lock.Lock() + defer t.lock.Unlock() + + params := &taskCreateArgs{} + err := sonic.UnmarshalString(argumentsInJSON, params) + if err != nil { + return "", err + } + + files, err := t.Backend.LsInfo(ctx, &LsInfoRequest{ + Path: t.BaseDir, + }) + if err != nil { + return "", fmt.Errorf("%s list files in %s failed, err: %w", TaskCreateToolName, t.BaseDir, err) + } + + highwatermark := int64(0) + for _, file := range files { + fileName := filepath.Base(file.Path) + if fileName == highWatermarkFileName { + content, readErr := t.Backend.Read(ctx, &ReadRequest{ + FilePath: file.Path, + }) + if readErr != nil { + return "", fmt.Errorf("%s read highwatermark file %s failed, err: %w", TaskCreateToolName, file.Path, readErr) + } + if content.Content != "" { + var val int64 + if _, scanErr := fmt.Sscanf(content.Content, "%d", &val); scanErr == nil { + highwatermark = val + } + } + break + } + } + + taskID := highwatermark + 1 + taskFileName := fmt.Sprintf("%d.json", taskID) + + for _, file := range files { + fileName := filepath.Base(file.Path) + if fileName == taskFileName { + return "", fmt.Errorf("task #%d already exists", taskID) + } + } + + newTask := &task{ + ID: fmt.Sprintf("%d", taskID), + Subject: params.Subject, + Description: params.Description, + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + ActiveForm: params.ActiveForm, + Metadata: params.Metadata, + } + + taskData, err := sonic.MarshalString(newTask) + if err != nil { + return "", fmt.Errorf("%s marshal task #%d failed, err: %w", TaskCreateToolName, taskID, err) + } + + // Write highwatermark file first + highwatermarkPath := filepath.Join(t.BaseDir, highWatermarkFileName) + err = t.Backend.Write(ctx, &WriteRequest{ + FilePath: highwatermarkPath, + Content: fmt.Sprintf("%d", taskID), + }) + if err != nil { + return "", fmt.Errorf("%s update highwatermark file %s failed, err: %w", TaskCreateToolName, highwatermarkPath, err) + } + + taskFilePath := filepath.Join(t.BaseDir, taskFileName) + err = t.Backend.Write(ctx, &WriteRequest{ + FilePath: taskFilePath, + Content: taskData, + }) + if err != nil { + return "", fmt.Errorf("%s create Task #%d failed, err: %w", TaskCreateToolName, taskID, err) + } + + resp := &taskOut{ + Result: fmt.Sprintf("Task #%d created successfully: %s", taskID, params.Subject), + } + + jsonResp, err := sonic.MarshalString(resp) + if err != nil { + return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskCreateToolName, err) + } + + return jsonResp, nil +} + +const TaskCreateToolName = "TaskCreate" +const taskCreateToolDesc = `Use this tool to create a structured task list for your current coding session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user. +It also helps the user understand the progress of the task and overall progress of their requests. + +## When to Use This Tool + +Use this tool proactively in these scenarios: + +- Complex multi-step tasks - When a task requires 3 or more distinct steps or actions +- Non-trivial and complex tasks - Tasks that require careful planning or multiple operations +- Plan mode - When using plan mode, create a task list to track the work +- User explicitly requests todo list - When the user directly asks you to use the todo list +- User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated) +- After receiving new instructions - Immediately capture user requirements as tasks +- When you start working on a task - Mark it as in_progress BEFORE beginning work +- After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation + +## When NOT to Use This Tool + +Skip using this tool when: +- There is only a single, straightforward task +- The task is trivial and tracking it provides no organizational benefit +- The task can be completed in less than 3 trivial steps +- The task is purely conversational or informational + +NOTE that you should not use this tool if there is only one trivial task to do. In this case you are better off just doing the task directly. + +## Task Fields + +- **subject**: A brief, actionable title in imperative form (e.g., "Fix authentication bug in login flow") +- **description**: Detailed description of what needs to be done, including context and acceptance criteria +- **activeForm**: Present continuous form shown in spinner when task is in_progress (e.g., "Fixing authentication bug"). This is displayed to the user while you work on the task. + +**IMPORTANT**: Always provide activeForm when creating tasks. The subject should be imperative ("Run tests") while activeForm should be present continuous ("Running tests"). All tasks are created with status "pending". + +## Tips + +- Create tasks with clear, specific subjects that describe the outcome +- Include enough detail in the description for another agent to understand and complete the task +- After creating tasks, use TaskUpdate to set up dependencies (blocks/blockedBy) if needed +- Check TaskList first to avoid creating duplicate tasks +` + +const taskCreateToolDescChinese = `使用此工具为当前编码会话创建结构化的任务列表。这有助于跟踪进度、组织复杂任务，并向用户展示工作的完整性。 +它还帮助用户了解任务的进度和请求的整体进展。 + +## 何时使用此工具 + +在以下场景中主动使用此工具： + +- 复杂的多步骤任务 - 当任务需要 3 个或更多不同的步骤或操作时 +- 非简单的复杂任务 - 需要仔细规划或多个操作的任务 +- 计划模式 - 使用计划模式时，创建任务列表来跟踪工作 +- 用户明确要求待办列表 - 当用户直接要求使用待办列表时 +- 用户提供多个任务 - 当用户提供待办事项列表时（编号或逗号分隔） +- 收到新指令后 - 立即将用户需求记录为任务 +- 开始处理任务时 - 在开始工作之前将其标记为 in_progress +- 完成任务后 - 将其标记为已完成，并添加实施过程中发现的任何后续任务 + +## 何时不使用此工具 + +在以下情况下跳过使用此工具： +- 只有一个简单直接的任务 +- 任务很简单，跟踪它没有组织上的好处 +- 任务可以在少于 3 个简单步骤内完成 +- 任务纯粹是对话性或信息性的 + +注意：如果只有一个简单任务要做，不应该使用此工具。在这种情况下，直接完成任务更好。 + +## 任务字段 + +- **subject**：简短的、可操作的标题，使用祈使句形式（例如，"修复登录流程中的认证错误"） +- **description**：需要完成的工作的详细描述，包括上下文和验收标准 +- **activeForm**：任务处于 in_progress 状态时在加载动画中显示的现在进行时形式（例如，"正在修复认证错误"）。这会在你处理任务时显示给用户。 + +**重要**：创建任务时始终提供 activeForm。subject 应该是祈使句（"运行测试"），而 activeForm 应该是现在进行时（"正在运行测试"）。所有任务创建时状态为 "pending"。 + +## 提示 + +- 创建具有清晰、具体主题的任务，描述预期结果 +- 在描述中包含足够的细节，以便其他代理能够理解并完成任务 +- 创建任务后，如果需要，使用 TaskUpdate 设置依赖关系（blocks/blockedBy） +- 先检查 TaskList 以避免创建重复任务 +` diff --git a/adk/middlewares/plantask/task_create_test.go b/adk/middlewares/plantask/task_create_test.go new file mode 100644 index 0000000..e451ffb --- /dev/null +++ b/adk/middlewares/plantask/task_create_test.go @@ -0,0 +1,91 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "path/filepath" + "sync" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" +) + +func TestTaskCreateTool(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + tool := newTaskCreateTool(backend, baseDir, lock) + + info, err := tool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, TaskCreateToolName, info.Name) + assert.Equal(t, taskCreateToolDesc, info.Desc) + + result, err := tool.InvokableRun(ctx, `{"subject": "Test Task", "description": "Test description", "activeForm": "Testing"}`) + assert.NoError(t, err) + assert.Equal(t, `{"result":"Task #1 created successfully: Test Task"}`, result) + + content, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.NoError(t, err) + + var taskData task + err = sonic.UnmarshalString(content.Content, &taskData) + assert.NoError(t, err) + assert.Equal(t, "1", taskData.ID) + assert.Equal(t, "Test Task", taskData.Subject) + assert.Equal(t, "Test description", taskData.Description) + assert.Equal(t, taskStatusPending, taskData.Status) + assert.Equal(t, "Testing", taskData.ActiveForm) + + hwContent, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, highWatermarkFileName)}) + assert.NoError(t, err) + assert.Equal(t, "1", hwContent.Content) + + result, err = tool.InvokableRun(ctx, `{"subject": "Second Task", "description": "Second description"}`) + assert.NoError(t, err) + assert.Equal(t, `{"result":"Task #2 created successfully: Second Task"}`, result) + + hwContent, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, highWatermarkFileName)}) + assert.NoError(t, err) + assert.Equal(t, "2", hwContent.Content) +} + +func TestTaskCreateToolWithMetadata(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + tool := newTaskCreateTool(backend, baseDir, lock) + + result, err := tool.InvokableRun(ctx, `{"subject": "Task with metadata", "description": "Has metadata", "metadata": {"key1": "value1", "key2": "value2"}}`) + assert.NoError(t, err) + assert.Contains(t, result, "Task #1 created successfully") + + content, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.NoError(t, err) + + var taskData task + err = sonic.UnmarshalString(content.Content, &taskData) + assert.NoError(t, err) + assert.Equal(t, "value1", taskData.Metadata["key1"]) + assert.Equal(t, "value2", taskData.Metadata["key2"]) +} diff --git a/adk/middlewares/plantask/task_get.go b/adk/middlewares/plantask/task_get.go new file mode 100644 index 0000000..55760c3 --- /dev/null +++ b/adk/middlewares/plantask/task_get.go @@ -0,0 +1,176 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "fmt" + "path/filepath" + "strings" + "sync" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func newTaskGetTool(backend Backend, baseDir string, lock *sync.Mutex) *taskGetTool { + return &taskGetTool{Backend: backend, BaseDir: baseDir, lock: lock} +} + +type taskGetTool struct { + Backend Backend + BaseDir string + lock *sync.Mutex +} + +func (t *taskGetTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + desc := internal.SelectPrompt(internal.I18nPrompts{ + English: taskGetToolDesc, + Chinese: taskGetToolDescChinese, + }) + + return &schema.ToolInfo{ + Name: TaskGetToolName, + Desc: desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "taskId": { + Type: schema.String, + Desc: "The ID of the task to retrieve", + Required: true, + }, + }), + }, nil +} + +type taskGetArgs struct { + TaskID string `json:"taskId"` +} + +func (t *taskGetTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + t.lock.Lock() + defer t.lock.Unlock() + + params := &taskGetArgs{} + err := sonic.UnmarshalString(argumentsInJSON, params) + if err != nil { + return "", err + } + + if !isValidTaskID(params.TaskID) { + return "", fmt.Errorf("%s validate task ID failed, err: invalid format: %s", TaskGetToolName, params.TaskID) + } + + taskFileName := fmt.Sprintf("%s.json", params.TaskID) + taskFilePath := filepath.Join(t.BaseDir, taskFileName) + + content, err := t.Backend.Read(ctx, &ReadRequest{ + FilePath: taskFilePath, + }) + if err != nil { + return "", fmt.Errorf("%s get Task #%s failed, err: %w", TaskGetToolName, params.TaskID, err) + } + + taskData := &task{} + err = sonic.UnmarshalString(content.Content, taskData) + if err != nil { + return "", fmt.Errorf("%s get Task #%s failed, err: %w", TaskGetToolName, params.TaskID, err) + } + + var result strings.Builder + result.WriteString(fmt.Sprintf("Task #%s: %s\n", taskData.ID, taskData.Subject)) + result.WriteString(fmt.Sprintf("Status: %s\n", taskData.Status)) + result.WriteString(fmt.Sprintf("Description: %s\n", taskData.Description)) + + if len(taskData.BlockedBy) > 0 { + blockedByIDs := make([]string, len(taskData.BlockedBy)) + for i, id := range taskData.BlockedBy { + blockedByIDs[i] = "#" + id + } + result.WriteString(fmt.Sprintf("Blocked by: %s\n", strings.Join(blockedByIDs, ", "))) + } + if len(taskData.Blocks) > 0 { + blocksIDs := make([]string, len(taskData.Blocks)) + for i, id := range taskData.Blocks { + blocksIDs[i] = "#" + id + } + result.WriteString(fmt.Sprintf("Blocks: %s\n", strings.Join(blocksIDs, ", "))) + } + if taskData.Owner != "" { + result.WriteString(fmt.Sprintf("Owner: %s\n", taskData.Owner)) + } + + resp := &taskOut{ + Result: result.String(), + } + + jsonResp, err := sonic.MarshalString(resp) + if err != nil { + return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskGetToolName, err) + } + + return jsonResp, nil +} + +const TaskGetToolName = "TaskGet" +const taskGetToolDesc = `Use this tool to retrieve a task by its ID from the task list. + +## When to Use This Tool + +- When you need the full description and context before starting work on a task +- To understand task dependencies (what it blocks, what blocks it) +- After being assigned a task, to get complete requirements + +## Output + +Returns full task details: +- **subject**: Task title +- **description**: Detailed requirements and context +- **status**: 'pending', 'in_progress', or 'completed' +- **blocks**: Tasks waiting on this one to complete +- **blockedBy**: Tasks that must complete before this one can start + +## Tips + +- After fetching a task, verify its blockedBy list is empty before beginning work. +- Use TaskList to see all tasks in summary form. +` + +const taskGetToolDescChinese = `使用此工具通过任务 ID 从任务列表中获取任务。 + +## 何时使用此工具 + +- 当你需要在开始处理任务之前获取完整的描述和上下文时 +- 了解任务依赖关系（它阻塞什么，什么阻塞它） +- 被分配任务后，获取完整的需求 + +## 输出 + +返回完整的任务详情： +- **subject**：任务标题 +- **description**：详细的需求和上下文 +- **status**：'pending'、'in_progress' 或 'completed' +- **blocks**：等待此任务完成的任务 +- **blockedBy**：必须在此任务开始之前完成的任务 + +## 提示 + +- 获取任务后，在开始工作之前验证其 blockedBy 列表是否为空。 +- 使用 TaskList 查看所有任务的摘要形式。 +` diff --git a/adk/middlewares/plantask/task_get_test.go b/adk/middlewares/plantask/task_get_test.go new file mode 100644 index 0000000..f1f9863 --- /dev/null +++ b/adk/middlewares/plantask/task_get_test.go @@ -0,0 +1,80 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "path/filepath" + "sync" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" +) + +func TestTaskGetTool(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + taskData := &task{ + ID: "1", + Subject: "Test Task", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{"2", "3"}, + BlockedBy: []string{"4"}, + } + taskJSON, _ := sonic.MarshalString(taskData) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: taskJSON}) + + tool := newTaskGetTool(backend, baseDir, lock) + + info, err := tool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, TaskGetToolName, info.Name) + assert.Equal(t, taskGetToolDesc, info.Desc) + + result, err := tool.InvokableRun(ctx, `{"taskId": "1"}`) + assert.NoError(t, err) + assert.Contains(t, result, "Task #1: Test Task") + assert.Contains(t, result, "Status: "+taskStatusPending) + assert.Contains(t, result, "Description: Test description") + assert.Contains(t, result, "Blocked by: #4") + assert.Contains(t, result, "Blocks: #2, #3") + + _, err = tool.InvokableRun(ctx, `{"taskId": "999"}`) + assert.Error(t, err) +} + +func TestTaskGetToolInvalidTaskID(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + tool := newTaskGetTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "../../../etc/passwd"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate task ID failed") + + _, err = tool.InvokableRun(ctx, `{"taskId": "abc"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate task ID failed") +} diff --git a/adk/middlewares/plantask/task_list.go b/adk/middlewares/plantask/task_list.go new file mode 100644 index 0000000..60a7d04 --- /dev/null +++ b/adk/middlewares/plantask/task_list.go @@ -0,0 +1,194 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "fmt" + "path/filepath" + "sort" + "strings" + "sync" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func newTaskListTool(backend Backend, baseDir string, lock *sync.Mutex) *taskListTool { + return &taskListTool{Backend: backend, BaseDir: baseDir, lock: lock} +} + +type taskListTool struct { + Backend Backend + BaseDir string + lock *sync.Mutex +} + +func (t *taskListTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + desc := internal.SelectPrompt(internal.I18nPrompts{ + English: taskListToolDesc, + Chinese: taskListToolDescChinese, + }) + + return &schema.ToolInfo{ + Name: TaskListToolName, + Desc: desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{}), + }, nil +} + +func listTasks(ctx context.Context, backend Backend, baseDir string) ([]*task, error) { + files, err := backend.LsInfo(ctx, &LsInfoRequest{ + Path: baseDir, + }) + if err != nil { + return nil, fmt.Errorf("%s list files in %s failed, err: %w", TaskListToolName, baseDir, err) + } + + var tasks []*task + for _, file := range files { + fileName := filepath.Base(file.Path) + if !strings.HasSuffix(fileName, ".json") { + continue + } + + taskID := strings.TrimSuffix(fileName, ".json") + if !isValidTaskID(taskID) { + continue + } + + content, err := backend.Read(ctx, &ReadRequest{ + FilePath: file.Path, + }) + if err != nil { + return nil, fmt.Errorf("%s read task file %s failed, err: %w", TaskListToolName, file.Path, err) + } + + taskData := &task{} + err = sonic.UnmarshalString(content.Content, taskData) + if err != nil { + return nil, fmt.Errorf("%s parse task file %s failed, err: %w", TaskListToolName, file.Path, err) + } + + tasks = append(tasks, taskData) + } + + // sort tasks by ID + sort.Slice(tasks, func(i, j int) bool { + return tasks[i].ID < tasks[j].ID + }) + + return tasks, nil +} + +func (t *taskListTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + t.lock.Lock() + defer t.lock.Unlock() + + tasks, err := listTasks(ctx, t.Backend, t.BaseDir) + if err != nil { + return "", err + } + + if len(tasks) == 0 { + resp := &taskOut{ + Result: "No tasks found.", + } + jsonResp, marshalErr := sonic.MarshalString(resp) + if marshalErr != nil { + return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskListToolName, marshalErr) + } + return jsonResp, nil + } + + var result strings.Builder + for i, taskData := range tasks { + if i > 0 { + result.WriteString("\n") + } + result.WriteString(fmt.Sprintf("#%s [%s] %s", taskData.ID, taskData.Status, taskData.Subject)) + if taskData.Owner != "" { + result.WriteString(fmt.Sprintf(" [owner: %s]", taskData.Owner)) + } + if len(taskData.BlockedBy) > 0 { + blockedByIDs := make([]string, len(taskData.BlockedBy)) + for j, id := range taskData.BlockedBy { + blockedByIDs[j] = "#" + id + } + result.WriteString(fmt.Sprintf(" [blocked by %s]", strings.Join(blockedByIDs, ", "))) + } + } + + resp := &taskOut{ + Result: result.String(), + } + + jsonResp, err := sonic.MarshalString(resp) + if err != nil { + return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskListToolName, err) + } + + return jsonResp, nil +} + +const TaskListToolName = "TaskList" +const taskListToolDesc = `Use this tool to list all tasks in the task list. + +## When to Use This Tool + +- To see what tasks are available to work on (status: 'pending', no owner, not blocked) +- To check overall progress on the project +- To find tasks that are blocked and need dependencies resolved +- After completing a task, to check for newly unblocked work or claim the next available task +- **Prefer working on tasks in ID order** (lowest ID first) when multiple tasks are available, as earlier tasks often set up context for later ones + +## Output + +Returns a summary of each task: +- **id**: Task identifier (use with TaskGet, TaskUpdate) +- **subject**: Brief description of the task +- **status**: 'pending', 'in_progress', or 'completed' +- **owner**: Agent ID if assigned, empty if available +- **blockedBy**: List of open task IDs that must be resolved first (tasks with blockedBy cannot be claimed until dependencies resolve) + +Use TaskGet with a specific task ID to view full details including description and comments. +` + +const taskListToolDescChinese = `使用此工具列出任务列表中的所有任务。 + +## 何时使用此工具 + +- 查看可以处理的任务（状态：'pending'，无所有者，未被阻塞） +- 检查项目的整体进度 +- 查找被阻塞且需要解决依赖关系的任务 +- 完成任务后，检查新解除阻塞的工作或认领下一个可用任务 +- **优先按 ID 顺序处理任务**（最小 ID 优先），当有多个任务可用时，因为较早的任务通常为后续任务建立上下文 + +## 输出 + +返回每个任务的摘要： +- **id**：任务标识符（与 TaskGet、TaskUpdate 一起使用） +- **subject**：任务的简要描述 +- **status**：'pending'、'in_progress' 或 'completed' +- **owner**：如果已分配则为代理 ID，如果可用则为空 +- **blockedBy**：必须首先解决的开放任务 ID 列表（具有 blockedBy 的任务在依赖关系解决之前无法被认领） + +使用 TaskGet 配合特定任务 ID 查看完整详情，包括描述和评论。 +` diff --git a/adk/middlewares/plantask/task_list_test.go b/adk/middlewares/plantask/task_list_test.go new file mode 100644 index 0000000..706f8c6 --- /dev/null +++ b/adk/middlewares/plantask/task_list_test.go @@ -0,0 +1,60 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "path/filepath" + "sync" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" +) + +func TestTaskListTool(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + tool := newTaskListTool(backend, baseDir, lock) + + info, err := tool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, TaskListToolName, info.Name) + assert.Equal(t, taskListToolDesc, info.Desc) + + result, err := tool.InvokableRun(ctx, `{}`) + assert.NoError(t, err) + assert.Equal(t, `{"result":"No tasks found."}`, result) + + task1 := &task{ID: "1", Subject: "Task 1", Status: taskStatusPending, BlockedBy: []string{"2"}} + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ID: "2", Subject: "Task 2", Status: taskStatusInProgress, Owner: "agent1"} + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + result, err = tool.InvokableRun(ctx, `{}`) + assert.NoError(t, err) + assert.Contains(t, result, "#1 ["+taskStatusPending+"] Task 1") + assert.Contains(t, result, "[blocked by #2]") + assert.Contains(t, result, "#2 ["+taskStatusInProgress+"] Task 2") + assert.Contains(t, result, "[owner: agent1]") +} diff --git a/adk/middlewares/plantask/task_update.go b/adk/middlewares/plantask/task_update.go new file mode 100644 index 0000000..7e9eb2d --- /dev/null +++ b/adk/middlewares/plantask/task_update.go @@ -0,0 +1,565 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "fmt" + "path/filepath" + "strings" + "sync" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func newTaskUpdateTool(backend Backend, baseDir string, lock *sync.Mutex) *taskUpdateTool { + return &taskUpdateTool{Backend: backend, BaseDir: baseDir, lock: lock} +} + +type taskUpdateTool struct { + Backend Backend + BaseDir string + lock *sync.Mutex +} + +type taskUpdateArgs struct { + TaskID string `json:"taskId"` + Subject string `json:"subject,omitempty"` + Description string `json:"description,omitempty"` + ActiveForm string `json:"activeForm,omitempty"` + Status string `json:"status,omitempty"` + AddBlocks []string `json:"addBlocks,omitempty"` + AddBlockedBy []string `json:"addBlockedBy,omitempty"` + Owner string `json:"owner,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +func (t *taskUpdateTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + desc := internal.SelectPrompt(internal.I18nPrompts{ + English: taskUpdateToolDesc, + Chinese: taskUpdateToolDescChinese, + }) + + return &schema.ToolInfo{ + Name: TaskUpdateToolName, + Desc: desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "taskId": { + Type: schema.String, + Desc: "The ID of the task to update", + Required: true, + }, + "subject": { + Type: schema.String, + Desc: "New subject for the task", + Required: false, + }, + "description": { + Type: schema.String, + Desc: "New description for the task", + Required: false, + }, + "activeForm": { + Type: schema.String, + Desc: "Present continuous form shown in spinner when in_progress (e.g., \"Running tests\")", + Required: false, + }, + "status": { + Type: schema.String, + Desc: "New status for the task: 'pending', 'in_progress', 'completed', or 'deleted'", + Enum: []string{"pending", "in_progress", "completed", "deleted"}, + Required: false, + }, + "addBlocks": { + Type: schema.Array, + Desc: "Task IDs that this task blocks", + ElemInfo: &schema.ParameterInfo{Type: schema.String}, + Required: false, + }, + "addBlockedBy": { + Type: schema.Array, + Desc: "Task IDs that block this task", + ElemInfo: &schema.ParameterInfo{Type: schema.String}, + Required: false, + }, + "owner": { + Type: schema.String, + Desc: "New owner for the task", + Required: false, + }, + "metadata": { + Type: schema.Object, + Desc: "Metadata keys to merge into the task. Set a key to null to delete it.", + SubParams: map[string]*schema.ParameterInfo{ + "propertyNames": { + Type: schema.String, + }, + }, + Required: false, + }, + }), + }, nil +} + +func (t *taskUpdateTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + t.lock.Lock() + defer t.lock.Unlock() + + params := &taskUpdateArgs{} + err := sonic.UnmarshalString(argumentsInJSON, params) + if err != nil { + return "", err + } + + if !isValidTaskID(params.TaskID) { + return "", fmt.Errorf("%s validate task ID failed, err: invalid format: %s", TaskUpdateToolName, params.TaskID) + } + + taskFileName := fmt.Sprintf("%s.json", params.TaskID) + taskFilePath := filepath.Join(t.BaseDir, taskFileName) + + if params.Status == taskStatusDeleted { + if removeErr := t.removeTaskFromDependencies(ctx, params.TaskID); removeErr != nil { + return "", fmt.Errorf("%s remove Task #%s from dependencies failed, err: %w", TaskUpdateToolName, params.TaskID, removeErr) + } + + err = t.Backend.Delete(ctx, &DeleteRequest{ + FilePath: taskFilePath, + }) + if err != nil { + return "", fmt.Errorf("%s delete Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err) + } + + resp := &taskOut{ + Result: fmt.Sprintf("Updated task #%s deleted", params.TaskID), + } + jsonResp, marshalErr := sonic.MarshalString(resp) + if marshalErr != nil { + return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskUpdateToolName, marshalErr) + } + return jsonResp, nil + } + + content, err := t.Backend.Read(ctx, &ReadRequest{ + FilePath: taskFilePath, + }) + if err != nil { + return "", fmt.Errorf("%s read Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err) + } + + taskData := &task{} + err = sonic.UnmarshalString(content.Content, taskData) + if err != nil { + return "", fmt.Errorf("%s parse Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err) + } + + var updatedFields []string + + if params.Subject != "" { + taskData.Subject = params.Subject + updatedFields = append(updatedFields, "subject") + } + if params.Description != "" { + taskData.Description = params.Description + updatedFields = append(updatedFields, "description") + } + if params.ActiveForm != "" { + taskData.ActiveForm = params.ActiveForm + updatedFields = append(updatedFields, "activeForm") + } + if params.Status != "" { + taskData.Status = params.Status + updatedFields = append(updatedFields, "status") + } + if len(params.AddBlocks) > 0 || len(params.AddBlockedBy) > 0 { + tasks, listErr := listTasks(ctx, t.Backend, t.BaseDir) + if listErr != nil { + return "", fmt.Errorf("%s list tasks failed, err: %w", TaskUpdateToolName, listErr) + } + taskMap := make(map[string]*task, len(tasks)) + for _, tk := range tasks { + taskMap[tk.ID] = tk + } + + if len(params.AddBlocks) > 0 { + for _, blockedTaskID := range params.AddBlocks { + if !isValidTaskID(blockedTaskID) { + return "", fmt.Errorf("%s validate blocked task ID failed, err: invalid format: %s", TaskUpdateToolName, blockedTaskID) + } + if hasCyclicDependency(taskMap, params.TaskID, blockedTaskID) { + return "", fmt.Errorf("%s adding Task #%s to blocks of Task #%s would create a cyclic dependency", TaskUpdateToolName, blockedTaskID, params.TaskID) + } + } + for _, blockedTaskID := range params.AddBlocks { + if addErr := t.addBlockedByToTask(ctx, blockedTaskID, params.TaskID); addErr != nil { + return "", fmt.Errorf("%s update Task #%s blocks failed, err: %w", TaskUpdateToolName, params.TaskID, addErr) + } + } + taskData.Blocks = appendUnique(taskData.Blocks, params.AddBlocks...) + updatedFields = append(updatedFields, "blocks") + } + if len(params.AddBlockedBy) > 0 { + for _, blockingTaskID := range params.AddBlockedBy { + if !isValidTaskID(blockingTaskID) { + return "", fmt.Errorf("%s validate blocking task ID failed, err: invalid format: %s", TaskUpdateToolName, blockingTaskID) + } + if hasCyclicDependency(taskMap, blockingTaskID, params.TaskID) { + return "", fmt.Errorf("%s adding Task #%s to blockedBy of Task #%s would create a cyclic dependency", TaskUpdateToolName, blockingTaskID, params.TaskID) + } + } + for _, blockingTaskID := range params.AddBlockedBy { + if addErr := t.addBlocksToTask(ctx, blockingTaskID, params.TaskID); addErr != nil { + return "", fmt.Errorf("%s update Task #%s blockedBy failed, err: %w", TaskUpdateToolName, params.TaskID, addErr) + } + } + taskData.BlockedBy = appendUnique(taskData.BlockedBy, params.AddBlockedBy...) + updatedFields = append(updatedFields, "blockedBy") + } + } + if params.Owner != "" { + taskData.Owner = params.Owner + updatedFields = append(updatedFields, "owner") + } + if params.Metadata != nil { + if taskData.Metadata == nil { + taskData.Metadata = make(map[string]any) + } + for k, v := range params.Metadata { + if v == nil { + delete(taskData.Metadata, k) + } else { + taskData.Metadata[k] = v + } + } + updatedFields = append(updatedFields, "metadata") + } + + updatedContent, err := sonic.MarshalString(taskData) + if err != nil { + return "", fmt.Errorf("%s marshal Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err) + } + + err = t.Backend.Write(ctx, &WriteRequest{ + FilePath: taskFilePath, + Content: updatedContent, + }) + if err != nil { + return "", fmt.Errorf("%s write Task #%s failed, err: %w", TaskUpdateToolName, params.TaskID, err) + } + + if params.Status == taskStatusCompleted { + if checkErr := t.checkIfNeedDeleteAllTasks(ctx); checkErr != nil { + return "", fmt.Errorf("%s check and delete all tasks failed, err: %w", TaskUpdateToolName, checkErr) + } + } + + resp := &taskOut{ + Result: fmt.Sprintf("Updated task #%s %s", params.TaskID, strings.Join(updatedFields, ", ")), + } + + jsonResp, err := sonic.MarshalString(resp) + if err != nil { + return "", fmt.Errorf("%s marshal taskOut failed, err: %w", TaskUpdateToolName, err) + } + + return jsonResp, nil +} + +func (t *taskUpdateTool) removeTaskFromDependencies(ctx context.Context, deletedTaskID string) error { + tasks, err := listTasks(ctx, t.Backend, t.BaseDir) + if err != nil { + return err + } + + for _, taskData := range tasks { + if taskData.ID == deletedTaskID { + continue + } + + modified := false + newBlocks := make([]string, 0, len(taskData.Blocks)) + for _, id := range taskData.Blocks { + if id != deletedTaskID { + newBlocks = append(newBlocks, id) + } else { + modified = true + } + } + + newBlockedBy := make([]string, 0, len(taskData.BlockedBy)) + for _, id := range taskData.BlockedBy { + if id != deletedTaskID { + newBlockedBy = append(newBlockedBy, id) + } else { + modified = true + } + } + + if modified { + taskData.Blocks = newBlocks + taskData.BlockedBy = newBlockedBy + + updatedContent, err := sonic.MarshalString(taskData) + if err != nil { + return fmt.Errorf("failed to marshal task #%s: %w", taskData.ID, err) + } + + taskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf("%s.json", taskData.ID)) + if err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil { + return fmt.Errorf("failed to write task #%s: %w", taskData.ID, err) + } + } + } + + return nil +} + +func (t *taskUpdateTool) addBlockedByToTask(ctx context.Context, targetTaskID, blockerTaskID string) error { + taskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf("%s.json", targetTaskID)) + + content, err := t.Backend.Read(ctx, &ReadRequest{FilePath: taskFilePath}) + if err != nil { + return fmt.Errorf("failed to read task #%s for updating blockedBy: %w", targetTaskID, err) + } + + targetTask := &task{} + if unmarshalErr := sonic.UnmarshalString(content.Content, targetTask); unmarshalErr != nil { + return fmt.Errorf("failed to parse task #%s: %w", targetTaskID, unmarshalErr) + } + + targetTask.BlockedBy = appendUnique(targetTask.BlockedBy, blockerTaskID) + + updatedContent, err := sonic.MarshalString(targetTask) + if err != nil { + return fmt.Errorf("failed to marshal task #%s: %w", targetTaskID, err) + } + + if err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil { + return fmt.Errorf("failed to write task #%s: %w", targetTaskID, err) + } + + return nil +} + +func (t *taskUpdateTool) addBlocksToTask(ctx context.Context, targetTaskID, blockedTaskID string) error { + taskFilePath := filepath.Join(t.BaseDir, fmt.Sprintf("%s.json", targetTaskID)) + + content, err := t.Backend.Read(ctx, &ReadRequest{FilePath: taskFilePath}) + if err != nil { + return fmt.Errorf("failed to read task #%s for updating blocks: %w", targetTaskID, err) + } + + targetTask := &task{} + if unmarshalErr := sonic.UnmarshalString(content.Content, targetTask); unmarshalErr != nil { + return fmt.Errorf("failed to parse task #%s: %w", targetTaskID, unmarshalErr) + } + + targetTask.Blocks = appendUnique(targetTask.Blocks, blockedTaskID) + + updatedContent, err := sonic.MarshalString(targetTask) + if err != nil { + return fmt.Errorf("failed to marshal task #%s: %w", targetTaskID, err) + } + + if err := t.Backend.Write(ctx, &WriteRequest{FilePath: taskFilePath, Content: updatedContent}); err != nil { + return fmt.Errorf("failed to write task #%s: %w", targetTaskID, err) + } + + return nil +} + +// checkIfNeedDeleteAllTasks checks if all tasks are completed, if so, it deletes all tasks +func (t *taskUpdateTool) checkIfNeedDeleteAllTasks(ctx context.Context) error { + tasks, err := listTasks(ctx, t.Backend, t.BaseDir) + if err != nil { + return err + } + + for _, task := range tasks { + if task.Status != taskStatusCompleted { + return nil + } + } + + for _, task := range tasks { + err := t.Backend.Delete(ctx, &DeleteRequest{ + FilePath: filepath.Join(t.BaseDir, task.ID+".json"), + }) + if err != nil { + return err + } + } + + return nil +} + +const TaskUpdateToolName = "TaskUpdate" +const taskUpdateToolDesc = `Use this tool to update a task in the task list. + +## When to Use This Tool + +**Mark tasks as resolved:** +- When you have completed the work described in a task +- When a task is no longer needed or has been superseded +- IMPORTANT: Always mark your assigned tasks as resolved when you finish them +- After resolving, call TaskList to find your next task + +- ONLY mark a task as completed when you have FULLY accomplished it +- If you encounter errors, blockers, or cannot finish, keep the task as in_progress +- When blocked, create a new task describing what needs to be resolved +- Never mark a task as completed if: + - Tests are failing + - Implementation is partial + - You encountered unresolved errors + - You couldn't find necessary files or dependencies + +**Delete tasks:** +- When a task is no longer relevant or was created in error +- Setting status to ` + "`deleted`" + ` permanently removes the task + +**Update task details:** +- When requirements change or become clearer +- When establishing dependencies between tasks + +## Fields You Can Update + +- **status**: The task status (see Status Workflow below) +- **subject**: Change the task title (imperative form, e.g., "Run tests") +- **description**: Change the task description +- **activeForm**: Present continuous form shown in spinner when in_progress (e.g., "Running tests") +- **owner**: Change the task owner (agent name) +- **metadata**: Merge metadata keys into the task (set a key to null to delete it) +- **addBlocks**: Mark tasks that cannot start until this one completes +- **addBlockedBy**: Mark tasks that must complete before this one can start + +## Status Workflow + +Status progresses: ` + "`pending`" + ` → ` + "`in_progress`" + ` → ` + "`completed`" + ` + +Use ` + "`deleted`" + ` to permanently remove a task. + +## Staleness + +Make sure to read a task's latest state using ` + "`TaskGet`" + ` before updating it. + +## Examples + +Mark task as in progress when starting work: +` + "```json" + ` +{"taskId": "1", "status": "in_progress"} +` + "```" + ` + +Mark task as completed after finishing work: +` + "```json" + ` +{"taskId": "1", "status": "completed"} +` + "```" + ` + +Delete a task: +` + "```json" + ` +{"taskId": "1", "status": "deleted"} +` + "```" + ` + +Claim a task by setting owner: +` + "```json" + ` +{"taskId": "1", "owner": "my-name"} +` + "```" + ` + +Set up task dependencies: +` + "```json" + ` +{"taskId": "2", "addBlockedBy": ["1"]} +` + "```" + ` +` + +const taskUpdateToolDescChinese = `使用此工具更新任务列表中的任务。 + +## 何时使用此工具 + +**将任务标记为已完成：** +- 当你完成了任务中描述的工作时 +- 当任务不再需要或已被取代时 +- 重要：完成分配给你的任务后，务必将其标记为已完成 +- 完成后，调用 TaskList 查找下一个任务 + +- 只有在完全完成任务时才将其标记为已完成 +- 如果遇到错误、阻塞或无法完成，请保持任务为 in_progress 状态 +- 当被阻塞时，创建一个新任务描述需要解决的问题 +- 在以下情况下不要将任务标记为已完成： + - 测试失败 + - 实现不完整 + - 遇到未解决的错误 + - 找不到必要的文件或依赖项 + +**删除任务：** +- 当任务不再相关或创建错误时 +- 将状态设置为 ` + "`deleted`" + ` 会永久删除任务 + +**更新任务详情：** +- 当需求变更或变得更清晰时 +- 当建立任务之间的依赖关系时 + +## 可更新的字段 + +- **status**：任务状态（参见下方状态流程） +- **subject**：更改任务标题（使用祈使句形式，例如"运行测试"） +- **description**：更改任务描述 +- **activeForm**：in_progress 状态时在加载动画中显示的现在进行时形式（例如"正在运行测试"） +- **owner**：更改任务所有者（代理名称） +- **metadata**：将元数据键合并到任务中（将键设置为 null 可删除它） +- **addBlocks**：标记在此任务完成之前无法开始的任务 +- **addBlockedBy**：标记必须在此任务开始之前完成的任务 + +## 状态流程 + +状态进展：` + "`pending`" + ` → ` + "`in_progress`" + ` → ` + "`completed`" + ` + +使用 ` + "`deleted`" + ` 永久删除任务。 + +## 过期性 + +更新任务前，请确保使用 ` + "`TaskGet`" + ` 读取任务的最新状态。 + +## 示例 + +开始工作时将任务标记为进行中： +` + "```json" + ` +{"taskId": "1", "status": "in_progress"} +` + "```" + ` + +完成工作后将任务标记为已完成： +` + "```json" + ` +{"taskId": "1", "status": "completed"} +` + "```" + ` + +删除任务： +` + "```json" + ` +{"taskId": "1", "status": "deleted"} +` + "```" + ` + +通过设置 owner 认领任务： +` + "```json" + ` +{"taskId": "1", "owner": "my-name"} +` + "```" + ` + +设置任务依赖关系： +` + "```json" + ` +{"taskId": "2", "addBlockedBy": ["1"]} +` + "```" + ` +` diff --git a/adk/middlewares/plantask/task_update_test.go b/adk/middlewares/plantask/task_update_test.go new file mode 100644 index 0000000..2869dd6 --- /dev/null +++ b/adk/middlewares/plantask/task_update_test.go @@ -0,0 +1,739 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package plantask + +import ( + "context" + "path/filepath" + "sync" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" +) + +func TestTaskUpdateTool(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + taskData := &task{ + ID: "1", + Subject: "Original Subject", + Description: "Original description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + taskJSON, _ := sonic.MarshalString(taskData) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: taskJSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + info, err := tool.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, TaskUpdateToolName, info.Name) + assert.Equal(t, taskUpdateToolDesc, info.Desc) + + result, err := tool.InvokableRun(ctx, `{"taskId": "1", "status": "in_progress"}`) + assert.NoError(t, err) + assert.Contains(t, result, "Updated task #1") + assert.Contains(t, result, "status") + + content, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.NoError(t, err) + var updated task + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, taskStatusInProgress, updated.Status) + + result, err = tool.InvokableRun(ctx, `{"taskId": "1", "subject": "New Subject", "description": "New description"}`) + assert.NoError(t, err) + assert.Contains(t, result, "subject") + assert.Contains(t, result, "description") + + content, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, "New Subject", updated.Subject) + assert.Equal(t, "New description", updated.Description) +} + +func TestTaskUpdateToolOwnerAndMetadata(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + taskData := &task{ + ID: "1", + Subject: "Test Task", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + taskJSON, _ := sonic.MarshalString(taskData) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: taskJSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + result, err := tool.InvokableRun(ctx, `{"taskId": "1", "owner": "agent1"}`) + assert.NoError(t, err) + assert.Contains(t, result, "owner") + + content, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updated task + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, "agent1", updated.Owner) + + result, err = tool.InvokableRun(ctx, `{"taskId": "1", "metadata": {"key1": "value1", "key2": "value2"}}`) + assert.NoError(t, err) + assert.Contains(t, result, "metadata") + + content, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, "value1", updated.Metadata["key1"]) + assert.Equal(t, "value2", updated.Metadata["key2"]) + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "metadata": {"key1": null, "key3": "value3"}}`) + assert.NoError(t, err) + + content, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updated2 task + _ = sonic.UnmarshalString(content.Content, &updated2) + _, key1Exists := updated2.Metadata["key1"] + assert.False(t, key1Exists) + assert.Equal(t, "value2", updated2.Metadata["key2"]) + assert.Equal(t, "value3", updated2.Metadata["key3"]) +} + +func TestTaskUpdateToolBlocks(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Test Task", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + task4 := &task{ + ID: "4", + Subject: "Task 4", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task4JSON, _ := sonic.MarshalString(task4) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "4.json"), Content: task4JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + result, err := tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["2", "3"]}`) + assert.NoError(t, err) + assert.Contains(t, result, "blocks") + + content, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updated task + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, []string{"2", "3"}, updated.Blocks) + + result, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["4"]}`) + assert.NoError(t, err) + assert.Contains(t, result, "blockedBy") + + content, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, []string{"4"}, updated.BlockedBy) +} + +func TestTaskUpdateToolDelete(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + taskData := &task{ + ID: "1", + Subject: "Test Task", + Description: "Test description", + Status: taskStatusPending, + } + taskJSON, _ := sonic.MarshalString(taskData) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: taskJSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + result, err := tool.InvokableRun(ctx, `{"taskId": "1", "status": "deleted"}`) + assert.NoError(t, err) + assert.Contains(t, result, "deleted") + + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.Error(t, err) +} + +func TestTaskUpdateToolInvalidTaskID(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "../../../etc/passwd", "status": "in_progress"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate task ID failed") + + _, err = tool.InvokableRun(ctx, `{"taskId": "abc", "status": "in_progress"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate task ID failed") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1.5", "status": "in_progress"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate task ID failed") + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["invalid"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate blocked task ID failed") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["invalid"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "validate blocking task ID failed") +} + +func TestTaskUpdateToolBlocksDeduplication(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{"1"}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{"1"}, + BlockedBy: []string{}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + task4 := &task{ + ID: "4", + Subject: "Task 4", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task4JSON, _ := sonic.MarshalString(task4) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "4.json"), Content: task4JSON}) + + task5 := &task{ + ID: "5", + Subject: "Task 5", + Description: "Test description", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task5JSON, _ := sonic.MarshalString(task5) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "5.json"), Content: task5JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["2", "4", "4"]}`) + assert.NoError(t, err) + + content, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updated task + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, []string{"2", "4"}, updated.Blocks) + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["3", "5", "5"]}`) + assert.NoError(t, err) + + content, _ = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + _ = sonic.UnmarshalString(content.Content, &updated) + assert.Equal(t, []string{"3", "5"}, updated.BlockedBy) +} + +func TestTaskUpdateToolBidirectionalBlocks(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Second task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Third task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["2", "3"]}`) + assert.NoError(t, err) + + content1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updatedTask1 task + _ = sonic.UnmarshalString(content1.Content, &updatedTask1) + assert.Equal(t, []string{"2", "3"}, updatedTask1.Blocks) + assert.Empty(t, updatedTask1.BlockedBy) + + content2, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "2.json")}) + var updatedTask2 task + _ = sonic.UnmarshalString(content2.Content, &updatedTask2) + assert.Empty(t, updatedTask2.Blocks) + assert.Equal(t, []string{"1"}, updatedTask2.BlockedBy) + + content3, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "3.json")}) + var updatedTask3 task + _ = sonic.UnmarshalString(content3.Content, &updatedTask3) + assert.Empty(t, updatedTask3.Blocks) + assert.Equal(t, []string{"1"}, updatedTask3.BlockedBy) +} + +func TestTaskUpdateToolBidirectionalBlockedBy(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Second task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Third task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "3", "addBlockedBy": ["1", "2"]}`) + assert.NoError(t, err) + + content3, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "3.json")}) + var updatedTask3 task + _ = sonic.UnmarshalString(content3.Content, &updatedTask3) + assert.Empty(t, updatedTask3.Blocks) + assert.Equal(t, []string{"1", "2"}, updatedTask3.BlockedBy) + + content1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updatedTask1 task + _ = sonic.UnmarshalString(content1.Content, &updatedTask1) + assert.Equal(t, []string{"3"}, updatedTask1.Blocks) + assert.Empty(t, updatedTask1.BlockedBy) + + content2, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "2.json")}) + var updatedTask2 task + _ = sonic.UnmarshalString(content2.Content, &updatedTask2) + assert.Equal(t, []string{"3"}, updatedTask2.Blocks) + assert.Empty(t, updatedTask2.BlockedBy) +} + +func TestTaskUpdateToolBidirectionalWithNonExistentTask(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["999"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "update Task #1 blocks failed") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["999"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "update Task #1 blockedBy failed") +} + +func TestTaskUpdateToolCyclicDependencyDetection(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Second task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Third task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["1"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cyclic dependency") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["1"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cyclic dependency") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlocks": ["2"]}`) + assert.NoError(t, err) + + _, err = tool.InvokableRun(ctx, `{"taskId": "2", "addBlocks": ["1"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cyclic dependency") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["2"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cyclic dependency") + + _, err = tool.InvokableRun(ctx, `{"taskId": "2", "addBlocks": ["3"]}`) + assert.NoError(t, err) + + _, err = tool.InvokableRun(ctx, `{"taskId": "3", "addBlocks": ["1"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cyclic dependency") + + _, err = tool.InvokableRun(ctx, `{"taskId": "1", "addBlockedBy": ["3"]}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "cyclic dependency") + + content1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updatedTask1 task + _ = sonic.UnmarshalString(content1.Content, &updatedTask1) + assert.Equal(t, []string{"2"}, updatedTask1.Blocks) + assert.Empty(t, updatedTask1.BlockedBy) + + content2, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "2.json")}) + var updatedTask2 task + _ = sonic.UnmarshalString(content2.Content, &updatedTask2) + assert.Equal(t, []string{"3"}, updatedTask2.Blocks) + assert.Equal(t, []string{"1"}, updatedTask2.BlockedBy) + + content3, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "3.json")}) + var updatedTask3 task + _ = sonic.UnmarshalString(content3.Content, &updatedTask3) + assert.Empty(t, updatedTask3.Blocks) + assert.Equal(t, []string{"2"}, updatedTask3.BlockedBy) +} + +func TestTaskUpdateToolDeleteCleansDependencies(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusPending, + Blocks: []string{"2", "3"}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Second task", + Status: taskStatusPending, + Blocks: []string{"3"}, + BlockedBy: []string{"1"}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Third task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{"1", "2"}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + result, err := tool.InvokableRun(ctx, `{"taskId": "1", "status": "deleted"}`) + assert.NoError(t, err) + assert.Contains(t, result, "deleted") + + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.Error(t, err) + + content2, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "2.json")}) + assert.NoError(t, err) + var updatedTask2 task + _ = sonic.UnmarshalString(content2.Content, &updatedTask2) + assert.Equal(t, []string{"3"}, updatedTask2.Blocks) + assert.Empty(t, updatedTask2.BlockedBy) + + content3, err := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "3.json")}) + assert.NoError(t, err) + var updatedTask3 task + _ = sonic.UnmarshalString(content3.Content, &updatedTask3) + assert.Empty(t, updatedTask3.Blocks) + assert.Equal(t, []string{"2"}, updatedTask3.BlockedBy) +} + +func TestTaskUpdateToolAutoDeleteAllTasksWhenAllCompleted(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusCompleted, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Second task", + Status: taskStatusCompleted, + Blocks: []string{}, + BlockedBy: []string{}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + task3 := &task{ + ID: "3", + Subject: "Task 3", + Description: "Third task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task3JSON, _ := sonic.MarshalString(task3) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "3.json"), Content: task3JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "3", "status": "completed"}`) + assert.NoError(t, err) + + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.Error(t, err) + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "2.json")}) + assert.Error(t, err) + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "3.json")}) + assert.Error(t, err) +} + +func TestTaskUpdateToolNoDeleteWhenNotAllCompleted(t *testing.T) { + ctx := context.Background() + backend := newInMemoryBackend() + baseDir := "/tmp/tasks" + lock := &sync.Mutex{} + + task1 := &task{ + ID: "1", + Subject: "Task 1", + Description: "First task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task1JSON, _ := sonic.MarshalString(task1) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "1.json"), Content: task1JSON}) + + task2 := &task{ + ID: "2", + Subject: "Task 2", + Description: "Second task", + Status: taskStatusPending, + Blocks: []string{}, + BlockedBy: []string{}, + } + task2JSON, _ := sonic.MarshalString(task2) + _ = backend.Write(ctx, &WriteRequest{FilePath: filepath.Join(baseDir, "2.json"), Content: task2JSON}) + + tool := newTaskUpdateTool(backend, baseDir, lock) + + _, err := tool.InvokableRun(ctx, `{"taskId": "1", "status": "completed"}`) + assert.NoError(t, err) + + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + assert.NoError(t, err) + _, err = backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "2.json")}) + assert.NoError(t, err) + + content1, _ := backend.Read(ctx, &ReadRequest{FilePath: filepath.Join(baseDir, "1.json")}) + var updatedTask1 task + _ = sonic.UnmarshalString(content1.Content, &updatedTask1) + assert.Equal(t, taskStatusCompleted, updatedTask1.Status) +} diff --git a/adk/middlewares/reduction/consts.go b/adk/middlewares/reduction/consts.go new file mode 100644 index 0000000..6f76588 --- /dev/null +++ b/adk/middlewares/reduction/consts.go @@ -0,0 +1,122 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package reduction provides middlewares to trim context and clear tool results. +package reduction + +import ( + "fmt" + + "github.com/cloudwego/eino/adk/internal" +) + +const ( + truncFmt = ` +Output too large ({original_size}). Full output saved to: {file_path} +Preview (first {preview_size}): +{preview_first} + +Preview (last {preview_size}): +{preview_last} + +Use {read_tool_name} to view +` + truncFmtZh = ` +输出结果过大 ({original_size}). 完整输出保存到: {file_path} +预览 (前 {preview_size}): +{preview_first} + +预览 (后 {preview_size}): +{preview_last} + +使用 {read_tool_name} 进行查看 +` +) + +const ( + streamTruncFmt = ` +Output truncated after {preview_size} bytes were streamed. {offload_notify} {error_msg_notify} +` + streamTruncFmtZh = ` +输出结果在流式传输 {preview_size} 字节后被截断。{offload_notify} {error_msg_notify} +` +) + +const ( + clearWithOffloadingFmt = `Tool result saved to: {file_path} +Use {read_tool_name} to view` + clearWithOffloadingFmtZh = `工具结果已保存至: {file_path} +使用 {read_tool_name} 进行查看` + + clearWithoutOffloadingFmt = `[Old tool result content cleared]` + clearWithoutOffloadingFmtZh = `[工具输出结果已清理]` +) + +const ( + msgClearedFlag = "_reduction_mw_processed" +) + +func getTruncFmt() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: truncFmt, + Chinese: truncFmtZh, + }) +} + +func getStreamTruncFmt() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: streamTruncFmt, + Chinese: streamTruncFmtZh, + }) +} + +func formatStreamOffloadSavedNotify(filePath, readFileToolName string) string { + return fmt.Sprintf(internal.SelectPrompt(internal.I18nPrompts{ + English: "Full output saved to: %s. Use %s to view.", + Chinese: "完整输出保存到: %s。使用 %s 进行查看。", + }), filePath, readFileToolName) +} + +func formatStreamOffloadFailedNotify(err error) string { + if err == nil { + return "" + } + return fmt.Sprintf(internal.SelectPrompt(internal.I18nPrompts{ + English: "Failed to save full output: %v.", + Chinese: "完整输出保存失败: %v。", + }), err) +} + +func getClearWithOffloadingFmt() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: clearWithOffloadingFmt, + Chinese: clearWithOffloadingFmtZh, + }) +} + +func getClearWithoutOffloadingFmt() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: clearWithoutOffloadingFmt, + Chinese: clearWithoutOffloadingFmtZh, + }) +} + +type scene int + +const ( + sceneTruncation scene = 1 + sceneClear scene = 2 +) diff --git a/adk/middlewares/reduction/internal/clear_tool_result.go b/adk/middlewares/reduction/internal/clear_tool_result.go new file mode 100644 index 0000000..224e1fb --- /dev/null +++ b/adk/middlewares/reduction/internal/clear_tool_result.go @@ -0,0 +1,165 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package internal provides middlewares to trim context and clear tool results. +package internal + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// ClearToolResultConfig configures the tool result clearing middleware. +// This middleware clears old tool results when their total token count exceeds a threshold, +// while protecting recent messages within a token budget. +type ClearToolResultConfig struct { + // ToolResultTokenThreshold is the threshold for total tool result tokens. + // When the sum of all tool result tokens exceeds this threshold, old tool results + // (outside the KeepRecentTokens range) will be replaced with a placeholder. + // Token estimation uses a simple heuristic: character count / 4. + // If 0, defaults to 20000. + ToolResultTokenThreshold int + + // KeepRecentTokens is the token budget for recent messages to keep intact. + // Messages within this token budget from the end will not have their tool results cleared, + // even if the total tool result tokens exceed the threshold. + // If 0, defaults to 40000. + KeepRecentTokens int + + // ClearToolResultPlaceholder is the text to replace old tool results with. + // If empty, defaults to "[Old tool result content cleared]". + ClearToolResultPlaceholder string + + // TokenCounter is a custom function to estimate token count for a message. + // If nil, uses the default counter (character count / 4). + TokenCounter func(msg *schema.Message) int + + // ExcludeTools is a list of tool names whose results should never be cleared. + ExcludeTools []string +} + +// NewClearToolResult creates a new middleware that clears old tool results +// based on token thresholds while protecting recent messages. +func NewClearToolResult(ctx context.Context, config *ClearToolResultConfig) (adk.AgentMiddleware, error) { + return adk.AgentMiddleware{ + BeforeChatModel: newClearToolResult(ctx, config), + }, nil +} + +func newClearToolResult(ctx context.Context, config *ClearToolResultConfig) func(ctx context.Context, state *adk.ChatModelAgentState) error { + if config == nil { + config = &ClearToolResultConfig{} + } + + // Set defaults + toolResultTokenThreshold := config.ToolResultTokenThreshold + if toolResultTokenThreshold == 0 { + toolResultTokenThreshold = 20000 + } + + keepRecentTokens := config.KeepRecentTokens + if keepRecentTokens == 0 { + keepRecentTokens = 40000 + } + + placeholder := config.ClearToolResultPlaceholder + if placeholder == "" { + placeholder = "[Old tool result content cleared]" + } + + // Set token estimator + counter := config.TokenCounter + if counter == nil { + counter = defaultTokenCounter + } + return func(ctx context.Context, state *adk.ChatModelAgentState) error { + return reduceByTokens(state, toolResultTokenThreshold, keepRecentTokens, placeholder, counter, config.ExcludeTools) + } +} + +// defaultTokenCounter estimates token count using character count / 4 +// This is a simple heuristic that works reasonably well for most languages +func defaultTokenCounter(msg *schema.Message) int { + count := len(msg.Content) + + // Also count tool call arguments if present + for _, tc := range msg.ToolCalls { + count += len(tc.Function.Arguments) + } + + // Estimate: roughly 4 characters per token + return (count + 3) / 4 +} + +// reduceByTokens reduces context based on tool result token threshold and recent message protection. +// It clears old tool results when: +// 1. The total tokens of all tool results exceed toolResultTokenThreshold +// 2. Only tool results outside the keepRecentTokens range (from the end) are cleared +func reduceByTokens(state *adk.ChatModelAgentState, toolResultTokenThreshold, keepRecentTokens int, placeholder string, counter func(*schema.Message) int, excludedTools []string) error { + if len(state.Messages) == 0 { + return nil + } + + // Step 1: Calculate total tool result tokens + totalToolResultTokens := 0 + for _, msg := range state.Messages { + if msg.Role == schema.Tool && msg.Content != placeholder { + totalToolResultTokens += counter(msg) + } + } + + // If total tool result tokens are under the threshold, no reduction needed + if totalToolResultTokens <= toolResultTokenThreshold { + return nil + } + + // Step 2: Calculate the index from which to protect recent messages + // We need to find the starting index where cumulative tokens from the end <= keepRecentTokens + recentStartIdx := len(state.Messages) + cumulativeTokens := 0 + + for i := len(state.Messages) - 1; i >= 0; i-- { + msgTokens := counter(state.Messages[i]) + if cumulativeTokens+msgTokens > keepRecentTokens { + // Adding this message would exceed the budget, so stop here + recentStartIdx = i + break + } + cumulativeTokens += msgTokens + recentStartIdx = i + } + + // Step 3: Clear tool results outside the protected range (before recentStartIdx) + for i := 0; i < recentStartIdx; i++ { + msg := state.Messages[i] + if msg.Role == schema.Tool && msg.Content != placeholder && !excluded(msg.ToolName, excludedTools) { + msg.Content = placeholder + } + } + + return nil +} + +func excluded(name string, exclude []string) bool { + for _, ex := range exclude { + if name == ex { + return true + } + } + return false +} diff --git a/adk/middlewares/reduction/internal/clear_tool_result_test.go b/adk/middlewares/reduction/internal/clear_tool_result_test.go new file mode 100644 index 0000000..ff9575f --- /dev/null +++ b/adk/middlewares/reduction/internal/clear_tool_result_test.go @@ -0,0 +1,302 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func Test_reduceByTokens(t *testing.T) { + type args struct { + state *adk.ChatModelAgentState + toolResultTokenThreshold int + keepRecentTokens int + placeholder string + estimator func(*schema.Message) int + } + tests := []struct { + name string + args args + wantErr assert.ErrorAssertionFunc + validateState func(*testing.T, *adk.ChatModelAgentState) + }{ + { + name: "no reduction when tool result tokens under threshold", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("hello"), + schema.AssistantMessage("hi", nil), + schema.ToolMessage("short tool result", "call-1", schema.WithToolName("tool1")), + }, + }, + toolResultTokenThreshold: 100, + keepRecentTokens: 500, + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + assert.Equal(t, "short tool result", state.Messages[2].Content) + }, + }, + { + name: "clear old tool results when total exceeds threshold", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("msg1"), + schema.ToolMessage(strings.Repeat("a", 40), "call-1", schema.WithToolName("tool1")), // ~10 tokens (old) + schema.UserMessage("msg2"), + schema.ToolMessage(strings.Repeat("b", 40), "call-2", schema.WithToolName("tool2")), // ~10 tokens (old) + schema.UserMessage("msg3"), + schema.ToolMessage(strings.Repeat("c", 40), "call-3", schema.WithToolName("tool3")), // ~10 tokens (recent, protected) + }, + }, + toolResultTokenThreshold: 20, + keepRecentTokens: 10, + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + assert.Equal(t, "[Old tool result content cleared]", state.Messages[1].Content) + assert.Equal(t, "[Old tool result content cleared]", state.Messages[3].Content) + assert.Equal(t, strings.Repeat("c", 40), state.Messages[5].Content) + }, + }, + { + name: "protect recent messages even when tool results exceed threshold", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("old msg"), + schema.ToolMessage(strings.Repeat("x", 100), "call-1", schema.WithToolName("tool1")), // ~25 tokens (old) + schema.UserMessage("recent msg"), + schema.ToolMessage(strings.Repeat("x", 100), "call-2", schema.WithToolName("tool2")), // ~25 tokens (recent, protected) + }, + }, + toolResultTokenThreshold: 10, + keepRecentTokens: 20, + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + // Total tool result tokens = 50, exceeds threshold of 10 + // But last 200 tokens are protected (includes last 2 messages) + // So only the first tool result should be cleared + assert.Equal(t, "[Old tool result content cleared]", state.Messages[1].Content) + assert.Equal(t, strings.Repeat("x", 100), state.Messages[3].Content) + }, + }, + { + name: "custom placeholder text", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("msg"), + schema.ToolMessage(strings.Repeat("x", 100), "call-1", schema.WithToolName("tool1")), + schema.UserMessage(strings.Repeat("x", 100)), + }, + }, + toolResultTokenThreshold: 10, + keepRecentTokens: 20, + placeholder: "[历史工具结果已清除]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + assert.Equal(t, "[历史工具结果已清除]", state.Messages[1].Content) + }, + }, + { + name: "no tool messages", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("msg 1"), + schema.AssistantMessage("response 1", nil), + schema.UserMessage("msg 2"), + schema.AssistantMessage("response 2", nil), + }, + }, + toolResultTokenThreshold: 10, + keepRecentTokens: 10, + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + // All messages should remain unchanged + assert.Equal(t, "msg 1", state.Messages[0].Content) + assert.Equal(t, "response 1", state.Messages[1].Content) + assert.Equal(t, "msg 2", state.Messages[2].Content) + assert.Equal(t, "response 2", state.Messages[3].Content) + }, + }, + { + name: "empty messages", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{}, + }, + toolResultTokenThreshold: 100, + keepRecentTokens: 500, + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + assert.Empty(t, state.Messages) + }, + }, + { + name: "custom token estimator - word count", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("hello world"), + schema.ToolMessage("this is a long tool result", "call-1", schema.WithToolName("tool1")), // 6 words (old) + schema.UserMessage("another message"), + schema.ToolMessage("recent tool result here", "call-2", schema.WithToolName("tool2")), // 4 words (recent) + }, + }, + toolResultTokenThreshold: 9, // 10 words total threshold + keepRecentTokens: 5, // 15 words protection budget + placeholder: "[Old tool result content cleared]", + estimator: func(msg *schema.Message) int { + if msg.Content == "" { + return 0 + } + words := 1 + for _, ch := range msg.Content { + if ch == ' ' { + words++ + } + } + return words + }, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + assert.Equal(t, "[Old tool result content cleared]", state.Messages[1].Content) + assert.Equal(t, "recent tool result here", state.Messages[3].Content) + }, + }, + { + name: "already cleared results are not counted", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("msg1"), + schema.ToolMessage("[Old tool result content cleared]", "call-1", schema.WithToolName("tool1")), // Already cleared + schema.UserMessage("msg2"), + schema.ToolMessage(strings.Repeat("a", 100), "call-2", schema.WithToolName("tool2")), // New long result + }, + }, + toolResultTokenThreshold: 10, + keepRecentTokens: 20, + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + // Only the new long result counts toward the threshold + // Both should have placeholder + assert.Equal(t, "[Old tool result content cleared]", state.Messages[1].Content) + assert.Equal(t, strings.Repeat("a", 100), state.Messages[3].Content) + }, + }, + { + name: "all tool results within protected range", + args: args{ + state: &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("msg1"), + schema.ToolMessage(strings.Repeat("a", 40), "call-1", schema.WithToolName("tool1")), // ~10 tokens + schema.UserMessage("msg2"), + schema.ToolMessage(strings.Repeat("b", 40), "call-2", schema.WithToolName("tool2")), // ~10 tokens + }, + }, + toolResultTokenThreshold: 10, // Low threshold (will exceed) + keepRecentTokens: 1000, // Very high protection (protects all) + placeholder: "[Old tool result content cleared]", + estimator: defaultTokenCounter, + }, + wantErr: assert.NoError, + validateState: func(t *testing.T, state *adk.ChatModelAgentState) { + // All messages are within protected range, nothing should be cleared + assert.Equal(t, strings.Repeat("a", 40), state.Messages[1].Content) + assert.Equal(t, strings.Repeat("b", 40), state.Messages[3].Content) + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := reduceByTokens(tt.args.state, tt.args.toolResultTokenThreshold, tt.args.keepRecentTokens, tt.args.placeholder, tt.args.estimator, []string{}) + tt.wantErr(t, err, fmt.Sprintf("reduceByTokens(%v, %v, %v, %v)", tt.args.state, tt.args.toolResultTokenThreshold, tt.args.keepRecentTokens, tt.args.placeholder)) + if tt.validateState != nil { + tt.validateState(t, tt.args.state) + } + }) + } +} + +func Test_newClearToolResult(t *testing.T) { + ctx := context.Background() + + t.Run("nil config uses defaults", func(t *testing.T) { + fn := newClearToolResult(ctx, nil) + assert.NotNil(t, fn) + + // Test that function works with nil config (uses defaults) + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("hello"), + schema.ToolMessage("short result", "call-1", schema.WithToolName("tool1")), + }, + } + err := fn(ctx, state) + assert.NoError(t, err) + // Default threshold is 20000, so short result should not be cleared + assert.Equal(t, "short result", state.Messages[1].Content) + }) + + t.Run("empty config uses defaults", func(t *testing.T) { + fn := newClearToolResult(ctx, &ClearToolResultConfig{}) + assert.NotNil(t, fn) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("hello"), + schema.ToolMessage("short result", "call-1", schema.WithToolName("tool1")), + }, + } + err := fn(ctx, state) + assert.NoError(t, err) + assert.Equal(t, "short result", state.Messages[1].Content) + }) +} diff --git a/adk/middlewares/reduction/internal/large_tool_result.go b/adk/middlewares/reduction/internal/large_tool_result.go new file mode 100644 index 0000000..c1e41b1 --- /dev/null +++ b/adk/middlewares/reduction/internal/large_tool_result.go @@ -0,0 +1,210 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "strings" + "unicode/utf8" + + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +const ( + tooLargeToolMessage = `Tool result too large, the result of this tool call {tool_call_id} was saved in the filesystem at this path: {file_path} +You can read the result from the filesystem by using the '{read_file_tool_name}' tool, but make sure to only read part of the result at a time. +You can do this by specifying an offset and limit in the '{read_file_tool_name}' tool call. +For example, to read the first 100 lines, you can use the '{read_file_tool_name}' tool with offset=0 and limit=100. + +Here are the first 10 lines of the result: +{content_sample}` + + tooLargeToolMessageChinese = `工具结果过大，此工具调用 {tool_call_id} 的结果已保存到文件系统的以下路径：{file_path} +你可以使用 '{read_file_tool_name}' 工具从文件系统读取结果，但请确保每次只读取部分结果。 +你可以通过在 '{read_file_tool_name}' 工具调用中指定 offset 和 limit 来实现。 +例如，要读取前 100 行，你可以使用 '{read_file_tool_name}' 工具，设置 offset=0 和 limit=100。 + +以下是结果的前 10 行： +{content_sample}` +) + +type toolResultOffloadingConfig struct { + Backend Backend + ReadFileToolName string + TokenLimit int + PathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error) + TokenCounter func(msg *schema.Message) int +} + +func newToolResultOffloading(_ context.Context, config *toolResultOffloadingConfig) compose.ToolMiddleware { + offloading := &toolResultOffloading{ + backend: config.Backend, + tokenLimit: config.TokenLimit, + pathGenerator: config.PathGenerator, + toolName: config.ReadFileToolName, + counter: config.TokenCounter, + } + + if offloading.tokenLimit == 0 { + offloading.tokenLimit = 20000 + } + + if offloading.pathGenerator == nil { + offloading.pathGenerator = func(ctx context.Context, input *compose.ToolInput) (string, error) { + return fmt.Sprintf("/large_tool_result/%s", input.CallID), nil + } + } + + if len(offloading.toolName) == 0 { + offloading.toolName = "read_file" + } + + if offloading.counter == nil { + offloading.counter = defaultTokenCounter + } + + return compose.ToolMiddleware{ + Invokable: offloading.invoke, + Streamable: offloading.stream, + } +} + +type toolResultOffloading struct { + backend Backend + tokenLimit int + pathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error) + toolName string + counter func(msg *schema.Message) int +} + +func (t *toolResultOffloading) invoke(endpoint compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + output, err := endpoint(ctx, input) + if err != nil { + return nil, err + } + result, err := t.handleResult(ctx, output.Result, input) + if err != nil { + return nil, err + } + return &compose.ToolOutput{Result: result}, nil + } +} + +func (t *toolResultOffloading) stream(endpoint compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + output, err := endpoint(ctx, input) + if err != nil { + return nil, err + } + result, err := concatString(output.Result) + if err != nil { + return nil, err + } + result, err = t.handleResult(ctx, result, input) + if err != nil { + return nil, err + } + return &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{result})}, nil + } +} + +func (t *toolResultOffloading) handleResult(ctx context.Context, result string, input *compose.ToolInput) (string, error) { + if t.counter(schema.ToolMessage(result, input.CallID, schema.WithToolName(input.Name))) > t.tokenLimit*4 { + path, err := t.pathGenerator(ctx, input) + if err != nil { + return "", err + } + + nResult := formatToolMessage(result) + tpl := internal.SelectPrompt(internal.I18nPrompts{ + English: tooLargeToolMessage, + Chinese: tooLargeToolMessageChinese, + }) + nResult, err = pyfmt.Fmt(tpl, map[string]any{ + "tool_call_id": input.CallID, + "file_path": path, + "content_sample": nResult, + "read_file_tool_name": t.toolName, + }) + if err != nil { + return "", err + } + + err = t.backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: path, + Content: result, + }) + if err != nil { + return "", err + } + + return nResult, nil + } + + return result, nil +} + +func concatString(sr *schema.StreamReader[string]) (string, error) { + if sr == nil { + return "", errors.New("stream is nil") + } + sb := strings.Builder{} + for { + str, err := sr.Recv() + if errors.Is(err, io.EOF) { + return sb.String(), nil + } + if err != nil { + return "", err + } + sb.WriteString(str) + } +} + +func formatToolMessage(s string) string { + reader := bufio.NewScanner(strings.NewReader(s)) + var b strings.Builder + + lineNum := 1 + for reader.Scan() { + if lineNum > 10 { + break + } + line := reader.Text() + + if utf8.RuneCountInString(line) > 1000 { + runes := []rune(line) + line = string(runes[:1000]) + } + + b.WriteString(fmt.Sprintf("%d: %s\n", lineNum, line)) + + lineNum++ + } + + return b.String() +} diff --git a/adk/middlewares/reduction/internal/large_tool_result_test.go b/adk/middlewares/reduction/internal/large_tool_result_test.go new file mode 100644 index 0000000..fafd004 --- /dev/null +++ b/adk/middlewares/reduction/internal/large_tool_result_test.go @@ -0,0 +1,568 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + "testing" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// mockBackend is a simple in-memory backend for testing +type mockBackend struct { + files map[string]string +} + +func newMockBackend() *mockBackend { + return &mockBackend{ + files: make(map[string]string), + } +} + +func (m *mockBackend) Write(_ context.Context, wr *filesystem.WriteRequest) error { + m.files[wr.FilePath] = wr.Content + return nil +} + +func TestToolResultOffloading_SmallResult(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 100, // Small limit for testing + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a mock endpoint that returns a small result + smallResult := "This is a small result" + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: smallResult}, nil + } + + // Wrap the endpoint with the middleware + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + // Execute + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_123", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Small result should pass through unchanged + if output.Result != smallResult { + t.Errorf("expected result %q, got %q", smallResult, output.Result) + } + + // No file should be written + if len(backend.files) != 0 { + t.Errorf("expected no files to be written, got %d files", len(backend.files)) + } +} + +func TestToolResultOffloading_LargeResult(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, // Very small limit to trigger offloading + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a large result (more than 10 * 4 = 40 bytes) + largeResult := strings.Repeat("This is a long line of text that will exceed the token limit.\n", 10) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_456", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Result should be replaced with a message + if !strings.Contains(output.Result, "Tool result too large") { + t.Errorf("expected result to contain 'Tool result too large', got %q", output.Result) + } + + if !strings.Contains(output.Result, "call_456") { + t.Errorf("expected result to contain call ID 'call_456', got %q", output.Result) + } + + if !strings.Contains(output.Result, "/large_tool_result/call_456") { + t.Errorf("expected result to contain file path, got %q", output.Result) + } + + // File should be written + if len(backend.files) != 1 { + t.Fatalf("expected 1 file to be written, got %d files", len(backend.files)) + } + + savedContent, ok := backend.files["/large_tool_result/call_456"] + if !ok { + t.Fatalf("expected file at /large_tool_result/call_456, got files: %v", backend.files) + } + + if savedContent != largeResult { + t.Errorf("saved content doesn't match original result") + } +} + +func TestToolResultOffloading_CustomPathGenerator(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + customPath := "/custom/path/result.txt" + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + PathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) { + return customPath, nil + }, + } + + middleware := newToolResultOffloading(ctx, config) + + largeResult := strings.Repeat("Large content ", 100) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_789", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Check custom path is used + if !strings.Contains(output.Result, customPath) { + t.Errorf("expected result to contain custom path %q, got %q", customPath, output.Result) + } + + // File should be written to custom path + savedContent, ok := backend.files[customPath] + if !ok { + t.Fatalf("expected file at %q, got files: %v", customPath, backend.files) + } + + if savedContent != largeResult { + t.Errorf("saved content doesn't match original result") + } +} + +func TestToolResultOffloading_PathGeneratorError(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + expectedErr := errors.New("path generation failed") + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + PathGenerator: func(ctx context.Context, input *compose.ToolInput) (string, error) { + return "", expectedErr + }, + } + + middleware := newToolResultOffloading(ctx, config) + + largeResult := strings.Repeat("Large content ", 100) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } +} + +func TestToolResultOffloading_EndpointError(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 100, + } + + middleware := newToolResultOffloading(ctx, config) + + expectedErr := errors.New("endpoint execution failed") + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return nil, expectedErr + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_endpoint_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } +} + +func TestToolResultOffloading_DefaultTokenLimit(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 0, // Should default to 20000 + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a result smaller than 20000 * 4 = 80000 bytes + smallResult := strings.Repeat("x", 1000) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: smallResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_default", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Should pass through unchanged + if output.Result != smallResult { + t.Errorf("expected result to pass through unchanged") + } + + // No file should be written + if len(backend.files) != 0 { + t.Errorf("expected no files to be written, got %d files", len(backend.files)) + } +} + +func TestToolResultOffloading_Stream(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + } + + middleware := newToolResultOffloading(ctx, config) + + // Create a streaming endpoint that returns large content + largeResult := strings.Repeat("Large streaming content ", 100) + mockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + // Split the result into chunks + chunks := []string{largeResult[:len(largeResult)/2], largeResult[len(largeResult)/2:]} + return &compose.StreamToolOutput{ + Result: schema.StreamReaderFromArray(chunks), + }, nil + } + + wrappedEndpoint := middleware.Streamable(mockStreamEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_stream", + } + output, err := wrappedEndpoint(ctx, input) + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Read the stream + var result strings.Builder + for { + chunk, err := output.Result.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("error reading stream: %v", err) + } + result.WriteString(chunk) + } + + resultStr := result.String() + + // Result should be replaced with a message + if !strings.Contains(resultStr, "Tool result too large") { + t.Errorf("expected result to contain 'Tool result too large', got %q", resultStr) + } + + if !strings.Contains(resultStr, "call_stream") { + t.Errorf("expected result to contain call ID 'call_stream', got %q", resultStr) + } + + // File should be written + if len(backend.files) != 1 { + t.Fatalf("expected 1 file to be written, got %d files", len(backend.files)) + } + + savedContent, ok := backend.files["/large_tool_result/call_stream"] + if !ok { + t.Fatalf("expected file at /large_tool_result/call_stream, got files: %v", backend.files) + } + + if savedContent != largeResult { + t.Errorf("saved content doesn't match original result") + } +} + +func TestToolResultOffloading_StreamError(t *testing.T) { + ctx := context.Background() + backend := newMockBackend() + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + } + + middleware := newToolResultOffloading(ctx, config) + + expectedErr := errors.New("stream endpoint failed") + mockStreamEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return nil, expectedErr + } + + wrappedEndpoint := middleware.Streamable(mockStreamEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_stream_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } +} + +func TestFormatToolMessage(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "single line", + input: "single line", + expected: "1: single line\n", + }, + { + name: "multiple lines", + input: "line1\nline2\nline3", + expected: "1: line1\n2: line2\n3: line3\n", + }, + { + name: "more than 10 lines", + input: "1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12", + expected: "1: 1\n2: 2\n3: 3\n4: 4\n5: 5\n6: 6\n7: 7\n8: 8\n9: 9\n10: 10\n", + }, + { + name: "long line truncation", + input: strings.Repeat("a", 1500), + expected: fmt.Sprintf("1: %s\n", strings.Repeat("a", 1000)), + }, + { + name: "unicode characters", + input: "你好世界\n测试", + expected: "1: 你好世界\n2: 测试\n", + }, + { + name: "long unicode line", + input: strings.Repeat("你", 1500), + expected: fmt.Sprintf("1: %s\n", strings.Repeat("你", 1000)), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := formatToolMessage(tt.input) + if result != tt.expected { + t.Errorf("formatToolMessage() = %q, want %q", result, tt.expected) + } + }) + } +} + +func TestConcatString(t *testing.T) { + tests := []struct { + name string + chunks []string + expected string + expectError bool + }{ + { + name: "single chunk", + chunks: []string{"hello"}, + expected: "hello", + }, + { + name: "multiple chunks", + chunks: []string{"hello", " ", "world"}, + expected: "hello world", + }, + { + name: "empty chunks", + chunks: []string{"", "", ""}, + expected: "", + }, + { + name: "mixed chunks", + chunks: []string{"a", "", "b", "c"}, + expected: "abc", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := schema.StreamReaderFromArray(tt.chunks) + result, err := concatString(sr) + + if tt.expectError { + if err == nil { + t.Error("expected error, got nil") + } + return + } + + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if result != tt.expected { + t.Errorf("concatString() = %q, want %q", result, tt.expected) + } + }) + } + + // Test nil stream + t.Run("nil stream", func(t *testing.T) { + _, err := concatString(nil) + if err == nil { + t.Error("expected error for nil stream, got nil") + } + if !strings.Contains(err.Error(), "stream is nil") { + t.Errorf("expected 'stream is nil' error, got %v", err) + } + }) +} + +func TestToolResultOffloading_BackendWriteError(t *testing.T) { + ctx := context.Background() + + // Create a backend that fails on write + backend := &failingBackend{ + writeErr: errors.New("write failed"), + } + + config := &toolResultOffloadingConfig{ + Backend: backend, + TokenLimit: 10, + } + + middleware := newToolResultOffloading(ctx, config) + + largeResult := strings.Repeat("Large content ", 100) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: largeResult}, nil + } + + wrappedEndpoint := middleware.Invokable(mockEndpoint) + + input := &compose.ToolInput{ + Name: "test_tool", + CallID: "call_write_error", + } + _, err := wrappedEndpoint(ctx, input) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !strings.Contains(err.Error(), "write failed") { + t.Errorf("expected 'write failed' error, got %v", err) + } +} + +// failingBackend is a mock backend that can be configured to fail +type failingBackend struct { + writeErr error +} + +func (f *failingBackend) Write(context.Context, *filesystem.WriteRequest) error { + if f.writeErr != nil { + return f.writeErr + } + return nil +} diff --git a/adk/middlewares/reduction/internal/tool_result.go b/adk/middlewares/reduction/internal/tool_result.go new file mode 100644 index 0000000..ec14ba9 --- /dev/null +++ b/adk/middlewares/reduction/internal/tool_result.go @@ -0,0 +1,121 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// Backend defines the interface provided by the user to implement file storage. +// It is used to save the content of large tool results to a persistent storage. +type Backend interface { + Write(context.Context, *filesystem.WriteRequest) error +} + +// ToolResultConfig configures the tool result reduction middleware. +type ToolResultConfig struct { + // ClearingTokenThreshold is the threshold for the total token count of all tool results. + // When the sum of all tool result tokens exceeds this threshold, old tool results + // (outside the KeepRecentTokens range) will be replaced with a placeholder. + // Token estimation uses a simple heuristic: character count / 4. + // optional, 20000 by default + ClearingTokenThreshold int + + // KeepRecentTokens is the token budget for recent messages to keep intact. + // Messages within this token budget from the end will not have their tool results cleared, + // even if the total tool result tokens exceed the threshold. + // optional, 40000 by default + KeepRecentTokens int + + // ClearToolResultPlaceholder is the text to replace old tool results with. + // optional, "[Old tool result content cleared]" by default + ClearToolResultPlaceholder string + + // TokenCounter is a custom function to estimate token count for a message. + // optional, uses the default counter (character count / 4) if nil + TokenCounter func(msg *schema.Message) int + + // ExcludeTools is a list of tool names whose results should never be cleared. + // optional + ExcludeTools []string + + // Backend is the storage backend for offloaded tool results. + // required + Backend Backend + + // OffloadingTokenLimit is the token threshold for a single tool result to trigger offloading. + // When a single tool result exceeds OffloadingTokenLimit * 4 characters, it will be + // offloaded to the filesystem. + // optional, 20000 by default + OffloadingTokenLimit int + + // ReadFileToolName is the name of the tool that LLM should use to read offloaded content. + // This name will be included in the summary message sent to the LLM. + // optional, "read_file" by default + // + // NOTE: If you are using the filesystem middleware, the read_file tool name + // is exactly "read_file", which matches the default value. + ReadFileToolName string + + // PathGenerator generates the write path for offloaded results. + // optional, "/large_tool_result/{ToolCallID}" by default + PathGenerator func(ctx context.Context, input *compose.ToolInput) (string, error) +} + +// NewToolResultMiddleware creates a tool result reduction middleware. +// This middleware combines two strategies to manage tool result tokens: +// +// 1. Clearing: Replaces old tool results with a placeholder when the total +// tool result tokens exceed the threshold, while protecting recent messages. +// +// 2. Offloading: Writes large individual tool results to the filesystem and +// returns a summary message guiding the LLM to read the full content. +// +// NOTE: If you are using the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem), +// this functionality is already included by default. Set Config.WithoutLargeToolResultOffloading = true +// in the filesystem middleware if you want to use this middleware separately instead. +// +// NOTE: This middleware only handles offloading results to the filesystem. +// You MUST also provide a read_file tool to your agent, otherwise the agent +// will not be able to read the offloaded content. You can either: +// - Use the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem) +// which provides the read_file tool automatically, OR +// - Implement your own read_file tool that reads from the same Backend +func NewToolResultMiddleware(ctx context.Context, cfg *ToolResultConfig) (adk.AgentMiddleware, error) { + bc := newClearToolResult(ctx, &ClearToolResultConfig{ + ToolResultTokenThreshold: cfg.ClearingTokenThreshold, + KeepRecentTokens: cfg.KeepRecentTokens, + ClearToolResultPlaceholder: cfg.ClearToolResultPlaceholder, + TokenCounter: cfg.TokenCounter, + ExcludeTools: cfg.ExcludeTools, + }) + tm := newToolResultOffloading(ctx, &toolResultOffloadingConfig{ + Backend: cfg.Backend, + ReadFileToolName: cfg.ReadFileToolName, + TokenLimit: cfg.OffloadingTokenLimit, + PathGenerator: cfg.PathGenerator, + }) + return adk.AgentMiddleware{ + BeforeChatModel: bc, + WrapToolCall: tm, + }, nil +} diff --git a/adk/middlewares/reduction/legacy.go b/adk/middlewares/reduction/legacy.go new file mode 100644 index 0000000..36c6d4f --- /dev/null +++ b/adk/middlewares/reduction/legacy.go @@ -0,0 +1,69 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package reduction + +import "github.com/cloudwego/eino/adk/middlewares/reduction/internal" + +// Package reduction provides historical compatibility exports for reduction middleware APIs. +// +// DEPRECATED: All top-level exports in this file are maintained exclusively for backward compatibility. +// New reduction middleware implementations are now developed and maintained in this package. +// It is STRONGLY RECOMMENDED that new code directly use the New instead. +// +// Existing code relying on these exports will continue to work indefinitely, +// but no new features or bug fixes will be backported to this compatibility layer. + +type ( + ClearToolResultConfig = internal.ClearToolResultConfig + ToolResultConfig = internal.ToolResultConfig + Backend = internal.Backend +) + +var ( + // NewClearToolResult creates a new middleware that clears old tool results + // based on token thresholds while protecting recent messages. + // + // Deprecated: Use New instead, which provides a more comprehensive tool result reduction + // middleware with both truncation and clearing strategies. New returns a ChatModelAgentMiddleware + // for better context propagation through wrapper methods. + NewClearToolResult = internal.NewClearToolResult + + // NewToolResultMiddleware creates a tool result reduction middleware. + // This middleware combines two strategies to manage tool result tokens: + // + // 1. Clearing: Replaces old tool results with a placeholder when the total + // tool result tokens exceed the threshold, while protecting recent messages. + // + // 2. Offloading: Writes large individual tool results to the filesystem and + // returns a summary message guiding the LLM to read the full content. + // + // NOTE: If you are using the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem), + // this functionality is already included by default. Set Config.WithoutLargeToolResultOffloading = true + // in the filesystem middleware if you want to use this middleware separately instead. + // + // NOTE: This middleware only handles offloading results to the filesystem. + // You MUST also provide a read_file tool to your agent, otherwise the agent + // will not be able to read the offloaded content. You can either: + // - Use the filesystem middleware (github.com/cloudwego/eino/adk/middlewares/filesystem) + // which provides the read_file tool automatically, OR + // - Implement your own read_file tool that reads from the same Backend + // + // Deprecated: Use New instead, which provides a more comprehensive tool result reduction + // middleware with both truncation and clearing strategies, per-tool configuration support, + // and returns a ChatModelAgentMiddleware for better context propagation through wrapper methods. + NewToolResultMiddleware = internal.NewToolResultMiddleware +) diff --git a/adk/middlewares/reduction/reduction.go b/adk/middlewares/reduction/reduction.go new file mode 100644 index 0000000..6a4203d --- /dev/null +++ b/adk/middlewares/reduction/reduction.go @@ -0,0 +1,2084 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package reduction + +import ( + "context" + "encoding/json" + "fmt" + "io" + "path/filepath" + "strings" + "unicode/utf8" + + "github.com/bytedance/sonic" + "github.com/google/uuid" + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +// TypedConfig is the configuration for tool reduction middleware. +// This middleware manages tool outputs in two phases to optimize context usage: +// +// 1. Truncation Phase: +// Triggered immediately after a tool execution completes. +// If the tool output length exceeds MaxLengthForTrunc, the full content is saved +// to the configured Backend, and the tool output is replaced with a truncated notice. +// This prevents immediate context overflow from a single large tool output. +// +// 2. Clear Phase: +// Triggered before sending messages to the model (in BeforeModelRewriteState). +// If the total token count exceeds MaxTokensForClear, the middleware iterates through +// historical messages. Based on GenOffloadFilePath (or RootDir when GenOffloadFilePath is nil) and +// ClearRetentionSuffixLimit, it offloads tool call arguments and results +// to the Backend to reduce token usage, keeping the conversation within limits while retaining access to the +// important information. After all, ClearPostProcess will be called, which you could save or notify current state. +type TypedConfig[M adk.MessageType] struct { + // Backend is the storage backend where offloaded content will be saved. + // Required when truncation is enabled (SkipTruncation is false). + // Optional for clear-only usage. If Backend is nil, clear will still replace tool outputs with placeholders + // but will not offload content. + Backend Backend + + // SkipTruncation skip truncating. + SkipTruncation bool + + // SkipClear skip clearing. + SkipClear bool + + // ReadFileToolName is tool name used to retrieve from file. + // After offloading content to file, you should give agent the same tool to retrieve content. + // Required. Default is "read_file". + ReadFileToolName string + + // RootDir root dir to save truncated/cleared content. + // Optional. + // Default is /tmp, truncated content saves to ${root_dir}/trunc/{tool_call_id}, cleared content saves to ${root_dir}/clear/{tool_call_id} + RootDir string + + // GenTruncOffloadFilePath is used to generate offload file path for truncated content. + // When GenTruncOffloadFilePath is configured, RootDir will be ignored. + // This is useful when tool_call_id is not unique, which may cause incorrect offload file overwrite. + // Optional. Default is nil. + GenTruncOffloadFilePath func(ctx context.Context, toolDetail *ToolDetail) (filePath string, err error) + + // GenClearOffloadFilePath is used to generate offload file path for truncated content. + // When GenClearOffloadFilePath is configured, RootDir will be ignored. + // This is useful when tool_call_id is not unique, which may cause incorrect offload file overwrite. + // Optional. Default is nil. + GenClearOffloadFilePath func(ctx context.Context, toolDetail *ToolDetail) (filePath string, err error) + + // MaxLengthForTrunc is the maximum allowed length of the tool output. + // If the output exceeds this length, it will be truncated. + // Required. Default is 50000. + MaxLengthForTrunc int + + // TruncExcludeTools is list of tool names whose tool results should never be truncated. + // Optional. Default is nil. + TruncExcludeTools []string + + // TokenCounter is used to count the number of tokens in the conversation messages. + // It is used to determine when to trigger clearing based on token usage, and token usage after clearing. + // Optional. If not provided, a default token counter will be used, which estimates tokens by counting 1 token per 4 characters. + TokenCounter func(ctx context.Context, msg []M, tools []*schema.ToolInfo) (int64, error) + + // MaxTokensForClear is the maximum number of tokens allowed in the conversation before clearing is attempted. + // Required. Default is 160000. + MaxTokensForClear int64 + + // ClearRetentionSuffixLimit is the number of most recent tool-call rounds to retain without clearing. + // A round consists of one assistant message (which may contain multiple tool calls) and its corresponding tool-result messages. + // This ensures the model has immediate context to process pending tool results. + // + // Example with ClearRetentionSuffixLimit = 2, the retained suffix looks like: + // Round 2: assistant [call_A, call_B] → result_A, result_B + // Round 1: assistant [call_C] → result_C + // (Round 1 is the most recent, closest to the end of the message list.) + // + // Optional. Default is 1. + ClearRetentionSuffixLimit int + + // ClearAtLeastTokens ensures a minimum number of tokens is cleared each time the strategy activates. + // If the strategy couldn't clear at least the specified amount, clear phase will not be applied. + // This helps determine if context clearing is worth breaking your prompt cache. + // Optional. Default is 0. + ClearAtLeastTokens int64 + + // ClearExcludeTools is list of tool names whose tool uses and results should never be cleared. + // Optional. Default is nil. + ClearExcludeTools []string + + // ClearMessageRewriter is a pre-process handler before clearing specific tool call and tool response pairs. + // You can rewrite tool call and tool messages extracted as parameters and return a rearranged message slice. + // This can be useful when you want to remove some tool calls (e.g., write_file / edit_file) and rewrite them + // as a user message (e.g. ). + // Returned messages will replace the original tool call and tool messages and will count towards ClearAtLeastTokens. + // If returned messagesAfterRewrite is nil, tool call and tool messages will be removed. + // Optional. Default is nil, which means no rewrite. + ClearMessageRewriter func(ctx context.Context, toolCallMsg M, toolResponseMsgs []M) (messagesAfterRewrite []M, err error) + + // ClearPostProcess is clear post process handler. + // Optional. + ClearPostProcess func(ctx context.Context, state *adk.TypedChatModelAgentState[M]) context.Context + + // ToolConfig is the specific configuration that applies to tools by name. + // This configuration takes precedence over GeneralConfig for the specified tools. + // Optional. + ToolConfig map[string]*ToolReductionConfig +} + +// Config is the backward-compatible alias for TypedConfig with *schema.Message. +type Config = TypedConfig[*schema.Message] + +type ToolReductionConfig struct { + // Backend is the storage backend where offloaded content will be saved. + // Required when truncation is enabled for this tool (SkipTruncation is false). + // Optional for clear-only usage. If Backend is nil, clear will still replace tool outputs with placeholders + // but will not offload content. + Backend Backend + + // SkipTruncation skip truncation for this tool. + SkipTruncation bool + + // TruncHandler is used to process tool call results during truncation. + // Optional. Default using defaultTruncHandler when SkipTruncation is false but TruncHandler is nil. + TruncHandler func(ctx context.Context, detail *ToolDetail) (*TruncResult, error) + + // SkipClear skip clear for this tool. + SkipClear bool + + // ClearHandler is used to process tool call arguments and results during clearing. + // Optional. Default using defaultClearHandler when SkipClear is false but ClearHandler is nil. + ClearHandler func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) +} + +type ToolDetail struct { + // ToolContext provides metadata about the tool call (e.g., tool name, call ID). + ToolContext *adk.ToolContext + + // ToolArgument contains the arguments passed to the tool. + ToolArgument *schema.ToolArgument + + // ToolResult contains the output returned by the invokable tool. + ToolResult *schema.ToolResult + + // StreamToolResult contains the output returned by the streamable tool. + StreamToolResult *schema.StreamReader[*schema.ToolResult] +} + +type TruncResult struct { + // NeedTrunc indicates whether the tool result should be truncated. + NeedTrunc bool + + // ToolResult contains the result returned by the invokable tool after trunc. + // Required when NeedTrunc is true and ToolDetail.ToolResult is not nil. + ToolResult *schema.ToolResult + + // StreamToolResult contains the output returned by the streamable tool after trunc. + // Required when NeedTrunc is true and ToolDetail.StreamToolResult is not nil. + StreamToolResult *schema.StreamReader[*schema.ToolResult] + + // NeedOffload indicates whether the tool result should be offloaded. + NeedOffload bool + + // OffloadFilePath is the path where the offloaded content should be stored. + // This path is typically relative to the backend's root. + // Required when NeedOffload is true. + OffloadFilePath string + + // OffloadContent is the actual content to be written to the storage backend. + // Required when NeedOffload is true. + OffloadContent string +} + +// ClearResult contains the result of the Handler's decision. +type ClearResult struct { + // NeedClear indicates whether the tool argument and result should be cleared. + NeedClear bool + + // ToolArgument contains the arguments passed to the tool after clear. + // Required when NeedClear is true. + ToolArgument *schema.ToolArgument + + // ToolResult contains the output returned by the tool after clear. + // Required when NeedClear is true + ToolResult *schema.ToolResult + + // NeedOffload indicates whether the tool argument and result should be offloaded. + NeedOffload bool + + // OffloadFilePath is the path where the offloaded content should be stored. + // This path is typically relative to the backend's root. + // Required when NeedOffload is true. + OffloadFilePath string + + // OffloadContent is the actual content to be written to the storage backend. + // Required when NeedOffload is true. + OffloadContent string +} + +func (t *TypedConfig[M]) copyAndFillDefaults() (*TypedConfig[M], error) { + cfg := &TypedConfig[M]{ + Backend: t.Backend, + SkipTruncation: t.SkipTruncation, + SkipClear: t.SkipClear, + ReadFileToolName: t.ReadFileToolName, + RootDir: t.RootDir, + GenTruncOffloadFilePath: t.GenTruncOffloadFilePath, + GenClearOffloadFilePath: t.GenClearOffloadFilePath, + MaxLengthForTrunc: t.MaxLengthForTrunc, + TruncExcludeTools: t.TruncExcludeTools, + TokenCounter: t.TokenCounter, + MaxTokensForClear: t.MaxTokensForClear, + ClearRetentionSuffixLimit: t.ClearRetentionSuffixLimit, + ClearAtLeastTokens: t.ClearAtLeastTokens, + ClearExcludeTools: t.ClearExcludeTools, + ClearMessageRewriter: t.ClearMessageRewriter, + ClearPostProcess: t.ClearPostProcess, + } + if cfg.TokenCounter == nil { + cfg.TokenCounter = getDefaultTokenCounter[M]() + } + if cfg.ClearRetentionSuffixLimit == 0 { + cfg.ClearRetentionSuffixLimit = 1 + } + if cfg.ReadFileToolName == "" { + cfg.ReadFileToolName = "read_file" + } + if cfg.RootDir == "" { + cfg.RootDir = "/tmp" + } + if cfg.GenTruncOffloadFilePath == nil { + cfg.GenTruncOffloadFilePath = func(ctx context.Context, toolDetail *ToolDetail) (filePath string, err error) { + tcID := toolDetail.ToolContext.CallID + if tcID == "" { + tcID = uuid.NewString() + } + return filepath.Join(cfg.RootDir, "trunc", tcID), nil + } + } + if cfg.GenClearOffloadFilePath == nil { + cfg.GenClearOffloadFilePath = func(ctx context.Context, toolDetail *ToolDetail) (filePath string, err error) { + tcID := toolDetail.ToolContext.CallID + if tcID == "" { + tcID = uuid.NewString() + } + return filepath.Join(cfg.RootDir, "clear", tcID), nil + } + } + if cfg.MaxLengthForTrunc == 0 { + cfg.MaxLengthForTrunc = 50000 + } + if cfg.MaxTokensForClear == 0 { + cfg.MaxTokensForClear = 160000 + } + if t.ToolConfig != nil { + cfg.ToolConfig = make(map[string]*ToolReductionConfig, len(t.ToolConfig)) + for toolName, trc := range t.ToolConfig { + cpConfig := &ToolReductionConfig{ + Backend: trc.Backend, + SkipTruncation: trc.SkipTruncation, + SkipClear: trc.SkipClear, + TruncHandler: trc.TruncHandler, + ClearHandler: trc.ClearHandler, + } + cfg.ToolConfig[toolName] = cpConfig + } + } + + return cfg, nil +} + +// NewTyped creates a generic tool reduction middleware from config. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +// Both message types support the full truncation and clear phases. +func NewTyped[M adk.MessageType](_ context.Context, config *TypedConfig[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + var err error + if config == nil { + return nil, fmt.Errorf("config must not be nil") + } + if config.Backend == nil && !config.SkipTruncation { + return nil, fmt.Errorf("backend must be set when not skipping truncation") + } + + config, err = config.copyAndFillDefaults() + if err != nil { + return nil, err + } + defaultReductionConfig := &ToolReductionConfig{ + Backend: config.Backend, + SkipTruncation: config.SkipTruncation, + SkipClear: config.SkipClear, + } + if !defaultReductionConfig.SkipTruncation { + defaultReductionConfig.TruncHandler = defaultTruncHandler(config.GenTruncOffloadFilePath, config.MaxLengthForTrunc, config.ReadFileToolName) + } + if !defaultReductionConfig.SkipClear { + defaultReductionConfig.ClearHandler = defaultClearHandler(config.GenClearOffloadFilePath, config.Backend != nil, config.ReadFileToolName) + } + excludeTruncTools := make(map[string]struct{}, len(config.TruncExcludeTools)) + for _, toolName := range config.TruncExcludeTools { + excludeTruncTools[toolName] = struct{}{} + } + excludeClearTools := make(map[string]struct{}, len(config.ClearExcludeTools)) + for _, toolName := range config.ClearExcludeTools { + excludeClearTools[toolName] = struct{}{} + } + + return &typedToolReductionMiddleware[M]{ + config: config, + defaultConfig: defaultReductionConfig, + excludeTruncTools: excludeTruncTools, + excludeClearTools: excludeClearTools, + }, nil +} + +// New creates tool reduction middleware from config +func New(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped(ctx, config) +} + +type typedToolReductionMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + + config *TypedConfig[M] + defaultConfig *ToolReductionConfig + + excludeTruncTools map[string]struct{} + excludeClearTools map[string]struct{} +} + +// getDefaultTokenCounter returns a default token counter function that operates on []M. +// For *schema.Message it delegates to defaultTokenCounter. +// For *schema.AgenticMessage it uses a simple character-based estimation. +func getDefaultTokenCounter[M adk.MessageType]() func(ctx context.Context, msgs []M, tools []*schema.ToolInfo) (int64, error) { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(func(ctx context.Context, msgs []*schema.Message, tools []*schema.ToolInfo) (int64, error) { + return defaultTokenCounter(ctx, msgs, tools) + }).(func(context.Context, []M, []*schema.ToolInfo) (int64, error)) + case *schema.AgenticMessage: + return any(func(ctx context.Context, msgs []*schema.AgenticMessage, tools []*schema.ToolInfo) (int64, error) { + return defaultAgenticTokenCounter(ctx, msgs, tools) + }).(func(context.Context, []M, []*schema.ToolInfo) (int64, error)) + } + panic("unreachable") +} + +func defaultAgenticTokenCounter(_ context.Context, msgs []*schema.AgenticMessage, tools []*schema.ToolInfo) (int64, error) { + var tokens int64 + for _, msg := range msgs { + if msg == nil { + continue + } + tokens += int64(len(msg.Role)) / 4 + for _, block := range msg.ContentBlocks { + if block != nil { + tokens += int64(len(block.String())) / 4 + } + } + } + for _, tl := range tools { + tl_ := *tl + tl_.Extra = nil + text, err := sonic.MarshalString(tl_) + if err != nil { + return 0, fmt.Errorf("failed to marshal tool info: %w", err) + } + tokens += int64(len(text) / 4) + } + return tokens, nil +} + +func (t *typedToolReductionMiddleware[M]) getToolConfig(toolName string, sc scene) *ToolReductionConfig { + if t.config.ToolConfig != nil { + if cfg, ok := t.config.ToolConfig[toolName]; ok { + if (sc == sceneTruncation && !cfg.SkipTruncation && cfg.TruncHandler == nil) || + (sc == sceneClear && !cfg.SkipClear && cfg.ClearHandler == nil) { + return t.defaultConfig + } + return cfg + } + } + return t.defaultConfig +} + +func (t *typedToolReductionMiddleware[M]) WrapInvokableToolCall(_ context.Context, endpoint adk.InvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.InvokableToolCallEndpoint, error) { + cfg := t.getToolConfig(tCtx.Name, sceneTruncation) + if cfg == nil || cfg.TruncHandler == nil { + return endpoint, nil + } + if _, excluded := t.excludeTruncTools[tCtx.Name]; excluded { + return endpoint, nil + } + + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + output, err := endpoint(ctx, argumentsInJSON, opts...) + if err != nil { + return "", err + } + detail := &ToolDetail{ + ToolContext: tCtx, + ToolArgument: &schema.ToolArgument{ + Text: argumentsInJSON, + }, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: output}, + }, + }, + } + truncResult, err := cfg.TruncHandler(ctx, detail) + if err != nil { + return "", err + } + if !truncResult.NeedTrunc { + return output, nil + } + if truncResult.NeedOffload { + if cfg.Backend == nil { + return "", fmt.Errorf("truncation: no backend for offload") + } + if err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: truncResult.OffloadFilePath, + Content: truncResult.OffloadContent, + }); err != nil { + return "", err + } + } + return truncResult.ToolResult.Parts[0].Text, nil + }, nil +} + +func (t *typedToolReductionMiddleware[M]) WrapStreamableToolCall(_ context.Context, endpoint adk.StreamableToolCallEndpoint, tCtx *adk.ToolContext) (adk.StreamableToolCallEndpoint, error) { + cfg := t.getToolConfig(tCtx.Name, sceneTruncation) + if cfg == nil || cfg.TruncHandler == nil { + return endpoint, nil + } + if _, excluded := t.excludeTruncTools[tCtx.Name]; excluded { + return endpoint, nil + } + + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + output, err := endpoint(ctx, argumentsInJSON, opts...) + if err != nil { + return nil, err + } + if cfg == t.defaultConfig { + return t.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, argumentsInJSON, output), nil + } + + readers := output.Copy(2) + output = readers[0] + origResp := readers[1] + + detail := &ToolDetail{ + ToolContext: tCtx, + ToolArgument: &schema.ToolArgument{ + Text: argumentsInJSON, + }, + StreamToolResult: schema.StreamReaderWithConvert(output, func(t string) (*schema.ToolResult, error) { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: t}}}, nil + }), + } + truncResult, err := cfg.TruncHandler(ctx, detail) + if err != nil { + origResp.Close() + return nil, err + } + if !truncResult.NeedTrunc { + return origResp, nil + } + origResp.Close() // close err resp when not using it + + if truncResult.NeedOffload { + if cfg.Backend == nil { + return nil, fmt.Errorf("truncation: no backend for offload") + } + if err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: truncResult.OffloadFilePath, + Content: truncResult.OffloadContent, + }); err != nil { + return nil, err + } + } + + sr := schema.StreamReaderWithConvert(truncResult.StreamToolResult, func(t *schema.ToolResult) (string, error) { + if t == nil || len(t.Parts) == 0 { + return "", nil + } + return t.Parts[0].Text, nil + }) + return sr, nil + }, nil +} + +func (t *typedToolReductionMiddleware[M]) wrapDefaultStreamableTruncation( + ctx context.Context, + cfg *ToolReductionConfig, + tCtx *adk.ToolContext, + argumentsInJSON string, + output *schema.StreamReader[string], +) *schema.StreamReader[string] { + var ( + fullOutput strings.Builder + curSize int + reachedLimit bool + truncTextPrefix string + truncTriggeredFullText string + sourceErr error + ) + + // Chunk mapping: + // + // Normal EOF: + // input : C1 -------- C2(crosses limit) -------- C3 -------- EOF + // output: C1 -------- [swallow C2,C3] ---------- prefix(C2)+notice + // + // Source error after truncation: + // input : C1 -------- C2(crosses limit) -------- C3 -------- ERR + // output: C1 -------- [swallow C2,C3] ---------- replay(C2+C3 tail) -- ERR + // + // Before truncation, chunks are forwarded unchanged. Once the limit is reached, + // the wrapper keeps reading and accumulating full output but suppresses chunks. + // On clean EOF, it offloads full output and emits one final truncation notice. + // On source error, it replays the original content swallowed since the truncation + // point, then surfaces the original source error via the outer convert layer. + convertHandler := func(chunk string) (string, error) { + if chunk == "" { + return "", schema.ErrNoValue + } + fullOutput.WriteString(chunk) + if reachedLimit { + return "", schema.ErrNoValue + } + + chunkSize := len(chunk) + if curSize+chunkSize < t.config.MaxLengthForTrunc { + curSize += chunkSize + return chunk, nil + } + + reachedLimit = true + truncTriggeredFullText = fullOutput.String() + + leftSize := t.config.MaxLengthForTrunc - curSize + if leftSize > 0 { + truncTextPrefix = clampPrefixToUTF8Boundary(chunk, leftSize) + } + return "", schema.ErrNoValue + } + + errHandler := func(err error) error { + if reachedLimit { + sourceErr = err + return nil + } + return err + } + + eofHandler := func() (any, error) { + if sourceErr != nil { + if fullOutput.Len() > curSize { + return fullOutput.String()[curSize:], nil + } + return nil, io.EOF + } + if !reachedLimit { + return nil, io.EOF + } + detail := &ToolDetail{ + ToolContext: tCtx, + ToolArgument: &schema.ToolArgument{ + Text: argumentsInJSON, + }, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: truncTriggeredFullText}, + }, + }, + } + + var ( + offloadNotify string + errorNotify string + ) + filePath, err := t.config.GenTruncOffloadFilePath(ctx, detail) + if err != nil { + errorNotify = formatStreamOffloadFailedNotify(err) + } else if cfg.Backend == nil { + errorNotify = formatStreamOffloadFailedNotify(fmt.Errorf("truncation: no backend for offload")) + } else if err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: filePath, + Content: fullOutput.String(), + }); err != nil { + errorNotify = formatStreamOffloadFailedNotify(err) + } else { + offloadNotify = formatStreamOffloadSavedNotify(filePath, t.config.ReadFileToolName) + } + + notice, err := formatStreamTruncNotice(t.config.MaxLengthForTrunc, offloadNotify, errorNotify) + if err != nil { + return nil, err + } + return truncTextPrefix + notice, nil + } + + truncatedStream := schema.StreamReaderWithConvert( + output, + convertHandler, + schema.WithErrWrapper(errHandler), + schema.WithOnEOF(eofHandler), + ) + + // StreamReaderWithConvert's WithOnEOF can inject only one value or one error. + // The source-error fallback needs two Recv results: + // 1. replay swallowed original output, + // 2. surface the original source error. + // The inner reader emits the replay value, and this outer reader emits the + // original error when the inner reader reaches EOF after replay. + return schema.StreamReaderWithConvert( + truncatedStream, + func(chunk string) (string, error) { + return chunk, nil + }, + schema.WithOnEOF(func() (any, error) { + if sourceErr != nil { + return nil, sourceErr + } + return nil, io.EOF + }), + ) +} + +func (t *typedToolReductionMiddleware[M]) WrapEnhancedInvokableToolCall(_ context.Context, endpoint adk.EnhancedInvokableToolCallEndpoint, tCtx *adk.ToolContext) (adk.EnhancedInvokableToolCallEndpoint, error) { + cfg := t.getToolConfig(tCtx.Name, sceneTruncation) + if cfg == nil || cfg.TruncHandler == nil { + return endpoint, nil + } + if _, excluded := t.excludeTruncTools[tCtx.Name]; excluded { + return endpoint, nil + } + + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + output, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + detail := &ToolDetail{ + ToolContext: tCtx, + ToolArgument: toolArgument, + ToolResult: output, + } + truncResult, err := cfg.TruncHandler(ctx, detail) + if err != nil { + return nil, err + } + if !truncResult.NeedTrunc { + return output, nil + } + if truncResult.NeedOffload { + if cfg.Backend == nil { + return nil, fmt.Errorf("truncation: no backend for offload") + } + if err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: truncResult.OffloadFilePath, + Content: truncResult.OffloadContent, + }); err != nil { + return nil, err + } + } + return truncResult.ToolResult, nil + }, nil +} + +func (t *typedToolReductionMiddleware[M]) WrapEnhancedStreamableToolCall(_ context.Context, endpoint adk.EnhancedStreamableToolCallEndpoint, tCtx *adk.ToolContext) (adk.EnhancedStreamableToolCallEndpoint, error) { + cfg := t.getToolConfig(tCtx.Name, sceneTruncation) + if cfg == nil || cfg.TruncHandler == nil { + return endpoint, nil + } + if _, excluded := t.excludeTruncTools[tCtx.Name]; excluded { + return endpoint, nil + } + + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + output, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + if cfg == t.defaultConfig { + return t.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArgument, output), nil + } + + readers := output.Copy(2) + output = readers[0] + origResp := readers[1] + + detail := &ToolDetail{ + ToolContext: tCtx, + ToolArgument: toolArgument, + StreamToolResult: output, + } + truncResult, err := cfg.TruncHandler(ctx, detail) + if err != nil { + origResp.Close() + return nil, err + } + if !truncResult.NeedTrunc { + return origResp, nil + } + origResp.Close() // close err resp when not using it + + if truncResult.NeedOffload { + if cfg.Backend == nil { + return nil, fmt.Errorf("truncation: no backend for offload") + } + if err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: truncResult.OffloadFilePath, + Content: truncResult.OffloadContent, + }); err != nil { + return nil, err + } + } + + return truncResult.StreamToolResult, nil + }, nil +} + +func (t *typedToolReductionMiddleware[M]) wrapDefaultEnhancedStreamableTruncation( + ctx context.Context, + cfg *ToolReductionConfig, + tCtx *adk.ToolContext, + toolArgument *schema.ToolArgument, + output *schema.StreamReader[*schema.ToolResult], +) *schema.StreamReader[*schema.ToolResult] { + var ( + chunks []*schema.ToolResult + replayChunks []*schema.ToolResult + fullLength int + reachedLimit bool + truncTextPrefix *schema.ToolResult + truncTriggeredToolResult *schema.ToolResult + sourceErr error + ) + + // Chunk mapping: + // + // Normal EOF: + // input : C1 -------- C2(crosses limit) -------- C3 -------- EOF + // output: C1 -------- [swallow C2,C3] ---------- prefix(C2)+notice + // + // Source error after truncation: + // input : C1 -------- C2(crosses limit) -------- C3 -------- ERR + // output: C1 -------- [swallow C2,C3] ---------- replay(C2+C3) -- ERR + // + // Replay uses one merged ToolResult because WithOnEOF can inject only one + // value. As in the string stream wrapper, the outer convert layer surfaces + // the original source error on the following Recv. + convertHandler := func(chunk *schema.ToolResult) (*schema.ToolResult, error) { + chunks = append(chunks, chunk) + chunkLength := toolResultTextLength(chunk) + if reachedLimit { + fullLength += chunkLength + replayChunks = append(replayChunks, chunk) + return nil, schema.ErrNoValue + } + if fullLength+chunkLength < t.config.MaxLengthForTrunc { + fullLength += chunkLength + return chunk, nil + } + + reachedLimit = true + replayChunks = append(replayChunks, chunk) + truncTextPrefix = toolResultTextPrefix(chunk, t.config.MaxLengthForTrunc-fullLength) + fullLength += chunkLength + truncTriggeredToolResult = &schema.ToolResult{Parts: flattenToolResultParts(chunks)} + return nil, schema.ErrNoValue + } + + errHandler := func(err error) error { + if reachedLimit { + sourceErr = err + return nil + } + return err + } + + eofHandler := func() (any, error) { + if sourceErr != nil { + replay := &schema.ToolResult{Parts: flattenToolResultParts(replayChunks)} + if len(replay.Parts) > 0 { + return replay, nil + } + return nil, io.EOF + } + if !reachedLimit { + return nil, io.EOF + } + + fullToolResult := &schema.ToolResult{Parts: flattenToolResultParts(chunks)} + detail := &ToolDetail{ + ToolContext: tCtx, + ToolArgument: toolArgument, + ToolResult: truncTriggeredToolResult, + } + + var ( + offloadNotify string + errorNotify string + ) + filePath, err := t.config.GenTruncOffloadFilePath(ctx, detail) + if err != nil { + errorNotify = formatStreamOffloadFailedNotify(err) + } else if cfg.Backend == nil { + errorNotify = formatStreamOffloadFailedNotify(fmt.Errorf("truncation: no backend for offload")) + } else if err = cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: filePath, + Content: stringifyToolOutputParts(fullToolResult.Parts), + }); err != nil { + errorNotify = formatStreamOffloadFailedNotify(err) + } else { + offloadNotify = formatStreamOffloadSavedNotify(filePath, t.config.ReadFileToolName) + } + + notice, err := formatStreamTruncNotice(t.config.MaxLengthForTrunc, offloadNotify, errorNotify) + if err != nil { + return nil, err + } + + parts := make([]schema.ToolOutputPart, 0, 1) + if truncTextPrefix != nil { + parts = append(parts, truncTextPrefix.Parts...) + } + parts = append(parts, schema.ToolOutputPart{Type: schema.ToolPartTypeText, Text: notice}) + return &schema.ToolResult{Parts: parts}, nil + } + + truncatedStream := schema.StreamReaderWithConvert( + output, + convertHandler, + schema.WithErrWrapper(errHandler), + schema.WithOnEOF(eofHandler), + ) + + return schema.StreamReaderWithConvert( + truncatedStream, + func(chunk *schema.ToolResult) (*schema.ToolResult, error) { + return chunk, nil + }, + schema.WithOnEOF(func() (any, error) { + if sourceErr != nil { + return nil, sourceErr + } + return nil, io.EOF + }), + ) +} + +func formatStreamTruncNotice(previewSize int, offloadNotify, errorMsgNotify string) (string, error) { + return pyfmt.Fmt(getStreamTruncFmt(), map[string]any{ + "preview_size": previewSize, + "offload_notify": offloadNotify, + "error_msg_notify": errorMsgNotify, + }) +} + +func toolResultTextPrefix(result *schema.ToolResult, maxLength int) *schema.ToolResult { + if result == nil || maxLength <= 0 { + return nil + } + var ( + length int + parts []schema.ToolOutputPart + ) + for _, part := range result.Parts { + if part.Type != schema.ToolPartTypeText { + parts = append(parts, part) + continue + } + remaining := maxLength - length + if remaining <= 0 { + break + } + if len(part.Text) <= remaining { + parts = append(parts, part) + length += len(part.Text) + continue + } + part.Text = clampPrefixToUTF8Boundary(part.Text, remaining) + if part.Text != "" { + parts = append(parts, part) + } + break + } + if len(parts) == 0 { + return nil + } + return &schema.ToolResult{Parts: parts} +} + +func toolResultTextLength(result *schema.ToolResult) int { + if result == nil { + return 0 + } + var length int + for _, part := range result.Parts { + if part.Type == schema.ToolPartTypeText { + length += len(part.Text) + } + } + return length +} + +func flattenToolResultParts(chunks []*schema.ToolResult) []schema.ToolOutputPart { + var parts []schema.ToolOutputPart + for _, chunk := range chunks { + if chunk == nil || len(chunk.Parts) == 0 { + continue + } + parts = append(parts, chunk.Parts...) + } + return parts +} + +func writeTruncOffload(ctx context.Context, cfg *ToolReductionConfig, truncResult *TruncResult) error { + if !truncResult.NeedOffload { + return nil + } + if cfg.Backend == nil { + return fmt.Errorf("truncation: no backend for offload") + } + return cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: truncResult.OffloadFilePath, + Content: truncResult.OffloadContent, + }) +} + +func (t *typedToolReductionMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], mc *adk.TypedModelContext[M]) ( + context.Context, *adk.TypedChatModelAgentState[M], error) { + + return t.beforeModelRewriteStateGeneric(ctx, state, mc) +} + +func (t *typedToolReductionMiddleware[M]) beforeModelRewriteStateGeneric(ctx context.Context, state *adk.TypedChatModelAgentState[M], _ *adk.TypedModelContext[M]) ( + context.Context, *adk.TypedChatModelAgentState[M], error) { + + var ( + err error + estimatedTokens int64 + ) + + // init msg tokens + estimatedTokens, err = t.config.TokenCounter(ctx, state.Messages, state.ToolInfos) + if err != nil { + return ctx, state, err + } + + if estimatedTokens < t.config.MaxTokensForClear { + return ctx, state, nil + } + + // calc range + var ( + start = 0 + end = len(state.Messages) + ) + for ; start < len(state.Messages); start++ { + msg := state.Messages[start] + if isAssistantMsg(msg) && !getMsgClearedFlagGeneric(msg) { + break + } + } + retention := t.config.ClearRetentionSuffixLimit + for ; retention > 0 && end > 0; end-- { + msg := state.Messages[end-1] + if isAssistantMsg(msg) && hasToolCalls(msg) { + retention-- + if retention == 0 { + end-- + break + } + } + } + if start >= end { + return ctx, state, nil + } + var ( + editTarget []M + clearAtLeastTokens = t.config.ClearAtLeastTokens + offloadStash []*offloadStashItem + ) + + editTarget, end, err = t.applyClearRewriteGeneric(ctx, state, start, end, clearAtLeastTokens) + if err != nil { + return ctx, state, err + } + + // recursively handle + toolCallMsgIndex := start + + for toolCallMsgIndex < end { + toolCallMsg := editTarget[toolCallMsgIndex] + toolCalls := getToolCallsGeneric(toolCallMsg) + if isAssistantMsg(toolCallMsg) && len(toolCalls) > 0 { + for _, tc := range toolCalls { + // Find the corresponding tool-result message by callID + resultMsgIndex, found := findToolResultByCallID(editTarget, toolCallMsgIndex+1, toolCallMsgIndex+1+len(toolCalls), tc.CallID) + if !found { + continue // No corresponding tool result found, skip + } + resultMsg := editTarget[resultMsgIndex] + + if _, found := t.excludeClearTools[tc.Name]; found { + continue + } + cfg := t.getToolConfig(tc.Name, sceneClear) + if cfg == nil || cfg.ClearHandler == nil { + continue + } + + toolResult, fromContent, toolResultErr := toolResultFromMsgGeneric(resultMsg) + if toolResultErr != nil { + return ctx, state, toolResultErr + } + + td := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: tc.Name, + CallID: tc.CallID, + }, + ToolArgument: &schema.ToolArgument{ + Text: tc.Arguments, + }, + ToolResult: toolResult, + } + + offloadInfo, offloadErr := cfg.ClearHandler(ctx, td) + if offloadErr != nil { + return ctx, state, offloadErr + } + if !offloadInfo.NeedClear { + continue + } + if offloadInfo.NeedOffload { + if cfg.Backend == nil { + return ctx, state, fmt.Errorf("clear: no backend for offload") + } + if clearAtLeastTokens > 0 { // delay clear offloading + offloadStash = append(offloadStash, &offloadStashItem{ + config: cfg, + offloadInfo: offloadInfo, + }) + } else { // instant clear offloading + writeErr := cfg.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: offloadInfo.OffloadFilePath, + Content: offloadInfo.OffloadContent, + }) + if writeErr != nil { + return ctx, state, writeErr + } + } + } + + setToolCallArguments(toolCallMsg, tc.BlockIndex, offloadInfo.ToolArgument.Text) + setToolResultContent(resultMsg, offloadInfo.ToolResult, fromContent) + } + + // set dedup flag + setMsgClearedFlagGeneric(toolCallMsg) + } + toolCallMsgIndex++ + } + + if clearAtLeastTokens > 0 { + estimatedTokensAfterClear, err := t.config.TokenCounter(ctx, editTarget, state.ToolInfos) + if err != nil { + return ctx, state, err + } + tokensCleared := estimatedTokens - estimatedTokensAfterClear + if tokensCleared < clearAtLeastTokens { + // clear not applied, post process won't apply as well. + return ctx, state, nil + } + for _, item := range offloadStash { + writeErr := item.config.Backend.Write(ctx, &filesystem.WriteRequest{ + FilePath: item.offloadInfo.OffloadFilePath, + Content: item.offloadInfo.OffloadContent, + }) + if writeErr != nil { + return ctx, state, writeErr + } + } + } + + state.Messages = editTarget // replace original state messages + + if t.config.ClearPostProcess != nil { + ctx = t.config.ClearPostProcess(ctx, state) + } + + return ctx, state, nil +} + +func (t *typedToolReductionMiddleware[M]) applyClearRewriteGeneric(ctx context.Context, state *adk.TypedChatModelAgentState[M], start, end int, clearAtLeastTokens int64) ( + []M, int, error) { + var ( + editTarget []M + needProcessPart []M + ) + + editTarget = append(editTarget, state.Messages[:start]...) + + if clearAtLeastTokens > 0 { + needProcessPart = copyMessagesGeneric(state.Messages[start:end]) + } else { + needProcessPart = state.Messages[start:end] + } + + if t.config.ClearMessageRewriter != nil { + var ( + rewritten []M + origLength = len(needProcessPart) + ) + for i := 0; i < len(needProcessPart); { + msg := needProcessPart[i] + if isSystemMsg(msg) || isUserMsg(msg) { + rewritten = append(rewritten, msg) + i++ + } else if isToolResultMsg(msg) { + // tool result message (schema.Tool role or agentic user msg carrying FunctionToolResult) + i++ + } else if isAssistantMsg(msg) { + toolCalls := getToolCallsGeneric(msg) + if len(toolCalls) == 0 { + rewritten = append(rewritten, msg) + i++ + continue + } + var ( + toolResponseMessages []M + trStart = i + 1 + trEnd = i + len(toolCalls) + 1 + ) + if trStart >= trEnd || trStart >= len(needProcessPart) || trEnd > len(needProcessPart) { + toolResponseMessages = nil + } else { + toolResponseMessages = needProcessPart[trStart:trEnd] + } + + rewrittenMessages, rewriteErr := t.config.ClearMessageRewriter(ctx, msg, toolResponseMessages) + if rewriteErr != nil { + return nil, 0, rewriteErr + } + rewritten = append(rewritten, rewrittenMessages...) + i = trEnd + } else { // unexpected + return nil, 0, fmt.Errorf("[applyClearRewrite] unexpected message: %v", any(msg)) + } + } + editTarget = append(editTarget, rewritten...) + editTarget = append(editTarget, state.Messages[end:]...) + end = end - origLength + len(rewritten) + } else { + editTarget = append(editTarget, needProcessPart...) + editTarget = append(editTarget, state.Messages[end:]...) + } + + return editTarget, end, nil +} + +func agenticResultCallID(block *schema.ContentBlock) (string, bool) { + if block == nil { + return "", false + } + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + return block.FunctionToolResult.CallID, true + } + if block.Type == schema.ContentBlockTypeToolSearchResult && block.ToolSearchFunctionToolResult != nil { + return block.ToolSearchFunctionToolResult.CallID, true + } + return "", false +} + +func findToolResultByCallID[M adk.MessageType](messages []M, startIndex, endIndex int, callID string) (int, bool) { + for i := startIndex; i < endIndex; i++ { + msg := messages[i] + if isToolResultMsg(msg) { + if getToolResultCallID(msg) == callID { + return i, true + } + } else { + return -1, false + } + } + return -1, false +} + +func getToolResultCallID[M adk.MessageType](msg M) string { + switch m := any(msg).(type) { + case *schema.Message: + if m.Role == schema.Tool { + return m.ToolCallID + } + case *schema.AgenticMessage: + if m.Role == schema.AgenticRoleTypeUser { + for _, block := range m.ContentBlocks { + if callID, ok := agenticResultCallID(block); ok { + return callID + } + } + } + } + return "" +} + +type offloadStashItem struct { + config *ToolReductionConfig + offloadInfo *ClearResult +} + +// toolCallInfo represents a tool call extracted from a message for generic processing. +type toolCallInfo struct { + // BlockIndex is the index used to locate the tool call within the message. + // For *schema.Message: index into msg.ToolCalls slice. + // For *schema.AgenticMessage: index into msg.ContentBlocks slice. + BlockIndex int + CallID string + Name string + Arguments string +} + +// isAssistantMsg checks if a message has assistant role. +func isAssistantMsg[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.Assistant + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeAssistant + } + return false +} + +// isSystemMsg checks if a message has system role. +func isSystemMsg[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.System + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeSystem + } + return false +} + +// isUserMsg checks if a message has user role (and is not a tool-result message). +func isUserMsg[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.User + case *schema.AgenticMessage: + if m.Role != schema.AgenticRoleTypeUser { + return false + } + // A user-role agentic message that contains any FunctionToolResult block + // is a tool result message, not a normal user message — even if it also + // carries UserInput blocks. This ensures the clear flow's tool-call grouping + // remains correctly aligned. + for _, block := range m.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult { + return false + } + } + return len(m.ContentBlocks) > 0 + } + return false +} + +// hasToolCalls checks if an assistant message contains tool calls. +func hasToolCalls[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return len(m.ToolCalls) > 0 + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall { + return true + } + } + } + return false +} + +// isToolResultMsg checks if a message is a tool result message. +// For *schema.Message: role == Tool. +// For *schema.AgenticMessage: user-role message with at least one FunctionToolResult block. +func isToolResultMsg[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.Tool + case *schema.AgenticMessage: + if m.Role != schema.AgenticRoleTypeUser { + return false + } + for _, block := range m.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolResult { + return true + } + } + } + return false +} + +// isToolResultOnlyMsg checks if a message is exclusively a tool result message +// (no other content besides tool results). +// For *schema.Message: role == Tool. +// For *schema.AgenticMessage: user-role message where ALL content blocks are FunctionToolResult. +func isToolResultOnlyMsg[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.Tool + case *schema.AgenticMessage: + if m.Role != schema.AgenticRoleTypeUser || len(m.ContentBlocks) == 0 { + return false + } + for _, block := range m.ContentBlocks { + if block == nil || block.Type != schema.ContentBlockTypeFunctionToolResult { + return false + } + } + return true + } + return false +} + +// getMsgClearedFlagGeneric checks if a message has the cleared flag set. +func getMsgClearedFlagGeneric[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return getMsgClearedFlag(m) + case *schema.AgenticMessage: + if m.Extra == nil { + return false + } + v, ok := m.Extra[msgClearedFlag].(bool) + return ok && v + } + return false +} + +// setMsgClearedFlagGeneric sets the cleared flag on a message. +func setMsgClearedFlagGeneric[M adk.MessageType](msg M) { + switch m := any(msg).(type) { + case *schema.Message: + setMsgClearedFlag(m) + case *schema.AgenticMessage: + if m.Extra == nil { + m.Extra = make(map[string]any) + } + m.Extra[msgClearedFlag] = true + } +} + +// getToolCallsGeneric extracts tool call info from an assistant message. +func getToolCallsGeneric[M adk.MessageType](msg M) []toolCallInfo { + switch m := any(msg).(type) { + case *schema.Message: + if len(m.ToolCalls) == 0 { + return nil + } + result := make([]toolCallInfo, 0, len(m.ToolCalls)) + for i, tc := range m.ToolCalls { + result = append(result, toolCallInfo{ + BlockIndex: i, + CallID: tc.ID, + Name: tc.Function.Name, + Arguments: tc.Function.Arguments, + }) + } + return result + case *schema.AgenticMessage: + var result []toolCallInfo + for i, block := range m.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil { + result = append(result, toolCallInfo{ + BlockIndex: i, + CallID: block.FunctionToolCall.CallID, + Name: block.FunctionToolCall.Name, + Arguments: block.FunctionToolCall.Arguments, + }) + } + } + return result + } + return nil +} + +// setToolCallArguments updates the arguments for a tool call at the given block index. +func setToolCallArguments[M adk.MessageType](msg M, blockIndex int, args string) { + switch m := any(msg).(type) { + case *schema.Message: + m.ToolCalls[blockIndex].Function.Arguments = args + case *schema.AgenticMessage: + if m.ContentBlocks[blockIndex].FunctionToolCall != nil { + m.ContentBlocks[blockIndex].FunctionToolCall.Arguments = args + } + } +} + +// toolResultFromMsgGeneric extracts tool result from a message as a *schema.ToolResult. +// For *schema.Message: delegates to existing toolResultFromMessage. +// For *schema.AgenticMessage: iterates FunctionToolResult blocks. +// The fromContent flag indicates whether the result came from simple content (true) +// or multi-part content (false), which affects how setToolResultContent writes it back. +func toolResultFromMsgGeneric[M adk.MessageType](msg M) (result *schema.ToolResult, fromContent bool, err error) { + switch m := any(msg).(type) { + case *schema.Message: + return toolResultFromMessage(m) + case *schema.AgenticMessage: + var found *schema.FunctionToolResult + for _, block := range m.ContentBlocks { + if block == nil || block.Type != schema.ContentBlockTypeFunctionToolResult || block.FunctionToolResult == nil { + continue + } + if found != nil { + return nil, false, fmt.Errorf("reduction: AgenticMessage contains multiple FunctionToolResult blocks; expected exactly one per message") + } + found = block.FunctionToolResult + } + if found == nil { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: ""}}}, true, nil + } + parts := toolResultToOutputParts(found) + if len(parts) == 0 { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: ""}}}, true, nil + } + isSimple := len(parts) == 1 && parts[0].Type == schema.ToolPartTypeText + return &schema.ToolResult{Parts: parts}, isSimple, nil + } + return nil, false, fmt.Errorf("unsupported message type") +} + +// setToolResultContent updates the tool result content in a message. +// For *schema.Message: sets msg.Content or msg.UserInputMultiContent. +// For *schema.AgenticMessage: reconstructs FunctionToolResult.Content. +func setToolResultContent[M adk.MessageType](msg M, toolResult *schema.ToolResult, fromContent bool) { + switch m := any(msg).(type) { + case *schema.Message: + if fromContent { + if len(toolResult.Parts) > 0 { + m.Content = toolResult.Parts[0].Text + } + } else { + convResult, convErr := toolResult.ToMessageInputParts() + if convErr == nil { + m.UserInputMultiContent = convResult + } + } + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if block == nil || block.Type != schema.ContentBlockTypeFunctionToolResult || block.FunctionToolResult == nil { + continue + } + setToolResultFromOutputParts(block.FunctionToolResult, toolResult.Parts) + return + } + } +} + +// copyMessagesGeneric deep-copies a slice of messages. +func copyMessagesGeneric[M adk.MessageType](msgs []M) []M { + var zero M + switch any(zero).(type) { + case *schema.Message: + origMsgs := any(msgs).([]*schema.Message) + copied := copyMessages(origMsgs) + return any(copied).([]M) + case *schema.AgenticMessage: + origMsgs := any(msgs).([]*schema.AgenticMessage) + copied := copyAgenticMessages(origMsgs) + return any(copied).([]M) + } + panic("unreachable") +} + +func copyAgenticMessages(msgs []*schema.AgenticMessage) []*schema.AgenticMessage { + resp := make([]*schema.AgenticMessage, len(msgs)) + for i, msg := range msgs { + if msg == nil { + continue + } + copied := &schema.AgenticMessage{ + Role: msg.Role, + ResponseMeta: msg.ResponseMeta, + } + if msg.ContentBlocks != nil { + copied.ContentBlocks = make([]*schema.ContentBlock, len(msg.ContentBlocks)) + for j, block := range msg.ContentBlocks { + if block == nil { + continue + } + cb := *block + // Deep copy mutable sub-fields + if block.FunctionToolCall != nil { + ftc := *block.FunctionToolCall + cb.FunctionToolCall = &ftc + } + if block.FunctionToolResult != nil { + ftr := *block.FunctionToolResult + if block.FunctionToolResult.Content != nil { + ftr.Content = make([]*schema.FunctionToolResultContentBlock, len(block.FunctionToolResult.Content)) + for k, rb := range block.FunctionToolResult.Content { + if rb != nil { + rbCopy := *rb // shallow copy: Image/Audio/Video/File sub-fields are not deep-copied. + // This is safe because the clear logic replaces entire blocks rather than + // mutating media fields in-place. Custom ClearHandlers should follow the same pattern. + if rb.Text != nil { + t := *rb.Text + rbCopy.Text = &t + } + ftr.Content[k] = &rbCopy + } + } + } + cb.FunctionToolResult = &ftr + } + if block.Extra != nil { + cb.Extra = make(map[string]any, len(block.Extra)) + for k, v := range block.Extra { + cb.Extra[k] = v + } + } + copied.ContentBlocks[j] = &cb + } + } + if msg.Extra != nil { + copied.Extra = make(map[string]any, len(msg.Extra)) + for k, v := range msg.Extra { + copied.Extra[k] = v + } + } + resp[i] = copied + } + return resp +} + +func copyMessages(msgs []*schema.Message) []*schema.Message { + resp := make([]*schema.Message, len(msgs)) + for i, msg := range msgs { + if msg == nil { + continue + } + copied := &schema.Message{ + Role: msg.Role, + Content: msg.Content, + MultiContent: msg.MultiContent, + UserInputMultiContent: msg.UserInputMultiContent, + AssistantGenMultiContent: msg.AssistantGenMultiContent, + Name: msg.Name, + ToolCalls: nil, + ToolCallID: msg.ToolCallID, + ToolName: msg.ToolName, + ResponseMeta: msg.ResponseMeta, + ReasoningContent: msg.ReasoningContent, + Extra: nil, + } + if msg.ToolCalls != nil { + copied.ToolCalls = append(make([]schema.ToolCall, 0, len(msg.ToolCalls)), msg.ToolCalls...) + } + if msg.Extra != nil { + copied.Extra = make(map[string]any, len(msg.Extra)) + for k, v := range msg.Extra { + copied.Extra[k] = v + } + } + resp[i] = copied + } + return resp +} + +// defaultTokenCounter estimates tokens, which treats one token as ~4 characters of text for common English text. +// github.com/tiktoken-go/tokenizer is highly recommended to replace it. +func defaultTokenCounter(_ context.Context, msgs []*schema.Message, tools []*schema.ToolInfo) (int64, error) { + var tokens int64 + for _, msg := range msgs { + if msg == nil { + continue + } + + var sb strings.Builder + sb.WriteString(string(msg.Role)) + sb.WriteString("\n") + sb.WriteString(msg.ReasoningContent) + sb.WriteString("\n") + sb.WriteString(msg.Content) + sb.WriteString("\n") + if msg.Role == schema.Assistant && len(msg.ToolCalls) > 0 { + for _, tc := range msg.ToolCalls { + sb.WriteString(tc.Function.Name) + sb.WriteString("\n") + sb.WriteString(tc.Function.Arguments) + } + } + + for _, mc := range msg.UserInputMultiContent { + switch mc.Type { + case schema.ChatMessagePartTypeText: + sb.WriteString(mc.Text) + sb.WriteString("\n") + default: + // do nothing for multi-modal content + } + } + + for _, mc := range msg.AssistantGenMultiContent { + switch mc.Type { + case schema.ChatMessagePartTypeText: + sb.WriteString(mc.Text) + sb.WriteString("\n") + default: + // do nothing for multi-modal content + } + } + + n := int64(len(sb.String()) / 4) + tokens += n + } + + for _, tl := range tools { + tl_ := *tl + tl_.Extra = nil + text, err := sonic.MarshalString(tl_) + if err != nil { + return 0, fmt.Errorf("failed to marshal tool info: %w", err) + } + + tokens += int64(len(text) / 4) + } + + return tokens, nil +} + +// defaultTruncHandler applies the shared buffered truncation strategy. The +// built-in default streaming wrappers bypass it so they can truncate +// incrementally with a streaming-specific notice. +// +// Processing steps: +// 1. Read and join tool output into a complete result: +// - Non-streaming: use ToolResult directly. +// - Streaming: consume the whole StreamToolResult, then concat all chunks. +// 2. If output is empty or total text length does not exceed truncMaxLength, +// return NeedTrunc=false. +// 3. If exceeded, replace oversized text parts with truncation notices and +// offload the full original content. +// +// Streaming-specific behavior: +// - Truncation is not incremental. The handler waits until the entire stream is read +// before deciding and producing output. +// - If stream Recv() returns a non-EOF error, getJointToolResult treats it as +// "skip processing" (needProcess=false, err=nil), so this handler returns +// NeedTrunc=false and does not propagate that recv error. +// - When truncation is applied to a streaming tool result, output is re-emitted as a +// buffered single-result stream (not original chunk-by-chunk streaming semantics). +// +// Tool-specific configurations that install this handler for streaming outputs +// keep the buffered behavior below. +func defaultTruncHandler( + genOffloadFilePathFn func(ctx context.Context, toolDetail *ToolDetail) (filePath string, err error), + truncMaxLength int, + readFileToolName string, +) func(ctx context.Context, detail *ToolDetail) (truncResult *TruncResult, err error) { + + return func(ctx context.Context, detail *ToolDetail) (offloadInfo *TruncResult, err error) { + isStreamResult := detail.StreamToolResult != nil + resultParts, needProcess, err := getJointToolResult(detail) + if err != nil { + return nil, err + } + if !needProcess { + return &TruncResult{NeedTrunc: false}, nil + } + + fullLength, textPartsCnt := 0, 0 + for _, part := range resultParts { + if part.Type == schema.ToolPartTypeText { + fullLength += len(part.Text) + textPartsCnt++ + } + } + if textPartsCnt == 0 || fullLength <= truncMaxLength { + return &TruncResult{NeedTrunc: false}, nil + } + + var ( + offloadContent = stringifyToolOutputParts(resultParts) + truncPartLength = truncMaxLength / textPartsCnt + previewSize = truncPartLength / 2 + ) + + filePath, err := genOffloadFilePathFn(ctx, detail) + if err != nil { + return nil, err + } + + for i, part := range resultParts { + text := part.Text + if part.Type != schema.ToolPartTypeText || + len(text) < truncPartLength { + continue + } + truncNotify, fmtErr := pyfmt.Fmt(getTruncFmt(), map[string]any{ + "original_size": len(part.Text), + "file_path": filePath, + "preview_size": previewSize, + "preview_first": clampPrefixToUTF8Boundary(text, previewSize), + "preview_last": clampSuffixToUTF8Boundary(text, previewSize), + "read_tool_name": readFileToolName, + }) + if fmtErr != nil { + return nil, fmtErr + } + resultParts[i].Text = truncNotify + } + + tr := &TruncResult{ + NeedTrunc: true, + NeedOffload: true, + OffloadFilePath: filePath, + OffloadContent: offloadContent, + } + if !isStreamResult { + tr.ToolResult = &schema.ToolResult{Parts: resultParts} + } else { + sr, sw := schema.Pipe[*schema.ToolResult](1) + sw.Send(&schema.ToolResult{Parts: resultParts}, nil) + sw.Close() + tr.StreamToolResult = sr + } + return tr, nil + } +} + +func clampPrefixToUTF8Boundary(s string, n int) string { + if n <= 0 { + return "" + } + if n >= len(s) { + return s + } + for n > 0 && !utf8.RuneStart(s[n]) { + n-- + } + return s[:n] +} + +func clampSuffixToUTF8Boundary(s string, n int) string { + if n <= 0 { + return "" + } + l := len(s) + if n >= l { + return s + } + start := l - n + for start < l && !utf8.RuneStart(s[start]) { + start++ + } + return s[start:] +} + +func defaultClearHandler( + genOffloadFilePathFn func(ctx context.Context, toolDetail *ToolDetail) (filePath string, err error), + needOffload bool, + readFileToolName string, +) func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) { + + return func(ctx context.Context, detail *ToolDetail) (clearResult *ClearResult, err error) { + resultParts, needProcess, err := getJointToolResult(detail) + if err != nil { + return nil, err + } + if !needProcess { + return &ClearResult{NeedClear: false}, nil + } + + if needOffload { + filePath, err := genOffloadFilePathFn(ctx, detail) + if err != nil { + return nil, err + } + textPlaceHolder, fmtErr := pyfmt.Fmt(getClearWithOffloadingFmt(), map[string]any{ + "file_path": filePath, + "read_tool_name": readFileToolName, + }) + if fmtErr != nil { + return nil, fmtErr + } + + offloadContent := stringifyToolOutputParts(resultParts) + for i, part := range resultParts { + if part.Type != schema.ToolPartTypeText { + continue + } + resultParts[i].Text = textPlaceHolder + } + clearResult = &ClearResult{ + NeedClear: true, + ToolArgument: detail.ToolArgument, + ToolResult: &schema.ToolResult{Parts: resultParts}, + NeedOffload: true, + OffloadFilePath: filePath, + OffloadContent: offloadContent, + } + } else { + textPlaceHolder := getClearWithoutOffloadingFmt() + for i, part := range resultParts { + if part.Type != schema.ToolPartTypeText { + continue + } + resultParts[i].Text = textPlaceHolder + } + clearResult = &ClearResult{ + NeedClear: true, + ToolArgument: detail.ToolArgument, + ToolResult: &schema.ToolResult{Parts: resultParts}, + NeedOffload: false, + } + } + + return clearResult, nil + } +} + +func getJointToolResult(toolDetail *ToolDetail) (toolOutputParts []schema.ToolOutputPart, needProcess bool, err error) { + if toolDetail.ToolResult == nil && toolDetail.StreamToolResult == nil { + return nil, false, fmt.Errorf("ToolResult and StreamToolResult are both nil") + } + + if toolDetail.ToolResult != nil { + toolOutputParts = toolDetail.ToolResult.Parts + } else { + var toolResultChunks []*schema.ToolResult + for { + toolResultChunk, recvErr := toolDetail.StreamToolResult.Recv() + if recvErr != nil { + if recvErr == io.EOF { + break + } + // return original stream reader, not sending recvErr + return nil, false, nil + } + toolResultChunks = append(toolResultChunks, toolResultChunk) + } + toolResult, concatErr := schema.ConcatToolResults(toolResultChunks) + if concatErr != nil { + return nil, false, concatErr + } + toolOutputParts = toolResult.Parts + } + + if len(toolOutputParts) == 0 { + return nil, false, nil + } + + return toolOutputParts, true, nil +} + +func stringifyToolOutputParts(toolOutputParts []schema.ToolOutputPart) string { + if len(toolOutputParts) == 0 { + return "" + } else if len(toolOutputParts) == 1 && toolOutputParts[0].Type == schema.ToolPartTypeText { + return toolOutputParts[0].Text + } else { + b, _ := json.MarshalIndent(toolOutputParts, "", "\t") + return string(b) + } +} + +func getMsgClearedFlag(msg *schema.Message) (offloaded bool) { + if msg.Extra == nil { + return false + } + v, ok := msg.Extra[msgClearedFlag].(bool) + if !ok { + return false + } + return v +} + +func setMsgClearedFlag(msg *schema.Message) { + if msg.Extra == nil { + msg.Extra = make(map[string]any) + } + msg.Extra[msgClearedFlag] = true +} + +func toolResultFromMessage(msg *schema.Message) (result *schema.ToolResult, fromContent bool, err error) { + if msg.Role != schema.Tool { + return nil, false, fmt.Errorf("message role %s is not a tool", msg.Role) + } + if len(msg.UserInputMultiContent) > 0 { + result = &schema.ToolResult{Parts: make([]schema.ToolOutputPart, 0, len(msg.UserInputMultiContent))} + for _, part := range msg.UserInputMultiContent { + top, convErr := convMessageInputPartToToolOutputPart(part) + if convErr != nil { + return nil, false, convErr + } + result.Parts = append(result.Parts, top) + } + return result, false, nil + } + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: msg.Content}}}, true, nil +} + +func convMessageInputPartToToolOutputPart(msgPart schema.MessageInputPart) (schema.ToolOutputPart, error) { + switch msgPart.Type { + case schema.ChatMessagePartTypeText: + return schema.ToolOutputPart{ + Type: schema.ToolPartTypeText, + Text: msgPart.Text, + }, nil + case schema.ChatMessagePartTypeImageURL: + return schema.ToolOutputPart{ + Type: schema.ToolPartTypeImage, + Image: &schema.ToolOutputImage{ + MessagePartCommon: msgPart.Image.MessagePartCommon, + }, + }, nil + case schema.ChatMessagePartTypeAudioURL: + return schema.ToolOutputPart{ + Type: schema.ToolPartTypeAudio, + Audio: &schema.ToolOutputAudio{ + MessagePartCommon: msgPart.Audio.MessagePartCommon, + }, + }, nil + case schema.ChatMessagePartTypeVideoURL: + return schema.ToolOutputPart{ + Type: schema.ToolPartTypeVideo, + Video: &schema.ToolOutputVideo{ + MessagePartCommon: msgPart.Video.MessagePartCommon, + }, + }, nil + case schema.ChatMessagePartTypeFileURL: + return schema.ToolOutputPart{ + Type: schema.ToolPartTypeFile, + File: &schema.ToolOutputFile{ + MessagePartCommon: msgPart.File.MessagePartCommon, + }, + }, nil + default: + return schema.ToolOutputPart{}, fmt.Errorf("unknown msg part type: %v", msgPart.Type) + } +} + +// toolResultToOutputParts converts a FunctionToolResult's Content blocks to ToolOutputPart slice. +func toolResultToOutputParts(f *schema.FunctionToolResult) []schema.ToolOutputPart { + var parts []schema.ToolOutputPart + for _, block := range f.Content { + if block == nil { + continue + } + if block.Text != nil { + parts = append(parts, schema.ToolOutputPart{Type: schema.ToolPartTypeText, Text: block.Text.Text}) + } else if block.Image != nil { + parts = append(parts, schema.ToolOutputPart{ + Type: schema.ToolPartTypeImage, + Image: &schema.ToolOutputImage{MessagePartCommon: schema.MessagePartCommon{URL: strPtr(block.Image.URL), MIMEType: block.Image.MIMEType}}, + }) + } else if block.Audio != nil { + parts = append(parts, schema.ToolOutputPart{ + Type: schema.ToolPartTypeAudio, + Audio: &schema.ToolOutputAudio{MessagePartCommon: schema.MessagePartCommon{URL: strPtr(block.Audio.URL), MIMEType: block.Audio.MIMEType}}, + }) + } else if block.Video != nil { + parts = append(parts, schema.ToolOutputPart{ + Type: schema.ToolPartTypeVideo, + Video: &schema.ToolOutputVideo{MessagePartCommon: schema.MessagePartCommon{URL: strPtr(block.Video.URL), MIMEType: block.Video.MIMEType}}, + }) + } else if block.File != nil { + parts = append(parts, schema.ToolOutputPart{ + Type: schema.ToolPartTypeFile, + File: &schema.ToolOutputFile{MessagePartCommon: schema.MessagePartCommon{URL: strPtr(block.File.URL), MIMEType: block.File.MIMEType}}, + }) + } + } + return parts +} + +// setToolResultFromOutputParts converts ToolOutputPart slice back to FunctionToolResultContentBlock +// slice and sets f.Content. +func setToolResultFromOutputParts(f *schema.FunctionToolResult, parts []schema.ToolOutputPart) { + var newBlocks []*schema.FunctionToolResultContentBlock + for _, part := range parts { + switch part.Type { + case schema.ToolPartTypeText: + newBlocks = append(newBlocks, &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: part.Text}, + }) + case schema.ToolPartTypeImage: + if part.Image != nil { + newBlocks = append(newBlocks, &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeImage, + Image: &schema.UserInputImage{URL: ptrStr(part.Image.URL), MIMEType: part.Image.MIMEType}, + }) + } + case schema.ToolPartTypeAudio: + if part.Audio != nil { + newBlocks = append(newBlocks, &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeAudio, + Audio: &schema.UserInputAudio{URL: ptrStr(part.Audio.URL), MIMEType: part.Audio.MIMEType}, + }) + } + case schema.ToolPartTypeVideo: + if part.Video != nil { + newBlocks = append(newBlocks, &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeVideo, + Video: &schema.UserInputVideo{URL: ptrStr(part.Video.URL), MIMEType: part.Video.MIMEType}, + }) + } + case schema.ToolPartTypeFile: + if part.File != nil { + newBlocks = append(newBlocks, &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeFile, + File: &schema.UserInputFile{URL: ptrStr(part.File.URL), MIMEType: part.File.MIMEType}, + }) + } + } + } + f.Content = newBlocks +} + +// strPtr returns a pointer to s, or nil if s is empty. +func strPtr(s string) *string { + if s == "" { + return nil + } + return &s +} + +// ptrStr safely dereferences a *string, returning "" if nil. +func ptrStr(p *string) string { + if p == nil { + return "" + } + return *p +} diff --git a/adk/middlewares/reduction/reduction_generic_test.go b/adk/middlewares/reduction/reduction_generic_test.go new file mode 100644 index 0000000..9953f47 --- /dev/null +++ b/adk/middlewares/reduction/reduction_generic_test.go @@ -0,0 +1,955 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package reduction + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// --------------------------------------------------------------------------- +// Generic message construction helpers +// --------------------------------------------------------------------------- + +type testToolCall struct { + ID string + Name string + Arguments string +} + +func makeUserMsgG[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(content)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(content)).(M) + } + panic("unreachable") +} + +func makeSystemMsgG[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{Role: schema.System, Content: content}).(M) + case *schema.AgenticMessage: + return any(schema.SystemAgenticMessage(content)).(M) + } + panic("unreachable") +} + +func makeAssistantMsgG[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(&schema.Message{Role: schema.Assistant, Content: content}).(M) + case *schema.AgenticMessage: + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.AssistantGenText{Text: content})}, + }).(M) + } + panic("unreachable") +} + +func makeAssistantMsgWithToolCallsG[M adk.MessageType](toolCalls []testToolCall) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + tcs := make([]schema.ToolCall, len(toolCalls)) + for i, tc := range toolCalls { + tcs[i] = schema.ToolCall{ + ID: tc.ID, + Type: "function", + Function: schema.FunctionCall{Name: tc.Name, Arguments: tc.Arguments}, + } + } + return any(schema.AssistantMessage("", tcs)).(M) + case *schema.AgenticMessage: + blocks := make([]*schema.ContentBlock, 0, len(toolCalls)) + for _, tc := range toolCalls { + blocks = append(blocks, schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: tc.ID, + Name: tc.Name, + Arguments: tc.Arguments, + })) + } + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: blocks, + }).(M) + } + panic("unreachable") +} + +func makeToolResultMsgG[M adk.MessageType](content string, callID string, toolName string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.ToolMessage(content, callID) + msg.ToolName = toolName + return any(msg).(M) + case *schema.AgenticMessage: + return any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: callID, + Name: toolName, + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: content}}, + }, + }), + }, + }).(M) + } + panic("unreachable") +} + +func getMsgContentG[M adk.MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return v.Content + case *schema.AgenticMessage: + for _, block := range v.ContentBlocks { + if block == nil { + continue + } + if block.UserInputText != nil { + return block.UserInputText.Text + } + if block.AssistantGenText != nil { + return block.AssistantGenText.Text + } + if block.FunctionToolResult != nil { + for _, b := range block.FunctionToolResult.Content { + if b != nil && b.Text != nil { + return b.Text.Text + } + } + } + } + return "" + } + panic("unreachable") +} + +// --------------------------------------------------------------------------- +// Part 1: Helper function tests +// --------------------------------------------------------------------------- + +func testHelperFunctions[M adk.MessageType](t *testing.T) { + t.Run("isAssistantMsg", func(t *testing.T) { + assistant := makeAssistantMsgG[M]("hello") + user := makeUserMsgG[M]("hello") + assert.True(t, isAssistantMsg(assistant)) + assert.False(t, isAssistantMsg(user)) + }) + + t.Run("isSystemMsg", func(t *testing.T) { + sys := makeSystemMsgG[M]("system prompt") + user := makeUserMsgG[M]("hello") + assert.True(t, isSystemMsg(sys)) + assert.False(t, isSystemMsg(user)) + }) + + t.Run("isUserMsg", func(t *testing.T) { + user := makeUserMsgG[M]("hello") + assert.True(t, isUserMsg(user)) + + // A user message that only has tool results should return false. + toolResultOnly := makeToolResultMsgG[M]("result", "call_1", "my_tool") + assert.False(t, isUserMsg(toolResultOnly)) + }) + + t.Run("hasToolCalls", func(t *testing.T) { + withTC := makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "c1", Name: "tool1", Arguments: `{"a":1}`}, + }) + assert.True(t, hasToolCalls(withTC)) + + noTC := makeAssistantMsgG[M]("plain response") + assert.False(t, hasToolCalls(noTC)) + }) + + t.Run("isToolResultMsg", func(t *testing.T) { + tr := makeToolResultMsgG[M]("result content", "call_1", "my_tool") + assert.True(t, isToolResultMsg(tr)) + + user := makeUserMsgG[M]("not a tool result") + assert.False(t, isToolResultMsg(user)) + }) + + t.Run("isToolResultOnlyMsg", func(t *testing.T) { + trOnly := makeToolResultMsgG[M]("result content", "call_1", "my_tool") + assert.True(t, isToolResultOnlyMsg(trOnly)) + + // A normal user message is not a tool-result-only message. + user := makeUserMsgG[M]("hello") + assert.False(t, isToolResultOnlyMsg(user)) + + // For AgenticMessage, a mixed message (user text + tool result) should return false. + var zero M + if _, ok := any(zero).(*schema.AgenticMessage); ok { + mixed := any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.UserInputText{Text: "hello"}), + schema.NewContentBlock(&schema.FunctionToolResult{CallID: "c1", Name: "tool1", Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "result"}}, + }}), + }, + }).(M) + assert.False(t, isToolResultOnlyMsg(mixed)) + } + }) + + t.Run("getMsgClearedFlagGeneric_setMsgClearedFlagGeneric", func(t *testing.T) { + msg := makeAssistantMsgG[M]("test") + assert.False(t, getMsgClearedFlagGeneric(msg)) + + setMsgClearedFlagGeneric(msg) + assert.True(t, getMsgClearedFlagGeneric(msg)) + }) + + t.Run("getToolCallsGeneric", func(t *testing.T) { + tcs := []testToolCall{ + {ID: "call_a", Name: "tool_alpha", Arguments: `{"x":1}`}, + {ID: "call_b", Name: "tool_beta", Arguments: `{"y":2}`}, + } + msg := makeAssistantMsgWithToolCallsG[M](tcs) + got := getToolCallsGeneric(msg) + require.Len(t, got, 2) + + assert.Equal(t, "call_a", got[0].CallID) + assert.Equal(t, "tool_alpha", got[0].Name) + assert.Equal(t, `{"x":1}`, got[0].Arguments) + assert.Equal(t, 0, got[0].BlockIndex) + + assert.Equal(t, "call_b", got[1].CallID) + assert.Equal(t, "tool_beta", got[1].Name) + assert.Equal(t, `{"y":2}`, got[1].Arguments) + assert.Equal(t, 1, got[1].BlockIndex) + + // Empty assistant message returns nil. + noTC := makeAssistantMsgG[M]("plain") + assert.Nil(t, getToolCallsGeneric(noTC)) + }) + + t.Run("setToolCallArguments", func(t *testing.T) { + tcs := []testToolCall{ + {ID: "call_a", Name: "tool_alpha", Arguments: `{"old":"args"}`}, + } + msg := makeAssistantMsgWithToolCallsG[M](tcs) + setToolCallArguments(msg, 0, `{"new":"args"}`) + + got := getToolCallsGeneric(msg) + require.Len(t, got, 1) + assert.Equal(t, `{"new":"args"}`, got[0].Arguments) + + // Verify AgenticMessage path writes to the ContentBlock directly. + if am, ok := any(msg).(*schema.AgenticMessage); ok { + require.NotNil(t, am.ContentBlocks[0].FunctionToolCall) + assert.Equal(t, `{"new":"args"}`, am.ContentBlocks[0].FunctionToolCall.Arguments) + } + }) + + t.Run("copyMessagesGeneric", func(t *testing.T) { + original := []M{ + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "c1", Name: "t1", Arguments: `{"k":"v"}`}, + }), + makeUserMsgG[M]("user text"), + } + copied := copyMessagesGeneric(original) + require.Len(t, copied, 2) + + // Modify the copy's tool call arguments. + setToolCallArguments(copied[0], 0, `{"modified":"true"}`) + + // Original must be unchanged. + origTCs := getToolCallsGeneric(original[0]) + require.Len(t, origTCs, 1) + assert.Equal(t, `{"k":"v"}`, origTCs[0].Arguments, "original must not be affected by copy mutation") + + copiedTCs := getToolCallsGeneric(copied[0]) + assert.Equal(t, `{"modified":"true"}`, copiedTCs[0].Arguments) + }) + + t.Run("getToolResultCallID", func(t *testing.T) { + // Tool result message with matching callID + tr := makeToolResultMsgG[M]("result content", "call_123", "my_tool") + assert.Equal(t, "call_123", getToolResultCallID(tr)) + + // Non-tool-result message should return empty string + user := makeUserMsgG[M]("hello") + assert.Equal(t, "", getToolResultCallID(user)) + }) + + t.Run("findToolResultByCallID", func(t *testing.T) { + messages := []M{ + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_1", Name: "tool1", Arguments: `{}`}, + }), + makeToolResultMsgG[M]("result 1", "call_1", "tool1"), + makeToolResultMsgG[M]("result 2", "call_2", "tool2"), + } + + // Find existing tool result + idx, found := findToolResultByCallID(messages, 1, 3, "call_2") + assert.True(t, found) + assert.Equal(t, 2, idx) + + // Find non-existent tool result + idx, found = findToolResultByCallID(messages, 1, 3, "call_999") + assert.False(t, found) + assert.Equal(t, -1, idx) + + // Stop searching when encountering a non-tool-result message + messages2 := []M{ + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_3", Name: "tool3", Arguments: `{}`}, + }), + makeToolResultMsgG[M]("result 3", "call_3", "tool3"), + makeUserMsgG[M]("not a tool result"), + makeToolResultMsgG[M]("result 4", "call_4", "tool4"), + } + idx, found = findToolResultByCallID(messages2, 1, 4, "call_4") + assert.False(t, found) + assert.Equal(t, -1, idx) + }) +} + +// --------------------------------------------------------------------------- +// Part 2: Clear rewrite flow +// --------------------------------------------------------------------------- + +func testClearFlowGeneric[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + // Token counter that always returns a high count to trigger clearing. + highTokenCounter := func(_ context.Context, _ []M, _ []*schema.ToolInfo) (int64, error) { + return 999999, nil + } + + // ClearRetentionSuffixLimit defaults to 1 in copyAndFillDefaults when set to 0, + // so we explicitly set it to 1. This means the last tool-call group (call_new) + // is retained and only the older group (call_old) is cleared. + config := &TypedConfig[M]{ + SkipTruncation: true, + TokenCounter: highTokenCounter, + MaxTokensForClear: 100, + ClearRetentionSuffixLimit: 1, + } + + mw, err := NewTyped(ctx, config) + require.NoError(t, err) + + // Messages: system, user, assistant+toolcalls(old), tool_result(old), user, assistant+toolcalls(new) + msgs := []M{ + makeSystemMsgG[M]("you are helpful"), + makeUserMsgG[M]("what's the weather?"), + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_old", Name: "get_weather", Arguments: `{"location":"London"}`}, + }), + makeToolResultMsgG[M]("Sunny and warm", "call_old", "get_weather"), + makeUserMsgG[M]("set thermostat"), + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_new", Name: "set_thermostat", Arguments: `{"temp":20}`}, + }), + } + + state := &adk.TypedChatModelAgentState[M]{Messages: msgs} + _, resultState, err := mw.BeforeModelRewriteState(ctx, state, &adk.TypedModelContext[M]{}) + require.NoError(t, err) + require.Equal(t, 6, len(resultState.Messages)) + + // The default ClearHandler preserves tool call arguments (sets them to the original). + // Verify they are unchanged. + oldTCs := getToolCallsGeneric(resultState.Messages[2]) + require.Len(t, oldTCs, 1) + assert.Equal(t, `{"location":"London"}`, oldTCs[0].Arguments, "default handler preserves tool call arguments") + + // The old tool result (index 3) should have its content replaced with a placeholder. + // The placeholder text is locale-dependent, so just verify it changed from the original. + oldResultContent := getMsgContentG(resultState.Messages[3]) + assert.NotEqual(t, "Sunny and warm", oldResultContent, "old tool result content should be replaced with placeholder") + + // The cleared flag should be set on the old assistant message. + assert.True(t, getMsgClearedFlagGeneric(resultState.Messages[2]), "cleared flag should be set on old assistant msg") + + // System message (index 0) should be untouched. + assert.Equal(t, "you are helpful", getMsgContentG(resultState.Messages[0])) + + // Recent messages (index 4, 5) should not be affected: the new tool-call group + // is in the retention window. + newTCs := getToolCallsGeneric(resultState.Messages[5]) + require.Len(t, newTCs, 1) + assert.Equal(t, `{"temp":20}`, newTCs[0].Arguments, "recent tool calls must not be cleared") +} + +// --------------------------------------------------------------------------- +// Part 3: Truncation flow +// --------------------------------------------------------------------------- + +func testTruncationGeneric[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + callCount := 0 + // Token counter returns decreasing counts as messages shrink. + tokenCounter := func(_ context.Context, msgs []M, _ []*schema.ToolInfo) (int64, error) { + callCount++ + // First call: over limit. After truncation (fewer msgs), under limit. + return int64(len(msgs)) * 100, nil + } + + config := &TypedConfig[M]{ + SkipTruncation: true, + SkipClear: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 250, // 5 messages * 100 = 500 > 250 + ClearRetentionSuffixLimit: 0, + } + + mw, err := NewTyped(ctx, config) + require.NoError(t, err) + + msgs := []M{ + makeSystemMsgG[M]("system prompt"), + makeUserMsgG[M]("old user message"), + makeAssistantMsgG[M]("old assistant response"), + makeUserMsgG[M]("new user message"), + makeAssistantMsgG[M]("new assistant response"), + } + + state := &adk.TypedChatModelAgentState[M]{Messages: msgs} + _, resultState, err := mw.BeforeModelRewriteState(ctx, state, &adk.TypedModelContext[M]{}) + require.NoError(t, err) + + // Since SkipClear is true, the clear path is entirely skipped. + // The middleware should return the state unchanged because clear is skipped + // (truncation in BeforeModelRewriteState is the clear phase, not the tool-output truncation). + // The messages are returned as-is since the clearing loop is the only message-removal mechanism. + assert.Equal(t, len(msgs), len(resultState.Messages)) +} + +// --------------------------------------------------------------------------- +// Part 4: ClearPostProcess callback +// --------------------------------------------------------------------------- + +func testClearPostProcessGeneric[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + postProcessCalled := false + highTokenCounter := func(_ context.Context, _ []M, _ []*schema.ToolInfo) (int64, error) { + return 999999, nil + } + + // ClearRetentionSuffixLimit=0 defaults to 1 via copyAndFillDefaults. + // We need at least 2 tool-call groups so that the first one gets cleared + // while the second is retained by the suffix limit. + config := &TypedConfig[M]{ + SkipTruncation: true, + TokenCounter: highTokenCounter, + MaxTokensForClear: 100, + ClearRetentionSuffixLimit: 1, + ClearPostProcess: func(ctx context.Context, state *adk.TypedChatModelAgentState[M]) context.Context { + postProcessCalled = true + return ctx + }, + } + + mw, err := NewTyped(ctx, config) + require.NoError(t, err) + + msgs := []M{ + makeSystemMsgG[M]("system"), + makeUserMsgG[M]("user"), + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_1", Name: "tool1", Arguments: `{"a":"b"}`}, + }), + makeToolResultMsgG[M]("result", "call_1", "tool1"), + makeUserMsgG[M]("another request"), + makeAssistantMsgWithToolCallsG[M]([]testToolCall{ + {ID: "call_2", Name: "tool2", Arguments: `{"c":"d"}`}, + }), + makeToolResultMsgG[M]("result2", "call_2", "tool2"), + } + + state := &adk.TypedChatModelAgentState[M]{Messages: msgs} + _, _, err = mw.BeforeModelRewriteState(ctx, state, &adk.TypedModelContext[M]{}) + require.NoError(t, err) + assert.True(t, postProcessCalled, "ClearPostProcess should have been called") +} + +// --------------------------------------------------------------------------- +// Part 5: AgenticMessage-specific coverage +// --------------------------------------------------------------------------- + +func TestGetDefaultTokenCounter_AgenticMessage(t *testing.T) { + ctx := context.Background() + counter := getDefaultTokenCounter[*schema.AgenticMessage]() + + msgs := []*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.UserInputText{Text: "Hello, world!"}), + }, + }, + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "Hi there!"}), + }, + }, + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolCall{CallID: "c1", Name: "my_tool", Arguments: `{"key":"value"}`}), + }, + }, + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: "c1", + Name: "my_tool", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "tool output text"}}, + }, + }), + }, + }, + nil, // nil message should be skipped + } + + tokens, err := counter(ctx, msgs, nil) + assert.NoError(t, err) + assert.Greater(t, tokens, int64(0), "should count tokens from content blocks") + + // Also test with tools + tools := []*schema.ToolInfo{ + {Name: "my_tool", Desc: "a test tool"}, + } + tokensWithTools, err := counter(ctx, msgs, tools) + assert.NoError(t, err) + assert.Greater(t, tokensWithTools, tokens, "tokens should increase with tool info") +} + +func TestCopyAgenticMessages_DeepCopy(t *testing.T) { + original := []*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolCall{ + CallID: "call_1", + Name: "tool_a", + Arguments: `{"x":1}`, + }), + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "call_1", + Name: "tool_a", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "original result"}}, + }, + }, + Extra: map[string]any{"meta": "data"}, + }, + }, + Extra: map[string]any{"msg_key": "msg_value"}, + }, + } + + copied := copyMessagesGeneric(original) + require.Len(t, copied, 1) + + // Mutate the copy and verify original is unchanged. + copied[0].ContentBlocks[0].FunctionToolCall.Arguments = `{"modified":true}` + assert.Equal(t, `{"x":1}`, original[0].ContentBlocks[0].FunctionToolCall.Arguments, + "original FunctionToolCall.Arguments must not be affected") + + copied[0].ContentBlocks[1].FunctionToolResult.Content[0].Text.Text = "modified result" + assert.Equal(t, "original result", original[0].ContentBlocks[1].FunctionToolResult.Content[0].Text.Text, + "original FunctionToolResult text must not be affected") + + copied[0].ContentBlocks[1].Extra["meta"] = "changed" + assert.Equal(t, "data", original[0].ContentBlocks[1].Extra["meta"], + "original ContentBlock.Extra must not be affected") + + copied[0].Extra["msg_key"] = "changed" + assert.Equal(t, "msg_value", original[0].Extra["msg_key"], + "original AgenticMessage.Extra must not be affected") +} + +func TestToolResultFromMsgGeneric_AgenticMessage(t *testing.T) { + t.Run("single text block returns fromContent=true", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "tool1", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "hello result"}}, + }, + }, + }, + }, + } + + result, fromContent, err := toolResultFromMsgGeneric(msg) + assert.NoError(t, err) + assert.True(t, fromContent, "single text part should be fromContent=true") + require.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "hello result", result.Parts[0].Text) + }) + + t.Run("multiple blocks returns fromContent=false", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "tool1", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "text part"}}, + {Text: &schema.UserInputText{Text: "another text part"}}, + }, + }, + }, + }, + } + + result, fromContent, err := toolResultFromMsgGeneric(msg) + assert.NoError(t, err) + assert.False(t, fromContent, "multiple parts should be fromContent=false") + require.Len(t, result.Parts, 2) + assert.Equal(t, "text part", result.Parts[0].Text) + assert.Equal(t, "another text part", result.Parts[1].Text) + }) + + t.Run("empty blocks returns empty text", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "tool1", + Content: nil, + }, + }, + }, + } + + result, fromContent, err := toolResultFromMsgGeneric(msg) + assert.NoError(t, err) + assert.True(t, fromContent) + require.Len(t, result.Parts, 1) + assert.Equal(t, "", result.Parts[0].Text) + }) +} + +func TestSetToolResultContent_AgenticMessage(t *testing.T) { + t.Run("fromContent=true sets text", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "tool1", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "old"}}, + }, + }, + }, + }, + } + + newResult := &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "new content"}, + }, + } + + setToolResultContent(msg, newResult, true) + + // Verify the block was updated + blocks := msg.ContentBlocks[0].FunctionToolResult.Content + require.Len(t, blocks, 1) + assert.Equal(t, "new content", blocks[0].Text.Text) + }) + + t.Run("fromContent=false sets multi-part", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "tool1", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "old"}}, + }, + }, + }, + }, + } + + imgURL := "https://example.com/img.png" + newResult := &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "text part"}, + {Type: schema.ToolPartTypeImage, Image: &schema.ToolOutputImage{ + MessagePartCommon: schema.MessagePartCommon{URL: &imgURL, MIMEType: "image/png"}, + }}, + }, + } + + setToolResultContent(msg, newResult, false) + + blocks := msg.ContentBlocks[0].FunctionToolResult.Content + require.Len(t, blocks, 2) + assert.Equal(t, "text part", blocks[0].Text.Text) + require.NotNil(t, blocks[1].Image) + assert.Equal(t, "https://example.com/img.png", blocks[1].Image.URL) + assert.Equal(t, "image/png", blocks[1].Image.MIMEType) + }) +} + +func TestToolResultFromMsgGeneric_MediaBlocks(t *testing.T) { + imgURL := "https://example.com/img.png" + audioURL := "https://example.com/audio.wav" + videoURL := "https://example.com/video.mp4" + fileURL := "https://example.com/doc.pdf" + + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "media_tool", + Content: []*schema.FunctionToolResultContentBlock{ + {Image: &schema.UserInputImage{URL: imgURL, MIMEType: "image/png"}}, + {Audio: &schema.UserInputAudio{URL: audioURL, MIMEType: "audio/wav"}}, + {Video: &schema.UserInputVideo{URL: videoURL, MIMEType: "video/mp4"}}, + {File: &schema.UserInputFile{URL: fileURL, MIMEType: "application/pdf"}}, + }, + }, + }, + }, + } + + result, fromContent, err := toolResultFromMsgGeneric(msg) + assert.NoError(t, err) + assert.False(t, fromContent, "multi-media should be fromContent=false") + require.Len(t, result.Parts, 4) + + assert.Equal(t, schema.ToolPartTypeImage, result.Parts[0].Type) + require.NotNil(t, result.Parts[0].Image) + require.NotNil(t, result.Parts[0].Image.URL) + assert.Equal(t, imgURL, *result.Parts[0].Image.URL) + + assert.Equal(t, schema.ToolPartTypeAudio, result.Parts[1].Type) + require.NotNil(t, result.Parts[1].Audio) + require.NotNil(t, result.Parts[1].Audio.URL) + assert.Equal(t, audioURL, *result.Parts[1].Audio.URL) + + assert.Equal(t, schema.ToolPartTypeVideo, result.Parts[2].Type) + require.NotNil(t, result.Parts[2].Video) + require.NotNil(t, result.Parts[2].Video.URL) + assert.Equal(t, videoURL, *result.Parts[2].Video.URL) + + assert.Equal(t, schema.ToolPartTypeFile, result.Parts[3].Type) + require.NotNil(t, result.Parts[3].File) + require.NotNil(t, result.Parts[3].File.URL) + assert.Equal(t, fileURL, *result.Parts[3].File.URL) +} + +func TestSetToolResultContent_MediaBlocks(t *testing.T) { + audioURL := "https://example.com/speech.mp3" + videoURL := "https://example.com/clip.mp4" + fileURL := "https://example.com/report.pdf" + + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "c1", + Name: "tool1", + Content: []*schema.FunctionToolResultContentBlock{ + {Text: &schema.UserInputText{Text: "old"}}, + }, + }, + }, + }, + } + + newResult := &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeAudio, Audio: &schema.ToolOutputAudio{ + MessagePartCommon: schema.MessagePartCommon{URL: &audioURL, MIMEType: "audio/mp3"}, + }}, + {Type: schema.ToolPartTypeVideo, Video: &schema.ToolOutputVideo{ + MessagePartCommon: schema.MessagePartCommon{URL: &videoURL, MIMEType: "video/mp4"}, + }}, + {Type: schema.ToolPartTypeFile, File: &schema.ToolOutputFile{ + MessagePartCommon: schema.MessagePartCommon{URL: &fileURL, MIMEType: "application/pdf"}, + }}, + }, + } + + setToolResultContent(msg, newResult, false) + + blocks := msg.ContentBlocks[0].FunctionToolResult.Content + require.Len(t, blocks, 3) + + require.NotNil(t, blocks[0].Audio) + assert.Equal(t, "https://example.com/speech.mp3", blocks[0].Audio.URL) + assert.Equal(t, "audio/mp3", blocks[0].Audio.MIMEType) + + require.NotNil(t, blocks[1].Video) + assert.Equal(t, "https://example.com/clip.mp4", blocks[1].Video.URL) + assert.Equal(t, "video/mp4", blocks[1].Video.MIMEType) + + require.NotNil(t, blocks[2].File) + assert.Equal(t, "https://example.com/report.pdf", blocks[2].File.URL) + assert.Equal(t, "application/pdf", blocks[2].File.MIMEType) +} + +func TestAgenticURLToMPC(t *testing.T) { + t.Run("non-empty URL", func(t *testing.T) { + ftr := &schema.FunctionToolResult{ + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeFile, File: &schema.UserInputFile{URL: "https://example.com/file.pdf", MIMEType: "application/pdf"}}, + }, + } + parts := toolResultToOutputParts(ftr) + require.Len(t, parts, 1) + require.NotNil(t, parts[0].File) + require.NotNil(t, parts[0].File.URL) + assert.Equal(t, "https://example.com/file.pdf", *parts[0].File.URL) + assert.Equal(t, "application/pdf", parts[0].File.MIMEType) + }) + + t.Run("empty URL", func(t *testing.T) { + ftr := &schema.FunctionToolResult{ + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeFile, File: &schema.UserInputFile{URL: "", MIMEType: "text/plain"}}, + }, + } + parts := toolResultToOutputParts(ftr) + require.Len(t, parts, 1) + require.NotNil(t, parts[0].File) + assert.Nil(t, parts[0].File.URL) + assert.Equal(t, "text/plain", parts[0].File.MIMEType) + }) +} + +func TestMpcURLToString(t *testing.T) { + t.Run("non-nil URL", func(t *testing.T) { + urlStr := "https://example.com" + tr := &schema.FunctionToolResult{} + setToolResultFromOutputParts(tr, []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeFile, File: &schema.ToolOutputFile{MessagePartCommon: schema.MessagePartCommon{URL: &urlStr, MIMEType: "text/plain"}}}, + }) + require.Len(t, tr.Content, 1) + require.NotNil(t, tr.Content[0].File) + assert.Equal(t, "https://example.com", tr.Content[0].File.URL) + }) + + t.Run("nil URL", func(t *testing.T) { + tr := &schema.FunctionToolResult{} + setToolResultFromOutputParts(tr, []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeFile, File: &schema.ToolOutputFile{MessagePartCommon: schema.MessagePartCommon{URL: nil, MIMEType: "text/plain"}}}, + }) + require.Len(t, tr.Content, 1) + require.NotNil(t, tr.Content[0].File) + assert.Equal(t, "", tr.Content[0].File.URL) + }) +} + +// --------------------------------------------------------------------------- +// Top-level test +// --------------------------------------------------------------------------- + +func TestReductionGeneric(t *testing.T) { + t.Run("Message", func(t *testing.T) { + t.Run("Helpers", testHelperFunctions[*schema.Message]) + t.Run("ClearFlow", testClearFlowGeneric[*schema.Message]) + t.Run("Truncation", testTruncationGeneric[*schema.Message]) + t.Run("ClearPostProcess", testClearPostProcessGeneric[*schema.Message]) + t.Run("CopyNilMessage", testCopyNilMessage[*schema.Message]) + }) + t.Run("AgenticMessage", func(t *testing.T) { + t.Run("Helpers", testHelperFunctions[*schema.AgenticMessage]) + t.Run("ClearFlow", testClearFlowGeneric[*schema.AgenticMessage]) + t.Run("Truncation", testTruncationGeneric[*schema.AgenticMessage]) + t.Run("ClearPostProcess", testClearPostProcessGeneric[*schema.AgenticMessage]) + t.Run("CopyNilMessage", testCopyNilMessage[*schema.AgenticMessage]) + }) +} + +// testCopyNilMessage verifies that copyMessagesGeneric does not panic when +// the input slice contains nil message elements (regression test). +func testCopyNilMessage[M adk.MessageType](t *testing.T) { + var zero M + var msgs []M + switch any(zero).(type) { + case *schema.Message: + msgs = any([]*schema.Message{ + schema.UserMessage("hello"), + nil, + schema.UserMessage("world"), + }).([]M) + case *schema.AgenticMessage: + msgs = any([]*schema.AgenticMessage{ + schema.UserAgenticMessage("hello"), + nil, + schema.UserAgenticMessage("world"), + }).([]M) + } + + assert.NotPanics(t, func() { + copied := copyMessagesGeneric(msgs) + assert.Len(t, copied, 3) + assert.Nil(t, copied[1], "nil element should be preserved as nil") + }) +} diff --git a/adk/middlewares/reduction/reduction_test.go b/adk/middlewares/reduction/reduction_test.go new file mode 100644 index 0000000..dee9bb4 --- /dev/null +++ b/adk/middlewares/reduction/reduction_test.go @@ -0,0 +1,3566 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package reduction + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/schema" +) + +func testTruncOffloadPath(root string) func(context.Context, *ToolDetail) (string, error) { + return func(_ context.Context, detail *ToolDetail) (string, error) { + callID := "generated" + if detail != nil && detail.ToolContext != nil && detail.ToolContext.CallID != "" { + callID = detail.ToolContext.CallID + } + return fmt.Sprintf("%s/trunc/%s", root, callID), nil + } +} + +func testClearOffloadPath(root string) func(context.Context, *ToolDetail) (string, error) { + return func(_ context.Context, detail *ToolDetail) (string, error) { + callID := "generated" + if detail != nil && detail.ToolContext != nil && detail.ToolContext.CallID != "" { + callID = detail.ToolContext.CallID + } + return fmt.Sprintf("%s/clear/%s", root, callID), nil + } +} + +func stringStream(chunks ...string) *schema.StreamReader[string] { + sr, sw := schema.Pipe[string](len(chunks)) + go func() { + defer sw.Close() + for _, chunk := range chunks { + sw.Send(chunk, nil) + } + }() + return sr +} + +func stringStreamWithError(chunks []string, err error) *schema.StreamReader[string] { + sr, sw := schema.Pipe[string](len(chunks) + 1) + go func() { + defer sw.Close() + for _, chunk := range chunks { + sw.Send(chunk, nil) + } + var zero string + sw.Send(zero, err) + }() + return sr +} + +func toolResultStream(chunks ...*schema.ToolResult) *schema.StreamReader[*schema.ToolResult] { + sr, sw := schema.Pipe[*schema.ToolResult](len(chunks)) + go func() { + defer sw.Close() + for _, chunk := range chunks { + sw.Send(chunk, nil) + } + }() + return sr +} + +func toolResultStreamWithError(chunks []*schema.ToolResult, err error) *schema.StreamReader[*schema.ToolResult] { + sr, sw := schema.Pipe[*schema.ToolResult](len(chunks) + 1) + go func() { + defer sw.Close() + for _, chunk := range chunks { + sw.Send(chunk, nil) + } + var zero *schema.ToolResult + sw.Send(zero, err) + }() + return sr +} + +func sendStreamChunk[T any](t *testing.T, sw *schema.StreamWriter[T], chunk T) bool { + t.Helper() + + done := make(chan bool, 1) + go func() { + done <- sw.Send(chunk, nil) + }() + + select { + case closed := <-done: + return closed + case <-time.After(2 * time.Second): + t.Fatal("timed out sending stream chunk") + return true + } +} + +type writeFailBackend struct { + filesystem.Backend + err error +} + +func (b *writeFailBackend) Write(context.Context, *filesystem.WriteRequest) error { + return b.err +} + +func TestReductionMiddlewareTrunc(t *testing.T) { + ctx := context.Background() + it := mockInvokableTool() + st := mockStreamableTool() + + t.Run("test invokable max length trunc", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "mock_invokable_tool", + CallID: "12345", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "mock_invokable_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: defaultTruncHandler(testTruncOffloadPath("/tmp"), 70, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + exp := "\nOutput too large (199). Full output saved to: /tmp/trunc/12345\nPreview (first 35):\nhello worldhello worldhello worldhe\n\nPreview (last 35):\nldhello worldhello worldhello world\n\nUse read_file to view\n" + + edp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, tCtx) + assert.NoError(t, err) + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + assert.Equal(t, exp, resp) + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/12345"}) + assert.NoError(t, err) + expOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world +hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world` + assert.Equal(t, expOrigContent, content.Content) + }) + + t.Run("test default truncation notice includes configured read tool name", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "mock_invokable_tool", + CallID: "12345", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + MaxLengthForTrunc: 70, + ReadFileToolName: "cat_file", + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + edp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, tCtx) + assert.NoError(t, err) + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + assert.Contains(t, resp, "Use cat_file to view") + }) + + t.Run("test streamable line and max length trunc", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "mock_streamable_tool", + CallID: "54321", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + ToolConfig: map[string]*ToolReductionConfig{ + "mock_streamable_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: defaultTruncHandler(testTruncOffloadPath("/tmp"), 70, "read_file"), + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + exp := "\nOutput too large (199). Full output saved to: /tmp/trunc/54321\nPreview (first 35):\nhello worldhello worldhello worldhe\n\nPreview (last 35):\nldhello worldhello worldhello world\n\nUse read_file to view\n" + + edp, err := mw.WrapStreamableToolCall(ctx, st.StreamableRun, tCtx) + assert.NoError(t, err) + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + s, err := resp.Recv() + assert.NoError(t, err) + resp.Close() + assert.Equal(t, exp, s) + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/54321"}) + assert.NoError(t, err) + expOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world +hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world` + assert.Equal(t, expOrigContent, content.Content) + }) + + t.Run("test streamable default truncation is nonblocking before source EOF", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "mock_streamable_tool", + CallID: "stream_nonblocking", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + MaxLengthForTrunc: 30, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + releaseOverflow := make(chan struct{}) + releaseEOF := make(chan struct{}) + endpoint := func(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) { + sr, sw := schema.Pipe[string](0) + go func() { + sw.Send("first chunk", nil) + <-releaseOverflow + sw.Send(strings.Repeat("x", 80), nil) + <-releaseEOF + sw.Close() + }() + return sr, nil + } + + edp, err := mw.WrapStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + + respCh := make(chan *schema.StreamReader[string], 1) + errCh := make(chan error, 1) + go func() { + resp, runErr := edp(ctx, `{"value":"asd"}`) + if runErr != nil { + errCh <- runErr + return + } + respCh <- resp + }() + + var resp *schema.StreamReader[string] + select { + case runErr := <-errCh: + close(releaseEOF) + assert.NoError(t, runErr) + return + case resp = <-respCh: + case <-time.After(2 * time.Second): + close(releaseEOF) + assert.Fail(t, "streamable truncation wrapper blocked until source EOF") + return + } + defer resp.Close() + + first, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "first chunk", first) + + close(releaseOverflow) + noticeCh := make(chan string, 1) + noticeErrCh := make(chan error, 1) + go func() { + chunk, recvErr := resp.Recv() + if recvErr != nil { + noticeErrCh <- recvErr + return + } + noticeCh <- chunk + }() + + select { + case got := <-noticeCh: + close(releaseEOF) + assert.Failf(t, "streamable truncation notice returned before source EOF", "got %q", got) + return + case recvErr := <-noticeErrCh: + close(releaseEOF) + assert.NoError(t, recvErr) + return + case <-time.After(100 * time.Millisecond): + } + + close(releaseEOF) + var remaining string + select { + case remaining = <-noticeCh: + case recvErr := <-noticeErrCh: + assert.NoError(t, recvErr) + return + case <-time.After(2 * time.Second): + assert.Fail(t, "streamable truncation notice blocked after source EOF") + return + } + assert.Contains(t, remaining, "Output truncated after") + assert.NotContains(t, remaining, "Preview (first") + assert.NotContains(t, remaining, "Preview (last") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/stream_nonblocking"}) + assert.NoError(t, err) + assert.Equal(t, "first chunk"+strings.Repeat("x", 80), content.Content) + }) + + t.Run("test enhanced streamable default truncation is nonblocking before source EOF", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "mock_enhanced_streamable_tool", + CallID: "enhanced_stream_nonblocking", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + MaxLengthForTrunc: 30, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + releaseOverflow := make(chan struct{}) + releaseEOF := make(chan struct{}) + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr, sw := schema.Pipe[*schema.ToolResult](0) + go func() { + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "first chunk"}}}, nil) + <-releaseOverflow + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: strings.Repeat("x", 80)}}}, nil) + <-releaseEOF + sw.Close() + }() + return sr, nil + } + + edp, err := mw.WrapEnhancedStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + + respCh := make(chan *schema.StreamReader[*schema.ToolResult], 1) + errCh := make(chan error, 1) + go func() { + resp, runErr := edp(ctx, &schema.ToolArgument{Text: `{"value":"asd"}`}) + if runErr != nil { + errCh <- runErr + return + } + respCh <- resp + }() + + var resp *schema.StreamReader[*schema.ToolResult] + select { + case runErr := <-errCh: + close(releaseEOF) + assert.NoError(t, runErr) + return + case resp = <-respCh: + case <-time.After(2 * time.Second): + close(releaseEOF) + assert.Fail(t, "enhanced streamable truncation wrapper blocked until source EOF") + return + } + defer resp.Close() + + first, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "first chunk", first.Parts[0].Text) + + close(releaseOverflow) + noticeCh := make(chan string, 1) + noticeErrCh := make(chan error, 1) + go func() { + chunk, recvErr := resp.Recv() + if recvErr != nil { + noticeErrCh <- recvErr + return + } + noticeCh <- stringifyToolOutputParts(chunk.Parts) + }() + + select { + case got := <-noticeCh: + close(releaseEOF) + assert.Failf(t, "enhanced streamable truncation notice returned before source EOF", "got %q", got) + return + case recvErr := <-noticeErrCh: + close(releaseEOF) + assert.NoError(t, recvErr) + return + case <-time.After(100 * time.Millisecond): + } + + close(releaseEOF) + var remaining string + select { + case remaining = <-noticeCh: + case recvErr := <-noticeErrCh: + assert.NoError(t, recvErr) + return + case <-time.After(2 * time.Second): + assert.Fail(t, "enhanced streamable truncation notice blocked after source EOF") + return + } + assert.Contains(t, remaining, "Output truncated after") + assert.NotContains(t, remaining, "Preview (first") + assert.NotContains(t, remaining, "Preview (last") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/enhanced_stream_nonblocking"}) + assert.NoError(t, err) + assert.Contains(t, content.Content, `"text": "first chunk"`) + assert.Contains(t, content.Content, fmt.Sprintf(`"text": "%s"`, strings.Repeat("x", 80))) + }) + + t.Run("test default streamable truncation terminal cases", func(t *testing.T) { + tCtx := &adk.ToolContext{Name: "mock_streamable_tool", CallID: "stream_branches"} + backend := filesystem.NewInMemoryBackend() + mw := &typedToolReductionMiddleware[*schema.Message]{ + config: &TypedConfig[*schema.Message]{ + MaxLengthForTrunc: 10, + GenTruncOffloadFilePath: testTruncOffloadPath("/tmp"), + }, + } + + t.Run("short output forwards and does not truncate", func(t *testing.T) { + cfg := &ToolReductionConfig{} + resp := mw.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStream("short")) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "short", got) + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("source error is forwarded", func(t *testing.T) { + wantErr := errors.New("source stream error") + cfg := &ToolReductionConfig{} + resp := mw.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStreamWithError([]string{"ok"}, wantErr)) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "ok", got) + _, err = resp.Recv() + assert.Equal(t, wantErr, err) + }) + + t.Run("path generation error is rendered in truncation notice", func(t *testing.T) { + wantErr := errors.New("path generation error") + mwWithPathErr := &typedToolReductionMiddleware[*schema.Message]{ + config: &TypedConfig[*schema.Message]{ + MaxLengthForTrunc: 10, + GenTruncOffloadFilePath: func(context.Context, *ToolDetail) (string, error) { + return "", wantErr + }, + }, + } + cfg := &ToolReductionConfig{Backend: backend} + resp := mwWithPathErr.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStream(strings.Repeat("x", 20))) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Contains(t, got, strings.Repeat("x", 10)) + assert.Contains(t, got, "Output truncated after 10 bytes were streamed") + assert.Contains(t, got, "Failed to save full output: path generation error.") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("missing backend error is rendered in truncation notice", func(t *testing.T) { + cfg := &ToolReductionConfig{} + resp := mw.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStream(strings.Repeat("x", 20))) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Contains(t, got, strings.Repeat("x", 10)) + assert.Contains(t, got, "Output truncated after 10 bytes were streamed") + assert.Contains(t, got, "Failed to save full output: truncation: no backend for offload.") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("truncates on source EOF and writes full output", func(t *testing.T) { + cfg := &ToolReductionConfig{Backend: backend} + resp := mw.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStream(strings.Repeat("x", 20), "tail")) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Contains(t, got, strings.Repeat("x", 10)) + assert.Contains(t, got, "Output truncated after 10 bytes were streamed") + assert.Contains(t, got, "Full output saved to: /tmp/trunc/stream_branches.") + assert.NotContains(t, got, "Preview (first") + assert.NotContains(t, got, "Preview (last") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/stream_branches"}) + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 20)+"tail", content.Content) + }) + + t.Run("truncation notice includes configured read tool name", func(t *testing.T) { + mwWithReadTool := &typedToolReductionMiddleware[*schema.Message]{ + config: &TypedConfig[*schema.Message]{ + MaxLengthForTrunc: 10, + GenTruncOffloadFilePath: testTruncOffloadPath("/tmp"), + ReadFileToolName: "cat_file", + }, + } + cfg := &ToolReductionConfig{Backend: backend} + resp := mwWithReadTool.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStream(strings.Repeat("x", 20))) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Contains(t, got, "Use cat_file to view") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("write error is rendered in truncation notice after source EOF", func(t *testing.T) { + wantErr := errors.New("stream write error") + cfg := &ToolReductionConfig{ + Backend: &writeFailBackend{ + Backend: filesystem.NewInMemoryBackend(), + err: wantErr, + }, + } + resp := mw.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStream(strings.Repeat("x", 20))) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Contains(t, got, strings.Repeat("x", 10)) + assert.Contains(t, got, "Output truncated after 10 bytes were streamed") + assert.Contains(t, got, "Failed to save full output: stream write error.") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("source error after truncation replays swallowed chunks before error", func(t *testing.T) { + wantErr := errors.New("stream source error after truncation") + cfg := &ToolReductionConfig{Backend: backend} + resp := mw.wrapDefaultStreamableTruncation(ctx, cfg, tCtx, `{}`, stringStreamWithError( + []string{"abc", "defghijklmnop", "tail"}, + wantErr, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "abc", got) + + replayed, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "defghijklmnoptail", replayed) + _, err = resp.Recv() + assert.Equal(t, wantErr, err) + }) + }) + + t.Run("test default enhanced streamable truncation terminal cases", func(t *testing.T) { + tCtx := &adk.ToolContext{Name: "mock_enhanced_streamable_tool", CallID: "enhanced_stream_branches"} + toolArg := &schema.ToolArgument{Text: `{}`} + backend := filesystem.NewInMemoryBackend() + mw := &typedToolReductionMiddleware[*schema.Message]{ + config: &TypedConfig[*schema.Message]{ + MaxLengthForTrunc: 10, + GenTruncOffloadFilePath: testTruncOffloadPath("/tmp"), + }, + } + longResult := &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: strings.Repeat("x", 20)}}} + + t.Run("short output forwards and does not truncate", func(t *testing.T) { + cfg := &ToolReductionConfig{} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream( + &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "short"}}}, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "short", got.Parts[0].Text) + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("source error is forwarded", func(t *testing.T) { + wantErr := errors.New("enhanced source stream error") + cfg := &ToolReductionConfig{} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStreamWithError( + []*schema.ToolResult{{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "ok"}}}}, + wantErr, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "ok", got.Parts[0].Text) + _, err = resp.Recv() + assert.Equal(t, wantErr, err) + }) + + t.Run("conflicting media parts do not prevent truncation notice", func(t *testing.T) { + audioA := "YXVkaW8x" + audioB := "YXVkaW8y" + cfg := &ToolReductionConfig{Backend: backend} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream( + &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: strings.Repeat("x", 20)}, + {Type: schema.ToolPartTypeAudio, Audio: &schema.ToolOutputAudio{MessagePartCommon: schema.MessagePartCommon{Base64Data: &audioA}}}, + }}, + &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeAudio, Audio: &schema.ToolOutputAudio{MessagePartCommon: schema.MessagePartCommon{Base64Data: &audioB}}}, + }}, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 10), got.Parts[0].Text) + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Output truncated after 10 bytes were streamed") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("conflicting media parts after forwarded chunk do not prevent truncation notice", func(t *testing.T) { + audioA := "YXVkaW8x" + audioB := "YXVkaW8y" + cfg := &ToolReductionConfig{Backend: backend} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream( + &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "ok"}, + {Type: schema.ToolPartTypeAudio, Audio: &schema.ToolOutputAudio{MessagePartCommon: schema.MessagePartCommon{Base64Data: &audioA}}}, + }}, + &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: strings.Repeat("x", 20)}, + {Type: schema.ToolPartTypeAudio, Audio: &schema.ToolOutputAudio{MessagePartCommon: schema.MessagePartCommon{Base64Data: &audioB}}}, + }}, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "ok", got.Parts[0].Text) + notice, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 8), notice.Parts[0].Text) + assert.Contains(t, stringifyToolOutputParts(notice.Parts), "Output truncated after 10 bytes were streamed") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("path generation error is rendered in truncation notice", func(t *testing.T) { + wantErr := errors.New("enhanced path generation error") + mwWithPathErr := &typedToolReductionMiddleware[*schema.Message]{ + config: &TypedConfig[*schema.Message]{ + MaxLengthForTrunc: 10, + GenTruncOffloadFilePath: func(context.Context, *ToolDetail) (string, error) { + return "", wantErr + }, + }, + } + cfg := &ToolReductionConfig{Backend: backend} + resp := mwWithPathErr.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream(longResult)) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 10), got.Parts[0].Text) + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Output truncated after 10 bytes were streamed") + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Failed to save full output: enhanced path generation error.") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("missing backend error is rendered in truncation notice", func(t *testing.T) { + cfg := &ToolReductionConfig{} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream(longResult)) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 10), got.Parts[0].Text) + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Output truncated after 10 bytes were streamed") + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Failed to save full output: truncation: no backend for offload.") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("truncates on source EOF and writes full output", func(t *testing.T) { + cfg := &ToolReductionConfig{Backend: backend} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream( + longResult, + &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "tail"}}}, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 10), got.Parts[0].Text) + gotText := stringifyToolOutputParts(got.Parts) + assert.Contains(t, gotText, "Output truncated after 10 bytes were streamed") + assert.Contains(t, gotText, "Full output saved to: /tmp/trunc/enhanced_stream_branches.") + assert.NotContains(t, gotText, "Preview (first") + assert.NotContains(t, gotText, "Preview (last") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/enhanced_stream_branches"}) + assert.NoError(t, err) + assert.Contains(t, content.Content, fmt.Sprintf(`"text": "%s"`, strings.Repeat("x", 20))) + assert.Contains(t, content.Content, `"text": "tail"`) + }) + + t.Run("truncation notice includes configured read tool name", func(t *testing.T) { + mwWithReadTool := &typedToolReductionMiddleware[*schema.Message]{ + config: &TypedConfig[*schema.Message]{ + MaxLengthForTrunc: 10, + GenTruncOffloadFilePath: testTruncOffloadPath("/tmp"), + ReadFileToolName: "cat_file", + }, + } + cfg := &ToolReductionConfig{Backend: backend} + resp := mwWithReadTool.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream(longResult)) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Use cat_file to view") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("write error is rendered in truncation notice after source EOF", func(t *testing.T) { + wantErr := errors.New("enhanced write error") + cfg := &ToolReductionConfig{ + Backend: &writeFailBackend{ + Backend: filesystem.NewInMemoryBackend(), + err: wantErr, + }, + } + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStream(longResult)) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, strings.Repeat("x", 10), got.Parts[0].Text) + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Output truncated after 10 bytes were streamed") + assert.Contains(t, stringifyToolOutputParts(got.Parts), "Failed to save full output: enhanced write error.") + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + }) + + t.Run("source error after truncation replays swallowed chunks before error", func(t *testing.T) { + wantErr := errors.New("enhanced source stream error after truncation") + cfg := &ToolReductionConfig{Backend: backend} + resp := mw.wrapDefaultEnhancedStreamableTruncation(ctx, cfg, tCtx, toolArg, toolResultStreamWithError( + []*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "abc"}}}, + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "defghijklmnop"}}}, + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "tail"}}}, + }, + wantErr, + )) + defer resp.Close() + + got, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "abc", got.Parts[0].Text) + + replayed, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "defghijklmnop", replayed.Parts[0].Text) + assert.Equal(t, "tail", replayed.Parts[1].Text) + _, err = resp.Recv() + assert.Equal(t, wantErr, err) + }) + }) + + t.Run("test streamable line and bypass error", func(t *testing.T) { + stWithErr := mockStreamableToolWithError() + tCtx := &adk.ToolContext{ + Name: "mock_streamable_tool", + CallID: "54321", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + ToolConfig: map[string]*ToolReductionConfig{ + "mock_streamable_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: defaultTruncHandler(testTruncOffloadPath("/tmp"), 70, "read_file"), + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + edp, err := mw.WrapStreamableToolCall(ctx, stWithErr.StreamableRun, tCtx) + assert.NoError(t, err) + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + cnt := 0 + gotError := false + for { + _, err := resp.Recv() + if err != nil { + if err == io.EOF { + break + } + assert.Equal(t, fmt.Errorf("mock error"), err) + gotError = true + break + } + cnt++ + } + assert.True(t, gotError) + assert.Equal(t, 10, cnt) + }) + + t.Run("test TruncExcludeTools with invokable tool", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "important_tool", + CallID: "12345", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + TruncExcludeTools: []string{"important_tool"}, + MaxLengthForTrunc: 70, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + edp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, tCtx) + assert.NoError(t, err) + + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + + expOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world +hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world` + assert.Equal(t, expOrigContent, resp) + + _, err = backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/12345"}) + assert.Error(t, err) + }) + + t.Run("test TruncExcludeTools with streamable tool", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "important_stream_tool", + CallID: "54321", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + TruncExcludeTools: []string{"important_stream_tool"}, + MaxLengthForTrunc: 70, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + edp, err := mw.WrapStreamableToolCall(ctx, st.StreamableRun, tCtx) + assert.NoError(t, err) + + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + + var fullOutput strings.Builder + for { + chunk, recvErr := resp.Recv() + if recvErr != nil { + if recvErr == io.EOF { + break + } + assert.Fail(t, "unexpected error", recvErr) + } + fullOutput.WriteString(chunk) + } + resp.Close() + + expOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world +hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world` + assert.Equal(t, expOrigContent, fullOutput.String()) + + _, err = backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/54321"}) + assert.Error(t, err) + }) + + t.Run("test mixed tools with TruncExcludeTools", func(t *testing.T) { + excludedToolCtx := &adk.ToolContext{ + Name: "excluded_tool", + CallID: "excluded_123", + } + normalToolCtx := &adk.ToolContext{ + Name: "normal_tool", + CallID: "normal_456", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + TruncExcludeTools: []string{"excluded_tool"}, + MaxLengthForTrunc: 70, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + excludedEdp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, excludedToolCtx) + assert.NoError(t, err) + excludedResp, err := excludedEdp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + expOrigContent := `hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world +hello worldhello worldhello worldhello worldhello worldhello worldhello worldhello world` + assert.Equal(t, expOrigContent, excludedResp) + + normalEdp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, normalToolCtx) + assert.NoError(t, err) + normalResp, err := normalEdp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + assert.NotEqual(t, expOrigContent, normalResp) + assert.Contains(t, normalResp, "persisted-output") + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/normal_456"}) + assert.NoError(t, err) + assert.Equal(t, expOrigContent, content.Content) + }) + + t.Run("test GenTruncOffloadFilePath used by default trunc handler", func(t *testing.T) { + tCtx := &adk.ToolContext{ + Name: "mock_invokable_tool", + CallID: "custom_123", + } + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + MaxLengthForTrunc: 70, + RootDir: "/tmp/ignored", + GenTruncOffloadFilePath: func(_ context.Context, detail *ToolDetail) (string, error) { + assert.Equal(t, "custom_123", detail.ToolContext.CallID) + return "/custom/trunc/custom_123", nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + edp, err := mw.WrapInvokableToolCall(ctx, it.InvokableRun, tCtx) + assert.NoError(t, err) + + resp, err := edp(ctx, `{"value":"asd"}`) + assert.NoError(t, err) + assert.Contains(t, resp, "/custom/trunc/custom_123") + + _, err = backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/custom/trunc/custom_123"}) + assert.NoError(t, err) + }) +} + +func TestReductionMiddlewareClear(t *testing.T) { + ctx := context.Background() + it := mockInvokableTool() + st := mockStreamableTool() + tools := []tool.BaseTool{it, st} + var toolsInfo []*schema.ToolInfo + for _, bt := range tools { + ti, _ := bt.Info(ctx) + toolsInfo = append(toolsInfo, ti) + } + type OffloadContent struct { + Arguments map[string]string `json:"arguments"` + Result string `json:"result"` + } + + t.Run("test default clear", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 0, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_987654321"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_123456789"), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[2].ToolCalls) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[4].ToolCalls) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_987654321\nUse read_file to view", s.Messages[3].Content) + fileContent, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_987654321", + }) + assert.NoError(t, err) + fileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), "1\t") + assert.Equal(t, "Sunny", fileContentStr) + }) + + t.Run("test GenClearOffloadFilePath used by default clear handler", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + SkipTruncation: true, + TokenCounter: func(context.Context, []adk.Message, []*schema.ToolInfo) (int64, error) { return 100, nil }, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 1, + RootDir: "/tmp/ignored", + GenClearOffloadFilePath: func(_ context.Context, detail *ToolDetail) (string, error) { + return "/custom/clear/" + detail.ToolContext.CallID, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("trigger clear"), + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "call_a", Type: "function", Function: schema.FunctionCall{Name: "get_weather", Arguments: `{}`}}, + }), + schema.ToolMessage("Sunny", "call_a"), + schema.AssistantMessage("", []schema.ToolCall{ + {ID: "call_b", Type: "function", Function: schema.FunctionCall{Name: "get_weather", Arguments: `{}`}}, + }), + schema.ToolMessage("Sunny", "call_b"), + }, + }, &adk.ModelContext{Tools: toolsInfo}) + assert.NoError(t, err) + + assert.Equal(t, "Tool result saved to: /custom/clear/call_a\nUse read_file to view", s.Messages[3].Content) + + fileContent, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/custom/clear/call_a", + }) + assert.NoError(t, err) + fileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), "1\t") + assert.Equal(t, "Sunny", fileContentStr) + }) + + t.Run("test default clear without offloading", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 0, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath(""), false, ""), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_987654321"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_123456789"), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[2].ToolCalls) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[4].ToolCalls) + assert.Equal(t, "[Old tool result content cleared]", s.Messages[3].Content) + }) + + t.Run("test clear with reordered tool results", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 0, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("get weather for two cities"), + // Assistant message with two tool calls + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "two_call_one", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London"}`}, + }, + { + ID: "two_call_two", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "Paris"}`}, + }, + }), + // Intentionally put tool result two BEFORE tool result one + schema.ToolMessage("Cloudy in Paris", "two_call_two"), + schema.ToolMessage("Sunny in London", "two_call_one"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + + // Verify tool call arguments are preserved (not offloaded since offload is true) + assert.Equal(t, `{"location": "London"}`, s.Messages[2].ToolCalls[0].Function.Arguments) + assert.Equal(t, `{"location": "Paris"}`, s.Messages[2].ToolCalls[1].Function.Arguments) + + // Verify tool result two is correctly matched and offloaded (it's at index 3) + assert.Equal(t, "Tool result saved to: /tmp/clear/two_call_two\nUse read_file to view", s.Messages[3].Content) + fileContent2, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/two_call_two", + }) + assert.NoError(t, err) + fileContentStr2 := strings.TrimPrefix(strings.TrimSpace(fileContent2.Content), "1\t") + assert.Equal(t, "Cloudy in Paris", fileContentStr2) + + // Verify tool result one is correctly matched and offloaded (it's at index 4) + assert.Equal(t, "Tool result saved to: /tmp/clear/two_call_one\nUse read_file to view", s.Messages[4].Content) + fileContent1, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/two_call_one", + }) + assert.NoError(t, err) + fileContentStr1 := strings.TrimPrefix(strings.TrimSpace(fileContent1.Content), "1\t") + assert.Equal(t, "Sunny in London", fileContentStr1) + }) + + t.Run("test clear", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + handler := func(ctx context.Context, detail *ToolDetail) (*ClearResult, error) { + arguments := make(map[string]string) + if err := json.Unmarshal([]byte(detail.ToolArgument.Text), &arguments); err != nil { + return nil, err + } + offloadContent := &OffloadContent{ + Arguments: arguments, + Result: detail.ToolResult.Parts[0].Text, + } + replacedArguments := make(map[string]string, len(arguments)) + filePath := fmt.Sprintf("/tmp/%s", detail.ToolContext.CallID) + for k := range arguments { + replacedArguments[k] = "argument offloaded" + } + return &ClearResult{ + ToolArgument: &schema.ToolArgument{Text: toJson(replacedArguments)}, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "result offloaded, retrieve it from " + filePath}, + }, + }, + NeedClear: true, + NeedOffload: true, + OffloadFilePath: filePath, + OffloadContent: toJson(offloadContent), + }, nil + } + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 1, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: handler, + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_987654321"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_123456789"), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location":"argument offloaded","unit":"argument offloaded"}`}, + }, + }, s.Messages[2].ToolCalls) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[4].ToolCalls) + assert.Equal(t, "result offloaded, retrieve it from /tmp/call_987654321", s.Messages[3].Content) + fileContent, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/call_987654321", + }) + assert.NoError(t, err) + fileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), "1\t") + oc := &OffloadContent{} + err = json.Unmarshal([]byte(fileContentStr), oc) + assert.NoError(t, err) + assert.Equal(t, &OffloadContent{ + Arguments: map[string]string{ + "location": "London, UK", + "unit": "c", + }, + Result: "Sunny", + }, oc) + }) + + t.Run("test skip handled ones", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 0, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + msgs := []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_987654321"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_123456789"), + } + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo}) + assert.NoError(t, err) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[2].ToolCalls) + assert.NotNil(t, msgs[2].Extra[msgClearedFlag]) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[4].ToolCalls) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_987654321\nUse read_file to view", s.Messages[3].Content) + fileContent, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_987654321", + }) + assert.NoError(t, err) + fileContentStr := strings.TrimPrefix(strings.TrimSpace(fileContent.Content), "1\t") + assert.Equal(t, "Sunny", fileContentStr) + + msgs = append(msgs, []*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_8877665544", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_8877665544"), + }...) + _, s, err = mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo}) + assert.NoError(t, err) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[2].ToolCalls) + assert.NotNil(t, msgs[2].Extra[msgClearedFlag]) + assert.Equal(t, []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }, s.Messages[4].ToolCalls) + assert.NotNil(t, msgs[4].Extra[msgClearedFlag]) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_987654321\nUse read_file to view", s.Messages[3].Content) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_123456789\nUse read_file to view", s.Messages[5].Content) + }) + + t.Run("test ClearExcludeTools", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: defaultTokenCounter, + MaxTokensForClear: 20, + ClearRetentionSuffixLimit: 0, + ClearExcludeTools: []string{"get_important_data"}, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + "get_important_data": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + msgs := []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_987654321"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_important_data", Arguments: `{"id": "123"}`}, + }, + }), + schema.ToolMessage("Important Data Content", "call_123456789"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_999", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_999"), + } + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo}) + assert.NoError(t, err) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_987654321\nUse read_file to view", s.Messages[3].Content) + assert.Equal(t, "Important Data Content", s.Messages[5].Content) + b, err := backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_987654321", + }) + assert.NoError(t, err) + assert.Equal(t, "Sunny", b.Content) + _, err = backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_123456789", + }) + assert.Error(t, err) + assert.Equal(t, "file not found: /tmp/clear/call_123456789", err.Error()) + }) + + t.Run("test ClearAtLeastTokens - not enough tokens cleared", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: func(_ context.Context, msgs []adk.Message, _ []*schema.ToolInfo) (int64, error) { + var size int + for _, msg := range msgs { + size += len(msg.Content) + for _, tc := range msg.ToolCalls { + size += len(tc.Function.Name) + size += len(tc.Function.Arguments) + } + } + return int64(size), nil + }, + MaxTokensForClear: 50, + ClearRetentionSuffixLimit: 1, + ClearAtLeastTokens: 100, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + "get_important_data": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + msgs := []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", "call_123456789"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_important_data", Arguments: `{"id": "123"}`}, + }, + }), + schema.ToolMessage("Important Data Content, qweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqwe", "call_123456789"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_999", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_999"), + } + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo}) + assert.NoError(t, err) + assert.Equal(t, "Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", s.Messages[3].Content) + assert.Equal(t, "Important Data Content, qweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqwe", s.Messages[5].Content) + _, err = backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_987654321", + }) + assert.Error(t, err) + }) + + t.Run("test ClearAtLeastTokens - enough tokens cleared", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + SkipTruncation: true, + TokenCounter: func(_ context.Context, msgs []adk.Message, _ []*schema.ToolInfo) (int64, error) { + var size int + for _, msg := range msgs { + size += len(msg.Content) + for _, tc := range msg.ToolCalls { + size += len(tc.Function.Name) + size += len(tc.Function.Arguments) + } + } + return int64(size), nil + }, + MaxTokensForClear: 50, + ClearRetentionSuffixLimit: 1, + ClearAtLeastTokens: 10, + ToolConfig: map[string]*ToolReductionConfig{ + "get_weather": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + "get_important_data": { + Backend: backend, + SkipClear: false, + ClearHandler: defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file"), + }, + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + msgs := []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("If it's warmer than 20°C in London, set the thermostat to 20°C, otherwise set it to 18°C."), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_987654321", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny Sunny", "call_987654321"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123456789", + Type: "function", + Function: schema.FunctionCall{Name: "get_important_data", Arguments: `{"id": "123"}`}, + }, + }), + schema.ToolMessage("Important Data Content, qweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqweqwe", "call_123456789"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_999", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London, UK", "unit": "c"}`}, + }, + }), + schema.ToolMessage("Sunny", "call_999"), + } + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{Messages: msgs}, &adk.ModelContext{Tools: toolsInfo}) + assert.NoError(t, err) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_987654321\nUse read_file to view", s.Messages[3].Content) + _, err = backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_987654321", + }) + assert.NoError(t, err) + assert.Equal(t, "Tool result saved to: /tmp/clear/call_123456789\nUse read_file to view", s.Messages[5].Content) + _, err = backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: "/tmp/clear/call_987654321", + }) + assert.NoError(t, err) + }) +} + +func TestStreamTruncationHelpers(t *testing.T) { + t.Run("toolResultTextPrefix trims text and preserves leading non-text parts", func(t *testing.T) { + assert.Nil(t, toolResultTextPrefix(nil, 10)) + assert.Nil(t, toolResultTextPrefix(&schema.ToolResult{}, 10)) + assert.Nil(t, toolResultTextPrefix(&schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "abc"}, + }}, 0)) + + got := toolResultTextPrefix(&schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeImage}, + {Type: schema.ToolPartTypeText, Text: "abc"}, + {Type: schema.ToolPartTypeText, Text: "def"}, + }}, 5) + assert.NotNil(t, got) + assert.Equal(t, []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeImage}, + {Type: schema.ToolPartTypeText, Text: "abc"}, + {Type: schema.ToolPartTypeText, Text: "de"}, + }, got.Parts) + + got = toolResultTextPrefix(&schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "abc"}, + {Type: schema.ToolPartTypeText, Text: "def"}, + }}, 3) + assert.NotNil(t, got) + assert.Equal(t, []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "abc"}, + }, got.Parts) + }) + + t.Run("toolResultTextLength handles nil result", func(t *testing.T) { + assert.Equal(t, 0, toolResultTextLength(nil)) + }) + + t.Run("writeTruncOffload handles disabled missing and successful offload", func(t *testing.T) { + ctx := context.Background() + + err := writeTruncOffload(ctx, &ToolReductionConfig{}, &TruncResult{NeedOffload: false}) + assert.NoError(t, err) + + err = writeTruncOffload(ctx, &ToolReductionConfig{}, &TruncResult{NeedOffload: true}) + assert.EqualError(t, err, "truncation: no backend for offload") + + backend := filesystem.NewInMemoryBackend() + err = writeTruncOffload(ctx, &ToolReductionConfig{Backend: backend}, &TruncResult{ + NeedOffload: true, + OffloadFilePath: "/tmp/trunc/helper", + OffloadContent: "full output", + }) + assert.NoError(t, err) + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/helper"}) + assert.NoError(t, err) + assert.Equal(t, "full output", content.Content) + }) +} + +func TestGetJointToolResult(t *testing.T) { + t.Run("test with ToolResult", func(t *testing.T) { + detail := &ToolDetail{ + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "hello world"}}}, + } + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.True(t, needProcess) + assert.Len(t, toolOutputParts, 1) + assert.Equal(t, schema.ToolPartTypeText, toolOutputParts[0].Type) + assert.Equal(t, "hello world", toolOutputParts[0].Text) + }) + + t.Run("test with multiple ToolResult parts", func(t *testing.T) { + detail := &ToolDetail{ + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "hello "}, + {Type: schema.ToolPartTypeText, Text: "world"}, + }}, + } + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.True(t, needProcess) + assert.Len(t, toolOutputParts, 2) + assert.Equal(t, schema.ToolPartTypeText, toolOutputParts[0].Type) + assert.Equal(t, "hello ", toolOutputParts[0].Text) + assert.Equal(t, schema.ToolPartTypeText, toolOutputParts[1].Type) + assert.Equal(t, "world", toolOutputParts[1].Text) + }) + + t.Run("test with empty ToolResult parts", func(t *testing.T) { + detail := &ToolDetail{ + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{}}, + } + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.False(t, needProcess) + assert.Nil(t, toolOutputParts) + }) + + t.Run("test with ToolResult multimodal", func(t *testing.T) { + detail := &ToolDetail{ + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeImage, Text: "https://example.com/image.png"}}}, + } + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.True(t, needProcess) + assert.Len(t, toolOutputParts, 1) + assert.Equal(t, schema.ToolPartTypeImage, toolOutputParts[0].Type) + assert.Equal(t, "https://example.com/image.png", toolOutputParts[0].Text) + }) + + t.Run("test with StreamToolResult", func(t *testing.T) { + sr, sw := schema.Pipe[*schema.ToolResult](10) + go func() { + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "hello "}}}, nil) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "world"}}}, nil) + sw.Close() + }() + + detail := &ToolDetail{ + StreamToolResult: sr, + } + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.True(t, needProcess) + assert.Len(t, toolOutputParts, 1) + assert.Equal(t, schema.ToolPartTypeText, toolOutputParts[0].Type) + assert.Equal(t, "hello world", toolOutputParts[0].Text) + }) + + t.Run("test with StreamToolResult error", func(t *testing.T) { + sr, sw := schema.Pipe[*schema.ToolResult](10) + go func() { + sw.Send(nil, fmt.Errorf("stream error")) + sw.Close() + }() + + detail := &ToolDetail{ + StreamToolResult: sr, + } + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.False(t, needProcess) + assert.Nil(t, toolOutputParts) + }) + + t.Run("test with both ToolResult and StreamToolResult nil (should error)", func(t *testing.T) { + detail := &ToolDetail{} + toolOutputParts, needProcess, err := getJointToolResult(detail) + assert.Error(t, err) + assert.False(t, needProcess) + assert.Nil(t, toolOutputParts) + }) +} + +func TestDefaultTruncHandlerWithStreamToolResult(t *testing.T) { + ctx := context.Background() + + t.Run("test short stream result no trunc", func(t *testing.T) { + sr, sw := schema.Pipe[*schema.ToolResult](10) + go func() { + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "short text"}}}, nil) + sw.Close() + }() + + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: "test", + CallID: "call_id", + }, + ToolArgument: &schema.ToolArgument{Text: "{}"}, + StreamToolResult: sr, + } + + fn := defaultTruncHandler(testTruncOffloadPath("/tmp"), 100, "read_file") + result, err := fn(ctx, detail) + assert.NoError(t, err) + assert.False(t, result.NeedTrunc) + }) + + t.Run("test long stream result need trunc", func(t *testing.T) { + sr, sw := schema.Pipe[*schema.ToolResult](10) + go func() { + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: strings.Repeat("hello world", 20)}}}, nil) + sw.Close() + }() + + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: "test", + CallID: "call_id", + }, + ToolArgument: &schema.ToolArgument{Text: "{}"}, + StreamToolResult: sr, + } + + fn := defaultTruncHandler(testTruncOffloadPath("/tmp"), 100, "read_file") + result, err := fn(ctx, detail) + assert.NoError(t, err) + assert.True(t, result.NeedTrunc) + }) + + t.Run("test gen offload file path error", func(t *testing.T) { + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: "test", + CallID: "call_id", + }, + ToolArgument: &schema.ToolArgument{Text: "{}"}, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: strings.Repeat("x", 1000)}}, + }, + } + + fn := defaultTruncHandler(func(context.Context, *ToolDetail) (string, error) { + return "", fmt.Errorf("gen path error") + }, 10, "read_file") + _, err := fn(ctx, detail) + assert.Error(t, err) + assert.Contains(t, err.Error(), "gen path error") + }) +} + +func TestDefaultClearHandlerWithStreamToolResult(t *testing.T) { + ctx := context.Background() + + t.Run("test with stream result need offload", func(t *testing.T) { + sr, sw := schema.Pipe[*schema.ToolResult](10) + go func() { + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "streaming content"}}}, nil) + sw.Close() + }() + + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: "test", + CallID: "stream_call_id", + }, + ToolArgument: &schema.ToolArgument{Text: "{}"}, + StreamToolResult: sr, + } + + fn := defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file") + result, err := fn(ctx, detail) + assert.NoError(t, err) + assert.True(t, result.NeedClear) + assert.True(t, result.NeedOffload) + assert.Equal(t, "streaming content", result.OffloadContent) + }) + + t.Run("test gen offload file path error", func(t *testing.T) { + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: "test", + CallID: "call_id", + }, + ToolArgument: &schema.ToolArgument{Text: "{}"}, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "content"}}, + }, + } + fn := defaultClearHandler(func(context.Context, *ToolDetail) (string, error) { + return "", fmt.Errorf("gen path error") + }, true, "read_file") + _, err := fn(ctx, detail) + assert.Error(t, err) + assert.Contains(t, err.Error(), "gen path error") + }) +} + +func TestDefaultOffloadHandler(t *testing.T) { + ctx := context.Background() + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + Name: "mock_name", + CallID: "mock_call_id_12345", + }, + ToolArgument: &schema.ToolArgument{Text: "anything"}, + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "hello"}}}, + } + + fn := defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file") + info, err := fn(ctx, detail) + assert.NoError(t, err) + assert.Equal(t, &ClearResult{ + ToolArgument: &schema.ToolArgument{Text: "anything"}, + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{ + { + Type: schema.ToolPartTypeText, + Text: "Tool result saved to: /tmp/clear/mock_call_id_12345\nUse read_file to view", + }, + }}, + NeedClear: true, + NeedOffload: true, + OffloadFilePath: "/tmp/clear/mock_call_id_12345", + OffloadContent: "hello", + }, info) + +} + +func mockInvokableTool() tool.InvokableTool { + type ContentContainer struct { + Value string `json:"value"` + } + s1 := strings.Repeat("hello world", 10) + "\n" + s2 := strings.Repeat("hello world", 8) + s3 := s1 + s2 + t, _ := utils.InferTool("mock_invokable_tool", "test desc", func(ctx context.Context, input *ContentContainer) (output string, err error) { + return s3, nil + }) + return t +} + +func mockStreamableTool() tool.StreamableTool { + type ContentContainer struct { + Value string `json:"value"` + } + s1 := strings.Repeat("hello world", 10) + "\n" + s2 := strings.Repeat("hello world", 8) + s3 := s1 + s2 + t, _ := utils.InferStreamTool("mock_streamable_tool", "test desc", func(ctx context.Context, input ContentContainer) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](11) + for _, part := range splitStrings(s3, 10) { + sw.Send(part, nil) + } + sw.Close() + return sr, nil + }) + return t +} + +func mockStreamableToolWithError() tool.StreamableTool { + type ContentContainer struct { + Value string `json:"value"` + } + s1 := strings.Repeat("hello world", 10) + "\n" + s2 := strings.Repeat("hello world", 8) + s3 := s1 + s2 + t, _ := utils.InferStreamTool("mock_streamable_tool", "test desc", func(ctx context.Context, input ContentContainer) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](11) + for _, part := range splitStrings(s3, 10) { + sw.Send(part, nil) + } + sw.Send("", fmt.Errorf("mock error")) + sw.Close() + return sr, nil + }) + return t +} + +func splitStrings(s string, n int) []string { + if n <= 0 { + n = 1 + } + if n == 1 { + return []string{s} + } + if len(s) <= n { + parts := make([]string, n) + for i := 0; i < len(s); i++ { + parts[i] = string(s[i]) + } + return parts + } + baseLen := len(s) / n + extra := len(s) % n + parts := make([]string, 0, n) + start := 0 + for i := 0; i < n; i++ { + end := start + baseLen + if i < extra { + end++ + } + parts = append(parts, s[start:end]) + start = end + } + return parts +} + +func toJson(v any) string { + b, _ := json.Marshal(v) + return string(b) +} + +func TestToolResultFromMessage(t *testing.T) { + t.Run("test from content", func(t *testing.T) { + msg := schema.ToolMessage("test content", "call_123") + result, fromContent, err := toolResultFromMessage(msg) + assert.NoError(t, err) + assert.True(t, fromContent) + assert.NotNil(t, result) + assert.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "test content", result.Parts[0].Text) + }) + + t.Run("test from user input multi content", func(t *testing.T) { + msg := schema.ToolMessage("", "call_456") + msg.UserInputMultiContent = []schema.MessageInputPart{ + { + Type: schema.ChatMessagePartTypeText, + Text: "test text", + }, + } + result, fromContent, err := toolResultFromMessage(msg) + assert.NoError(t, err) + assert.False(t, fromContent) + assert.NotNil(t, result) + assert.Len(t, result.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "test text", result.Parts[0].Text) + }) + + t.Run("test invalid role", func(t *testing.T) { + msg := schema.UserMessage("test user message") + _, _, err := toolResultFromMessage(msg) + assert.Error(t, err) + assert.Contains(t, err.Error(), "message role") + }) +} + +func TestConvMessageInputPartToToolOutputPart(t *testing.T) { + t.Run("test text type", func(t *testing.T) { + part := schema.MessageInputPart{ + Type: schema.ChatMessagePartTypeText, + Text: "test text", + } + result, err := convMessageInputPartToToolOutputPart(part) + assert.NoError(t, err) + assert.Equal(t, schema.ToolPartTypeText, result.Type) + assert.Equal(t, "test text", result.Text) + }) + + t.Run("test image url type", func(t *testing.T) { + part := schema.MessageInputPart{ + Type: schema.ChatMessagePartTypeImageURL, + Image: &schema.MessageInputImage{}, + } + result, err := convMessageInputPartToToolOutputPart(part) + assert.NoError(t, err) + assert.Equal(t, schema.ToolPartTypeImage, result.Type) + assert.NotNil(t, result.Image) + }) + + t.Run("test audio url type", func(t *testing.T) { + part := schema.MessageInputPart{ + Type: schema.ChatMessagePartTypeAudioURL, + Audio: &schema.MessageInputAudio{}, + } + result, err := convMessageInputPartToToolOutputPart(part) + assert.NoError(t, err) + assert.Equal(t, schema.ToolPartTypeAudio, result.Type) + assert.NotNil(t, result.Audio) + }) + + t.Run("test video url type", func(t *testing.T) { + part := schema.MessageInputPart{ + Type: schema.ChatMessagePartTypeVideoURL, + Video: &schema.MessageInputVideo{}, + } + result, err := convMessageInputPartToToolOutputPart(part) + assert.NoError(t, err) + assert.Equal(t, schema.ToolPartTypeVideo, result.Type) + assert.NotNil(t, result.Video) + }) + + t.Run("test file url type", func(t *testing.T) { + part := schema.MessageInputPart{ + Type: schema.ChatMessagePartTypeFileURL, + File: &schema.MessageInputFile{}, + } + result, err := convMessageInputPartToToolOutputPart(part) + assert.NoError(t, err) + assert.Equal(t, schema.ToolPartTypeFile, result.Type) + assert.NotNil(t, result.File) + }) + + t.Run("test unknown type", func(t *testing.T) { + part := schema.MessageInputPart{ + Type: "unknown_type", + } + _, err := convMessageInputPartToToolOutputPart(part) + assert.Error(t, err) + assert.Contains(t, err.Error(), "unknown msg part type") + }) +} + +func TestGetSetMsgOffloadedFlag(t *testing.T) { + t.Run("test get offloaded flag - not set", func(t *testing.T) { + msg := schema.UserMessage("test") + assert.False(t, getMsgClearedFlag(msg)) + }) + + t.Run("test get offloaded flag - set", func(t *testing.T) { + msg := schema.UserMessage("test") + setMsgClearedFlag(msg) + assert.True(t, getMsgClearedFlag(msg)) + }) + + t.Run("test set offloaded flag - nil extra", func(t *testing.T) { + msg := schema.UserMessage("test") + setMsgClearedFlag(msg) + assert.True(t, getMsgClearedFlag(msg)) + }) + + t.Run("test set offloaded flag - existing extra", func(t *testing.T) { + msg := schema.UserMessage("test") + msg.Extra = map[string]any{"existing": "value"} + setMsgClearedFlag(msg) + assert.True(t, getMsgClearedFlag(msg)) + assert.Equal(t, "value", msg.Extra["existing"]) + }) +} + +func TestNewErrors(t *testing.T) { + ctx := context.Background() + + t.Run("test nil config", func(t *testing.T) { + _, err := New(ctx, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config must not be nil") + }) + + t.Run("test no backend when not skipping truncation", func(t *testing.T) { + config := &Config{ + Backend: nil, + SkipTruncation: false, + } + _, err := New(ctx, config) + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend must be set") + }) +} + +func TestGetToolConfig(t *testing.T) { + ctx := context.Background() + backend := filesystem.NewInMemoryBackend() + + t.Run("test no tool config", func(t *testing.T) { + config := &Config{ + Backend: backend, + SkipTruncation: true, + SkipClear: true, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + trmw, ok := mw.(*typedToolReductionMiddleware[*schema.Message]) + assert.True(t, ok) + + cfg := trmw.getToolConfig("non_existent_tool", sceneTruncation) + assert.NotNil(t, cfg) + }) + + t.Run("test with tool config", func(t *testing.T) { + config := &Config{ + Backend: backend, + SkipTruncation: true, + SkipClear: true, + ToolConfig: map[string]*ToolReductionConfig{ + "test_tool": { + SkipTruncation: true, + SkipClear: true, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + trmw, ok := mw.(*typedToolReductionMiddleware[*schema.Message]) + assert.True(t, ok) + + cfg := trmw.getToolConfig("test_tool", sceneTruncation) + assert.NotNil(t, cfg) + assert.True(t, cfg.SkipTruncation) + }) + + t.Run("test with tool config needing default handler", func(t *testing.T) { + config := &Config{ + Backend: backend, + SkipTruncation: false, + ToolConfig: map[string]*ToolReductionConfig{ + "test_tool": { + SkipTruncation: false, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + trmw, ok := mw.(*typedToolReductionMiddleware[*schema.Message]) + assert.True(t, ok) + + cfg := trmw.getToolConfig("test_tool", sceneTruncation) + assert.NotNil(t, cfg) + assert.NotNil(t, cfg.TruncHandler) + }) +} + +func TestCopyAndFillDefaults(t *testing.T) { + t.Run("test empty config", func(t *testing.T) { + cfg := &Config{} + result, err := cfg.copyAndFillDefaults() + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Equal(t, "/tmp", result.RootDir) + assert.Equal(t, "read_file", result.ReadFileToolName) + assert.Equal(t, 50000, result.MaxLengthForTrunc) + assert.Equal(t, 1, result.ClearRetentionSuffixLimit) + assert.NotNil(t, result.TokenCounter) + assert.NotNil(t, result.GenTruncOffloadFilePath) + assert.NotNil(t, result.GenClearOffloadFilePath) + + truncPath, err := result.GenTruncOffloadFilePath(context.Background(), &ToolDetail{ + ToolContext: &adk.ToolContext{CallID: "call_1"}, + }) + assert.NoError(t, err) + assert.Equal(t, "/tmp/trunc/call_1", truncPath) + + clearPath, err := result.GenClearOffloadFilePath(context.Background(), &ToolDetail{ + ToolContext: &adk.ToolContext{CallID: "call_1"}, + }) + assert.NoError(t, err) + assert.Equal(t, "/tmp/clear/call_1", clearPath) + }) + + t.Run("test default offload path generators create uuid when CallID is empty", func(t *testing.T) { + cfg := &Config{} + result, err := cfg.copyAndFillDefaults() + assert.NoError(t, err) + + truncPath, err := result.GenTruncOffloadFilePath(context.Background(), &ToolDetail{ + ToolContext: &adk.ToolContext{CallID: ""}, + }) + assert.NoError(t, err) + assert.True(t, strings.HasPrefix(truncPath, "/tmp/trunc/")) + assert.NotEqual(t, "/tmp/trunc/", truncPath) + + clearPath, err := result.GenClearOffloadFilePath(context.Background(), &ToolDetail{ + ToolContext: &adk.ToolContext{CallID: ""}, + }) + assert.NoError(t, err) + assert.True(t, strings.HasPrefix(clearPath, "/tmp/clear/")) + assert.NotEqual(t, "/tmp/clear/", clearPath) + }) + + t.Run("test with tool config", func(t *testing.T) { + cfg := &Config{ + ToolConfig: map[string]*ToolReductionConfig{ + "test_tool": { + SkipTruncation: true, + }, + }, + } + result, err := cfg.copyAndFillDefaults() + assert.NoError(t, err) + assert.NotNil(t, result.ToolConfig) + assert.True(t, result.ToolConfig["test_tool"].SkipTruncation) + }) + + t.Run("test custom offload path generators preserved", func(t *testing.T) { + cfg := &Config{ + RootDir: "/tmp/ignored", + GenTruncOffloadFilePath: func(_ context.Context, detail *ToolDetail) (string, error) { + return "/custom/trunc/" + detail.ToolContext.CallID, nil + }, + GenClearOffloadFilePath: func(_ context.Context, detail *ToolDetail) (string, error) { + return "/custom/clear/" + detail.ToolContext.CallID, nil + }, + } + result, err := cfg.copyAndFillDefaults() + assert.NoError(t, err) + + filePath, err := result.GenClearOffloadFilePath(context.Background(), &ToolDetail{ + ToolContext: &adk.ToolContext{CallID: "call_custom"}, + }) + assert.NoError(t, err) + assert.Equal(t, "/custom/clear/call_custom", filePath) + }) +} + +func TestDefaultTokenCounter(t *testing.T) { + ctx := context.Background() + + t.Run("test with nil messages", func(t *testing.T) { + msgs := []*schema.Message{nil} + tokens, err := defaultTokenCounter(ctx, msgs, nil) + assert.NoError(t, err) + assert.GreaterOrEqual(t, tokens, int64(0)) + }) + + t.Run("test with tool info", func(t *testing.T) { + toolInfo := &schema.ToolInfo{ + Name: "test_tool", + Desc: "test description", + } + tokens, err := defaultTokenCounter(ctx, nil, []*schema.ToolInfo{toolInfo}) + assert.NoError(t, err) + assert.GreaterOrEqual(t, tokens, int64(0)) + }) +} + +func TestDefaultClearHandler(t *testing.T) { + ctx := context.Background() + + t.Run("test empty parts", func(t *testing.T) { + handler := defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file") + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + CallID: "test_call", + }, + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{}}, + } + result, err := handler(ctx, detail) + assert.NoError(t, err) + assert.False(t, result.NeedClear) + }) + + t.Run("test multimodal content", func(t *testing.T) { + handler := defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file") + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{ + CallID: "test_call", + }, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeImage}}, + }, + } + result, err := handler(ctx, detail) + assert.NoError(t, err) + assert.Equal(t, detail.ToolResult, result.ToolResult) + }) + + t.Run("test no call id", func(t *testing.T) { + handler := defaultClearHandler(testClearOffloadPath("/tmp"), true, "read_file") + detail := &ToolDetail{ + ToolContext: &adk.ToolContext{}, + ToolResult: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "test"}}, + }, + } + result, err := handler(ctx, detail) + assert.NoError(t, err) + assert.True(t, result.NeedClear) + assert.Equal(t, "/tmp/clear/generated", result.OffloadFilePath) + }) +} + +func TestStringifyToolOutputParts(t *testing.T) { + t.Run("test empty parts", func(t *testing.T) { + result := stringifyToolOutputParts([]schema.ToolOutputPart{}) + assert.Equal(t, "", result) + }) + + t.Run("test single text part", func(t *testing.T) { + result := stringifyToolOutputParts([]schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "hello world"}, + }) + assert.Equal(t, "hello world", result) + }) + + t.Run("test multiple text parts", func(t *testing.T) { + result := stringifyToolOutputParts([]schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "hello "}, + {Type: schema.ToolPartTypeText, Text: "world"}, + }) + expected := `[ + { + "type": "text", + "text": "hello " + }, + { + "type": "text", + "text": "world" + } +]` + assert.JSONEq(t, expected, result) + }) + + t.Run("test with image part", func(t *testing.T) { + result := stringifyToolOutputParts([]schema.ToolOutputPart{ + {Type: schema.ToolPartTypeImage, Text: "https://example.com/image.png"}, + }) + expected := `[ + { + "type": "image", + "text": "https://example.com/image.png" + } +]` + assert.JSONEq(t, expected, result) + }) + + t.Run("test mixed parts", func(t *testing.T) { + result := stringifyToolOutputParts([]schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "hello world"}, + {Type: schema.ToolPartTypeImage, Text: "https://example.com/image.png"}, + }) + expected := `[ + { + "type": "text", + "text": "hello world" + }, + { + "type": "image", + "text": "https://example.com/image.png" + } +]` + assert.JSONEq(t, expected, result) + }) +} + +func TestClampPrefixToUTF8Boundary(t *testing.T) { + t.Run("test n <= 0", func(t *testing.T) { + result := clampPrefixToUTF8Boundary("hello world", 0) + assert.Equal(t, "", result) + + result = clampPrefixToUTF8Boundary("hello world", -5) + assert.Equal(t, "", result) + }) + + t.Run("test n >= len(s)", func(t *testing.T) { + result := clampPrefixToUTF8Boundary("hello", 10) + assert.Equal(t, "hello", result) + }) + + t.Run("test ASCII string at rune boundary", func(t *testing.T) { + result := clampPrefixToUTF8Boundary("hello world", 5) + assert.Equal(t, "hello", result) + }) + + t.Run("test ASCII string not at rune boundary (should not happen)", func(t *testing.T) { + result := clampPrefixToUTF8Boundary("world", 3) + assert.Equal(t, "wor", result) + }) + + t.Run("test multi-byte UTF-8 characters at boundary", func(t *testing.T) { + s := "你好世界" + result := clampPrefixToUTF8Boundary(s, 3) + assert.Equal(t, "你", result) + }) + + t.Run("test multi-byte UTF-8 characters not at boundary", func(t *testing.T) { + s := "你好世界" + result := clampPrefixToUTF8Boundary(s, 2) + assert.Equal(t, "", result) + + result = clampPrefixToUTF8Boundary(s, 4) + assert.Equal(t, "你", result) + }) + + t.Run("test mixed ASCII and multi-byte characters", func(t *testing.T) { + s := "hello你好world" + result := clampPrefixToUTF8Boundary(s, 5+3) + assert.Equal(t, "hello你", result) + + result = clampPrefixToUTF8Boundary(s, 5+4) + assert.Equal(t, "hello你", result) + + result = clampPrefixToUTF8Boundary(s, 5+6) + assert.Equal(t, "hello你好", result) + }) +} + +func TestClampSuffixToUTF8Boundary(t *testing.T) { + t.Run("test n <= 0", func(t *testing.T) { + result := clampSuffixToUTF8Boundary("hello world", 0) + assert.Equal(t, "", result) + + result = clampSuffixToUTF8Boundary("hello world", -5) + assert.Equal(t, "", result) + }) + + t.Run("test n >= len(s)", func(t *testing.T) { + result := clampSuffixToUTF8Boundary("hello", 10) + assert.Equal(t, "hello", result) + }) + + t.Run("test ASCII string at rune boundary", func(t *testing.T) { + result := clampSuffixToUTF8Boundary("hello world", 5) + assert.Equal(t, "world", result) + }) + + t.Run("test ASCII string not at rune boundary (should not happen)", func(t *testing.T) { + result := clampSuffixToUTF8Boundary("hello", 3) + assert.Equal(t, "llo", result) + }) + + t.Run("test multi-byte UTF-8 characters at boundary", func(t *testing.T) { + s := "你好世界" + assert.Equal(t, 12, len(s)) + result := clampSuffixToUTF8Boundary(s, 3) + assert.Equal(t, "界", result) + }) + + t.Run("test multi-byte UTF-8 characters not at boundary", func(t *testing.T) { + s := "你好世界" + result := clampSuffixToUTF8Boundary(s, 2) + assert.Equal(t, "", result) + + result = clampSuffixToUTF8Boundary(s, 4) + assert.Equal(t, "界", result) + + result = clampSuffixToUTF8Boundary(s, 6) + assert.Equal(t, "世界", result) + }) + + t.Run("test mixed ASCII and multi-byte characters", func(t *testing.T) { + s := "hello你好world" + result := clampSuffixToUTF8Boundary(s, 5+3) + assert.Equal(t, "好world", result) + + result = clampSuffixToUTF8Boundary(s, 5+4) + assert.Equal(t, "好world", result) + + result = clampSuffixToUTF8Boundary(s, 5+6) + assert.Equal(t, "你好world", result) + }) +} + +func TestGetMsgClearedFlag(t *testing.T) { + t.Run("test non-bool value in extra", func(t *testing.T) { + msg := schema.UserMessage("test") + msg.Extra = map[string]any{msgClearedFlag: "not a bool"} + flag := getMsgClearedFlag(msg) + assert.False(t, flag) + }) +} + +func TestCopyMessages(t *testing.T) { + t.Run("test empty messages", func(t *testing.T) { + result := copyMessages(nil) + assert.Len(t, result, 0) + + result = copyMessages([]*schema.Message{}) + assert.Len(t, result, 0) + }) + + t.Run("test message with tool calls and extra", func(t *testing.T) { + msg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_123", + Type: "function", + Function: schema.FunctionCall{Name: "test_func", Arguments: `{"key": "value"}`}, + }, + }) + msg.Extra = map[string]any{"test_key": "test_value"} + msg.UserInputMultiContent = []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "test text"}, + } + msg.MultiContent = []schema.ChatMessagePart{ + {Type: schema.ChatMessagePartTypeText, Text: "multi content"}, + } + msg.AssistantGenMultiContent = []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "assistant gen"}, + } + + result := copyMessages([]*schema.Message{msg}) + assert.Len(t, result, 1) + assert.NotEqual(t, &msg, result[0]) + assert.Len(t, result[0].ToolCalls, 1) + assert.Equal(t, "call_123", result[0].ToolCalls[0].ID) + assert.Equal(t, "test_value", result[0].Extra["test_key"]) + assert.Len(t, result[0].UserInputMultiContent, 1) + assert.Len(t, result[0].MultiContent, 1) + assert.Len(t, result[0].AssistantGenMultiContent, 1) + }) +} + +func TestReductionMiddlewareEnhancedTrunc(t *testing.T) { + ctx := context.Background() + + t.Run("enhanced invokable - passthrough when truncation skipped", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + SkipTruncation: true, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + + called := 0 + endpoint := func(_ context.Context, ta *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + called++ + assert.Same(t, toolArg, ta) + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "ok"}}}, nil + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + got, err := wrapped(ctx, toolArg) + assert.NoError(t, err) + assert.Equal(t, 1, called) + assert.Equal(t, "ok", got.Parts[0].Text) + }) + + t.Run("enhanced invokable - excluded tool bypasses truncation", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + TruncExcludeTools: []string{"enh_tool"}, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + t.Fatalf("trunc handler should not be called when tool is excluded") + return nil, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + got, err := wrapped(ctx, toolArg) + assert.NoError(t, err) + assert.Equal(t, "orig", got.Parts[0].Text) + }) + + t.Run("enhanced invokable - endpoint error returned", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + wantErr := errors.New("endpoint error") + + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + t.Fatalf("trunc handler should not be called on endpoint error") + return nil, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return nil, wantErr + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + _, err = wrapped(ctx, toolArg) + assert.Equal(t, wantErr, err) + }) + + t.Run("enhanced invokable - trunc handler error returned", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + wantErr := errors.New("handler error") + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + return nil, wantErr + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + _, err = wrapped(ctx, toolArg) + assert.Equal(t, wantErr, err) + }) + + t.Run("enhanced invokable - no trunc returns original result", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, detail *ToolDetail) (*TruncResult, error) { + assert.NotNil(t, detail.ToolResult) + assert.Equal(t, "orig", detail.ToolResult.Parts[0].Text) + return &TruncResult{NeedTrunc: false}, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + got, err := wrapped(ctx, toolArg) + assert.NoError(t, err) + assert.Equal(t, "orig", got.Parts[0].Text) + }) + + t.Run("enhanced invokable - offload without backend returns error", func(t *testing.T) { + // config.Backend must be non-nil to pass New(), but the per-tool cfg.Backend is nil. + configBackend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: configBackend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_tool": { + Backend: nil, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + return &TruncResult{ + NeedTrunc: true, + NeedOffload: true, + OffloadFilePath: "/tmp/trunc/cid", + OffloadContent: "big", + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "trunc"}}}, + }, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + _, err = wrapped(ctx, toolArg) + assert.EqualError(t, err, "truncation: no backend for offload") + }) + + t.Run("enhanced invokable - trunc + offload writes and returns trunc result", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, detail *ToolDetail) (*TruncResult, error) { + assert.NotNil(t, detail.ToolContext) + assert.NotNil(t, detail.ToolArgument) + assert.NotNil(t, detail.ToolResult) + return &TruncResult{ + NeedTrunc: true, + NeedOffload: true, + OffloadFilePath: "/tmp/trunc/cid", + OffloadContent: "full content", + ToolResult: &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "truncated"}}}, + }, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_tool", CallID: "cid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, ta *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + assert.Same(t, toolArg, ta) + return &schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil + } + + wrapped, err := mw.WrapEnhancedInvokableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + got, err := wrapped(ctx, toolArg) + assert.NoError(t, err) + assert.Equal(t, "truncated", got.Parts[0].Text) + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/cid"}) + assert.NoError(t, err) + assert.Equal(t, "full content", content.Content) + }) + + t.Run("enhanced streamable - no trunc returns original stream (Copy semantics)", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_stream_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, detail *ToolDetail) (*TruncResult, error) { + // Consume the "analysis" copy; the returned reader should still have the full stream. + parts, needProcess, err := getJointToolResult(detail) + assert.NoError(t, err) + assert.True(t, needProcess) + assert.Len(t, parts, 1) + assert.Equal(t, "hello world", parts[0].Text) + return &TruncResult{NeedTrunc: false}, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_stream_tool", CallID: "scid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr, sw := schema.Pipe[*schema.ToolResult](4) + go func() { + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "hello "}}}, nil) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "world"}}}, nil) + sw.Close() + }() + return sr, nil + } + + wrapped, err := mw.WrapEnhancedStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + resp, err := wrapped(ctx, toolArg) + assert.NoError(t, err) + defer resp.Close() + + var chunks []*schema.ToolResult + for { + tr, recvErr := resp.Recv() + if recvErr != nil { + assert.Equal(t, io.EOF, recvErr) + break + } + chunks = append(chunks, tr) + } + joined, err := schema.ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Equal(t, "hello world", joined.Parts[0].Text) + }) + + t.Run("enhanced streamable - trunc + offload writes and returns trunc stream", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_stream_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + sr, sw := schema.Pipe[*schema.ToolResult](1) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "truncated"}}}, nil) + sw.Close() + return &TruncResult{ + NeedTrunc: true, + NeedOffload: true, + OffloadFilePath: "/tmp/trunc/scid", + OffloadContent: "full content", + StreamToolResult: sr, + }, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_stream_tool", CallID: "scid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr, sw := schema.Pipe[*schema.ToolResult](2) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil) + sw.Close() + return sr, nil + } + + wrapped, err := mw.WrapEnhancedStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + resp, err := wrapped(ctx, toolArg) + assert.NoError(t, err) + defer resp.Close() + + tr, err := resp.Recv() + assert.NoError(t, err) + assert.Equal(t, "truncated", tr.Parts[0].Text) + + _, err = resp.Recv() + assert.Equal(t, io.EOF, err) + + content, err := backend.Read(ctx, &filesystem.ReadRequest{FilePath: "/tmp/trunc/scid"}) + assert.NoError(t, err) + assert.Equal(t, "full content", content.Content) + }) + + t.Run("enhanced streamable - endpoint error returned", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + wantErr := errors.New("endpoint error") + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_stream_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + t.Fatalf("trunc handler should not be called on endpoint error") + return nil, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_stream_tool", CallID: "scid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return nil, wantErr + } + + wrapped, err := mw.WrapEnhancedStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + _, err = wrapped(ctx, toolArg) + assert.Equal(t, wantErr, err) + }) + + t.Run("enhanced streamable - trunc handler error returned", func(t *testing.T) { + backend := filesystem.NewInMemoryBackend() + wantErr := errors.New("handler error") + config := &Config{ + Backend: backend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_stream_tool": { + Backend: backend, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + return nil, wantErr + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_stream_tool", CallID: "scid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr, sw := schema.Pipe[*schema.ToolResult](1) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil) + sw.Close() + return sr, nil + } + + wrapped, err := mw.WrapEnhancedStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + _, err = wrapped(ctx, toolArg) + assert.Equal(t, wantErr, err) + }) + + t.Run("enhanced streamable - offload without backend returns error", func(t *testing.T) { + // config.Backend must be non-nil to pass New(), but the per-tool cfg.Backend is nil. + configBackend := filesystem.NewInMemoryBackend() + config := &Config{ + Backend: configBackend, + ToolConfig: map[string]*ToolReductionConfig{ + "enh_stream_tool": { + Backend: nil, + SkipTruncation: false, + TruncHandler: func(_ context.Context, _ *ToolDetail) (*TruncResult, error) { + sr, sw := schema.Pipe[*schema.ToolResult](1) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "truncated"}}}, nil) + sw.Close() + return &TruncResult{ + NeedTrunc: true, + NeedOffload: true, + OffloadFilePath: "/tmp/trunc/scid", + OffloadContent: "full content", + StreamToolResult: sr, + }, nil + }, + }, + }, + } + mw, err := New(ctx, config) + assert.NoError(t, err) + + tCtx := &adk.ToolContext{Name: "enh_stream_tool", CallID: "scid"} + toolArg := &schema.ToolArgument{Text: `{"k":"v"}`} + endpoint := func(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr, sw := schema.Pipe[*schema.ToolResult](1) + sw.Send(&schema.ToolResult{Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "orig"}}}, nil) + sw.Close() + return sr, nil + } + + wrapped, err := mw.WrapEnhancedStreamableToolCall(ctx, endpoint, tCtx) + assert.NoError(t, err) + _, err = wrapped(ctx, toolArg) + assert.EqualError(t, err, "truncation: no backend for offload") + }) +} + +func TestClearRewriteMessagesHandler(t *testing.T) { + ctx := context.Background() + it := mockInvokableTool() + st := mockStreamableTool() + tools := []tool.BaseTool{it, st} + var toolsInfo []*schema.ToolInfo + for _, bt := range tools { + ti, _ := bt.Info(ctx) + toolsInfo = append(toolsInfo, ti) + } + + tokenCounter := func(_ context.Context, msgs []adk.Message, _ []*schema.ToolInfo) (int64, error) { + return int64(1000), nil + } + + t.Run("test all tools hit - rewrite to system reminder", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{ + schema.UserMessage("write_file and edit_file executed successfully"), + }, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt", "content": "hello"}`}, + }, + { + ID: "call_2", + Type: "function", + Function: schema.FunctionCall{Name: "edit_file", Arguments: `{"file": "test.txt", "changes": "add"}`}, + }, + }), + schema.ToolMessage("write success", "call_1"), + schema.ToolMessage("edit success", "call_2"), + schema.AssistantMessage("done", nil), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Len(t, s.Messages, 5) + assert.Equal(t, schema.User, s.Messages[2].Role) + assert.Equal(t, "write_file and edit_file executed successfully", s.Messages[2].Content) + assert.Equal(t, schema.Assistant, s.Messages[3].Role) + assert.Equal(t, "done", s.Messages[3].Content) + }) + + t.Run("test remove tool call and tool responses", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{}, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt", "content": "hello"}`}, + }, + { + ID: "call_2", + Type: "function", + Function: schema.FunctionCall{Name: "edit_file", Arguments: `{"file": "test.txt", "changes": "add"}`}, + }, + }), + schema.ToolMessage("write success", "call_1"), + schema.ToolMessage("edit success", "call_2"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + schema.ToolMessage("dummy dummy", "dummy"), + schema.AssistantMessage("done", nil), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Len(t, s.Messages, 5) + assert.Equal(t, schema.Assistant, s.Messages[2].Role) + assert.Equal(t, schema.Tool, s.Messages[3].Role) + assert.Equal(t, "dummy dummy", s.Messages[3].Content) + assert.Equal(t, schema.Assistant, s.Messages[4].Role) + assert.Equal(t, "done", s.Messages[4].Content) + }) + + t.Run("test partial tools hit - keep unhit tools and rewrite hit ones", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearExcludeTools: []string{"get_weather"}, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + hitToolIDs := map[string]bool{"call_1": true} + + newToolCalls := make([]schema.ToolCall, 0) + newToolResponses := make([]adk.Message, 0) + + for i, tc := range assistantMessage.ToolCalls { + if !hitToolIDs[tc.ID] { + newToolCalls = append(newToolCalls, tc) + if i < len(toolResponses) { + newToolResponses = append(newToolResponses, toolResponses[i]) + } + } + } + + result := []adk.Message{ + schema.UserMessage("write_file executed successfully"), + } + if len(newToolCalls) > 0 { + result = append(result, schema.AssistantMessage("", newToolCalls)) + result = append(result, newToolResponses...) + } + return result, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt", "content": "hello"}`}, + }, + { + ID: "call_2", + Type: "function", + Function: schema.FunctionCall{Name: "get_weather", Arguments: `{"location": "London"}`}, + }, + }), + schema.ToolMessage("write success", "call_1"), + schema.ToolMessage("sunny", "call_2"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Len(t, s.Messages, 6) + assert.Equal(t, schema.User, s.Messages[2].Role) + assert.Equal(t, "write_file executed successfully", s.Messages[2].Content) + assert.Equal(t, schema.Assistant, s.Messages[3].Role) + assert.Len(t, s.Messages[3].ToolCalls, 1) + assert.Equal(t, "call_2", s.Messages[3].ToolCalls[0].ID) + assert.Equal(t, schema.Tool, s.Messages[4].Role) + assert.Equal(t, "sunny", s.Messages[4].Content) + }) + + t.Run("test mixed messages with system/user before", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{ + schema.UserMessage("tool executed"), + }, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("first request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt"}`}, + }, + }), + schema.ToolMessage("success", "call_1"), + schema.UserMessage("second request"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Len(t, s.Messages, 5) + assert.Equal(t, schema.System, s.Messages[0].Role) + assert.Equal(t, schema.User, s.Messages[1].Role) + assert.Equal(t, schema.User, s.Messages[2].Role) + assert.Equal(t, "tool executed", s.Messages[2].Content) + assert.Equal(t, schema.User, s.Messages[3].Role) + }) + + t.Run("test mixed messages with system/user before", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{ + schema.UserMessage("tool executed"), + }, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("first request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt"}`}, + }, + }), + schema.ToolMessage("success", "call_1"), + schema.UserMessage("second request"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Len(t, s.Messages, 5) + assert.Equal(t, schema.System, s.Messages[0].Role) + assert.Equal(t, schema.User, s.Messages[1].Role) + assert.Equal(t, schema.User, s.Messages[2].Role) + assert.Equal(t, "tool executed", s.Messages[2].Content) + assert.Equal(t, schema.User, s.Messages[3].Role) + }) + + t.Run("test assistant message without tool calls - keep as is", func(t *testing.T) { + handlerCalled := false + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + handlerCalled = true + return nil, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("regular response without tools", nil), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.False(t, handlerCalled) + assert.Len(t, s.Messages, 4) + assert.Equal(t, schema.Assistant, s.Messages[2].Role) + assert.Equal(t, "regular response without tools", s.Messages[2].Content) + }) + + t.Run("test clear not meeting threshold - should keep original messages", func(t *testing.T) { + originalTokenCount := int64(1000) + rewrittenTokenCount := int64(999) + callCount := 0 + + config := &Config{ + SkipTruncation: true, + TokenCounter: func(_ context.Context, msgs []adk.Message, _ []*schema.ToolInfo) (int64, error) { + callCount++ + if callCount == 1 { + return originalTokenCount, nil + } + return rewrittenTokenCount, nil + }, + MaxTokensForClear: 10, + ClearAtLeastTokens: 10, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{ + schema.UserMessage("tool executed"), + }, nil + }, + } + + originalMessages := []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt"}`}, + }, + }), + schema.ToolMessage("success", "call_1"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: originalMessages, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.Len(t, s.Messages, len(originalMessages)) + for i := range originalMessages { + assert.Equal(t, originalMessages[i].Role, s.Messages[i].Role) + assert.Equal(t, originalMessages[i].Content, s.Messages[i].Content) + } + }) + + t.Run("test applyClearRewrite with default role", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{}, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + unknownRoleMsg := &schema.Message{ + Role: "unknown_role", + Content: "unknown", + } + + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + unknownRoleMsg, + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.NotNil(t, s) + }) + + t.Run("test applyClearRewrite with rewrite error", func(t *testing.T) { + expectedErr := fmt.Errorf("rewrite error") + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return nil, expectedErr + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + + _, _, err = mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt"}`}, + }, + }), + schema.ToolMessage("success", "call_1"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.Error(t, err) + assert.Equal(t, expectedErr, err) + }) + + t.Run("test applyClearRewrite with missing tool responses", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 0, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + assert.Nil(t, toolResponses) + return []adk.Message{ + schema.UserMessage("tool executed"), + }, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt"}`}, + }, + }), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.NotNil(t, s) + }) + + t.Run("test applyClearRewrite with clearAtLeastTokens > 0", func(t *testing.T) { + config := &Config{ + SkipTruncation: true, + TokenCounter: tokenCounter, + MaxTokensForClear: 10, + ClearAtLeastTokens: 10, + ClearMessageRewriter: func(ctx context.Context, assistantMessage adk.Message, toolResponses []adk.Message) (messagesAfterRewrite []adk.Message, err error) { + return []adk.Message{ + schema.UserMessage("tool executed"), + }, nil + }, + } + + mw, err := New(ctx, config) + assert.NoError(t, err) + _, s, err := mw.BeforeModelRewriteState(ctx, &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("you are a helpful assistant"), + schema.UserMessage("user request"), + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{Name: "write_file", Arguments: `{"file": "test.txt"}`}, + }, + }), + schema.ToolMessage("success", "call_1"), + schema.AssistantMessage("", []schema.ToolCall{{ID: "dummy", Type: "function", Function: schema.FunctionCall{Name: "dummy_tool"}}}), + }, + }, &adk.ModelContext{ + Tools: toolsInfo, + }) + assert.NoError(t, err) + assert.NotNil(t, s) + }) +} + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped(ctx, &TypedConfig[*schema.AgenticMessage]{ + SkipTruncation: true, + SkipClear: true, + }) + assert.NoError(t, err) + assert.NotNil(t, mw) + + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} diff --git a/adk/middlewares/skill/doc.go b/adk/middlewares/skill/doc.go new file mode 100644 index 0000000..9eea039 --- /dev/null +++ b/adk/middlewares/skill/doc.go @@ -0,0 +1,47 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package skill provides a Skill middleware, types, and a local filesystem backend. +// +// # Overview +// +// The Skill middleware is a ChatModelAgentMiddleware implementation that injects: +// - a system instruction (Skills System) +// - a tool (default name: "skill") to load and execute skills +// +// Skill definitions are stored in SKILL.md files with a YAML frontmatter. The frontmatter is +// parsed into FrontMatter (name/description/context/agent/model), and the remaining markdown +// is treated as the skill content. +// +// # Execution modes +// +// Skill execution is controlled by frontmatter "context": +// - inline (default): returns the skill content as the tool result in the current agent +// - fork: runs the skill with a new sub-agent without parent message history +// - fork_with_context: runs the skill with a new sub-agent carrying parent message history +// +// # Extension points +// +// - CustomToolParams customizes the tool parameter schema. +// - BuildContent customizes how skill content is generated from raw tool arguments. +// - BuildForkMessages customizes the initial messages passed to the sub-agent in fork modes. +// - FormatForkResult customizes how sub-agent outputs are formatted back to the caller. +// +// # Filesystem backend +// +// NewBackendFromFilesystem loads skills from a filesystem backend. It scans only first-level +// subdirectories under BaseDir and reads each /SKILL.md as a skill definition. +package skill diff --git a/adk/middlewares/skill/filesystem_backend.go b/adk/middlewares/skill/filesystem_backend.go new file mode 100644 index 0000000..a337a5a --- /dev/null +++ b/adk/middlewares/skill/filesystem_backend.go @@ -0,0 +1,186 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package skill + +import ( + "context" + "fmt" + "path/filepath" + "strings" + + "gopkg.in/yaml.v3" + + "github.com/cloudwego/eino/adk/filesystem" +) + +const skillFileName = "SKILL.md" + +type filesystemBackend struct { + backend filesystem.Backend + baseDir string +} + +// BackendFromFilesystemConfig contains configuration for NewBackendFromFilesystem. +type BackendFromFilesystemConfig struct { + // Backend is the filesystem.Backend implementation used for file operations. + Backend filesystem.Backend + // BaseDir is the base directory where skill directories are located. + // Each skill should be in a subdirectory containing a SKILL.md file. + BaseDir string +} + +// NewBackendFromFilesystem creates a new Backend implementation that reads skills from a filesystem. +// It searches for SKILL.md files in immediate subdirectories of the configured BaseDir. +// Only first-level subdirectories are scanned; deeply nested SKILL.md files are ignored. +func NewBackendFromFilesystem(_ context.Context, config *BackendFromFilesystemConfig) (Backend, error) { + if config == nil { + return nil, fmt.Errorf("config is required") + } + if config.Backend == nil { + return nil, fmt.Errorf("backend is required") + } + if config.BaseDir == "" { + return nil, fmt.Errorf("baseDir is required") + } + + return &filesystemBackend{ + backend: config.Backend, + baseDir: config.BaseDir, + }, nil +} + +func (b *filesystemBackend) List(ctx context.Context) ([]FrontMatter, error) { + skills, err := b.list(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list skills: %w", err) + } + + matters := make([]FrontMatter, 0, len(skills)) + for _, skill := range skills { + matters = append(matters, skill.FrontMatter) + } + + return matters, nil +} + +func (b *filesystemBackend) Get(ctx context.Context, name string) (Skill, error) { + skills, err := b.list(ctx) + if err != nil { + return Skill{}, fmt.Errorf("failed to list skills: %w", err) + } + + for _, skill := range skills { + if skill.Name == name { + return skill, nil + } + } + + return Skill{}, fmt.Errorf("skill not found: %s", name) +} + +func (b *filesystemBackend) list(ctx context.Context) ([]Skill, error) { + var skills []Skill + + pattern := "*/" + skillFileName + entries, err := b.backend.GlobInfo(ctx, &filesystem.GlobInfoRequest{ + Pattern: pattern, + Path: b.baseDir, + }) + if err != nil { + return nil, fmt.Errorf("failed to glob skill files: %w", err) + } + + for _, entry := range entries { + filePath := entry.Path + if !filepath.IsAbs(filePath) { + filePath = filepath.Join(b.baseDir, filePath) + } + skill, loadErr := b.loadSkillFromFile(ctx, filePath) + if loadErr != nil { + return nil, fmt.Errorf("failed to load skill from %s: %w", filePath, loadErr) + } + + skills = append(skills, skill) + } + + return skills, nil +} + +func (b *filesystemBackend) loadSkillFromFile(ctx context.Context, path string) (Skill, error) { + fileContent, err := b.backend.Read(ctx, &filesystem.ReadRequest{ + FilePath: path, + }) + if err != nil { + return Skill{}, fmt.Errorf("failed to read file: %w", err) + } + + data := stripLineNumbers(fileContent.Content) + + frontmatter, content, err := parseFrontmatter(data) + if err != nil { + return Skill{}, fmt.Errorf("failed to parse frontmatter: %w", err) + } + + var fm FrontMatter + if err = yaml.Unmarshal([]byte(frontmatter), &fm); err != nil { + return Skill{}, fmt.Errorf("failed to unmarshal frontmatter: %w", err) + } + + absDir := filepath.Dir(path) + + return Skill{ + FrontMatter: fm, + Content: strings.TrimSpace(content), + BaseDirectory: absDir, + }, nil +} + +func stripLineNumbers(data string) string { + lines := strings.Split(data, "\n") + result := make([]string, 0, len(lines)) + for _, line := range lines { + idx := strings.Index(line, "\t") + if idx != -1 { + line = line[idx+1:] + } + result = append(result, line) + } + return strings.Join(result, "\n") +} + +func parseFrontmatter(data string) (frontmatter string, content string, err error) { + const delimiter = "---" + + data = strings.TrimSpace(data) + + if !strings.HasPrefix(data, delimiter) { + return "", "", fmt.Errorf("file does not start with frontmatter delimiter") + } + + rest := data[len(delimiter):] + endIdx := strings.Index(rest, "\n"+delimiter) + if endIdx == -1 { + return "", "", fmt.Errorf("frontmatter closing delimiter not found") + } + + frontmatter = strings.TrimSpace(rest[:endIdx]) + content = rest[endIdx+len("\n"+delimiter):] + + content = strings.TrimPrefix(content, "\n") + + return frontmatter, content, nil +} diff --git a/adk/middlewares/skill/filesystem_backend_test.go b/adk/middlewares/skill/filesystem_backend_test.go new file mode 100644 index 0000000..3b8743c --- /dev/null +++ b/adk/middlewares/skill/filesystem_backend_test.go @@ -0,0 +1,500 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package skill + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk/filesystem" +) + +func TestNewBackendFromFilesystem(t *testing.T) { + ctx := context.Background() + + t.Run("nil config returns error", func(t *testing.T) { + backend, err := NewBackendFromFilesystem(ctx, nil) + assert.Nil(t, backend) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config is required") + }) + + t.Run("nil backend returns error", func(t *testing.T) { + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + BaseDir: "/skills", + }) + assert.Nil(t, backend) + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend is required") + }) + + t.Run("empty baseDir returns error", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "", + }) + assert.Nil(t, backend) + assert.Error(t, err) + assert.Contains(t, err.Error(), "baseDir is required") + }) + + t.Run("valid config succeeds", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + assert.NoError(t, err) + assert.NotNil(t, backend) + }) +} + +func TestFilesystemBackend_List(t *testing.T) { + ctx := context.Background() + + t.Run("empty directory returns empty list", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/.keep", + Content: "", + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + assert.Empty(t, skills) + }) + + t.Run("directory with no SKILL.md files returns empty list", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/subdir/other.txt", + Content: "some content", + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + assert.Empty(t, skills) + }) + + t.Run("files in root directory are ignored", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/SKILL.md", + Content: `--- +name: root-skill +description: Root skill +--- +Content`, + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + assert.Empty(t, skills) + }) + + t.Run("valid skill directory returns skill", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/my-skill/SKILL.md", + Content: `--- +name: pdf-processing +description: Extract text and tables from PDF files, fill forms, merge documents. +license: Apache-2.0 +metadata: + author: example-org + version: "1.0" +--- +This is the skill content.`, + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + require.Len(t, skills, 1) + assert.Equal(t, "pdf-processing", skills[0].Name) + assert.Equal(t, "Extract text and tables from PDF files, fill forms, merge documents.", skills[0].Description) + }) + + t.Run("multiple skill directories returns all skills", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/skill-1/SKILL.md", + Content: `--- +name: skill-1 +description: First skill +--- +Content 1`, + }) + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/skill-2/SKILL.md", + Content: `--- +name: skill-2 +description: Second skill +--- +Content 2`, + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + assert.Len(t, skills, 2) + + names := []string{skills[0].Name, skills[1].Name} + assert.Contains(t, names, "skill-1") + assert.Contains(t, names, "skill-2") + }) + + t.Run("invalid SKILL.md returns error", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/invalid-skill/SKILL.md", + Content: `No frontmatter here`, + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.Error(t, err) + assert.Nil(t, skills) + assert.Contains(t, err.Error(), "failed to load skill") + }) + + t.Run("non-existent baseDir returns empty list", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/path/that/does/not/exist", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + assert.Empty(t, skills) + }) + + t.Run("deeply nested SKILL.md is ignored", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/valid-skill/SKILL.md", + Content: `--- +name: valid-skill +description: Valid skill +--- +Content`, + }) + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/deep/nested/SKILL.md", + Content: `--- +name: nested-skill +description: Nested skill +--- +Content`, + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skills, err := backend.List(ctx) + assert.NoError(t, err) + assert.Len(t, skills, 1) + assert.Equal(t, "valid-skill", skills[0].Name) + }) +} + +func TestFilesystemBackend_Get(t *testing.T) { + ctx := context.Background() + + t.Run("skill not found returns error", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/.keep", + Content: "", + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skill, err := backend.Get(ctx, "non-existent") + assert.Error(t, err) + assert.Empty(t, skill) + assert.Contains(t, err.Error(), "skill not found") + }) + + t.Run("existing skill is returned", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/test-skill/SKILL.md", + Content: `--- +name: test-skill +description: Test skill description +--- +Test content here.`, + }) + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skill, err := backend.Get(ctx, "test-skill") + assert.NoError(t, err) + assert.Equal(t, "test-skill", skill.Name) + assert.Equal(t, "Test skill description", skill.Description) + assert.Equal(t, "Test content here.", skill.Content) + }) + + t.Run("get specific skill from multiple", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + for _, name := range []string{"alpha", "beta", "gamma"} { + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/" + name + "/SKILL.md", + Content: `--- +name: ` + name + ` +description: Skill ` + name + ` +--- +Content for ` + name, + }) + } + + backend, err := NewBackendFromFilesystem(ctx, &BackendFromFilesystemConfig{ + Backend: fsBackend, + BaseDir: "/skills", + }) + require.NoError(t, err) + + skill, err := backend.Get(ctx, "beta") + assert.NoError(t, err) + assert.Equal(t, "beta", skill.Name) + assert.Equal(t, "Skill beta", skill.Description) + assert.Equal(t, "Content for beta", skill.Content) + }) +} + +func TestParseFrontmatter(t *testing.T) { + t.Run("valid frontmatter", func(t *testing.T) { + data := `--- +name: test +description: test description +--- +This is the content.` + + fm, content, err := parseFrontmatter(data) + assert.NoError(t, err) + assert.Equal(t, "name: test\ndescription: test description", fm) + assert.Equal(t, "This is the content.", content) + }) + + t.Run("frontmatter with multiline content", func(t *testing.T) { + data := `--- +name: test +--- +Line 1 +Line 2 +Line 3` + + fm, content, err := parseFrontmatter(data) + assert.NoError(t, err) + assert.Equal(t, "name: test", fm) + assert.Equal(t, "Line 1\nLine 2\nLine 3", content) + }) + + t.Run("frontmatter with leading/trailing whitespace", func(t *testing.T) { + data := ` +--- +name: test +--- +Content ` + + fm, content, err := parseFrontmatter(data) + assert.NoError(t, err) + assert.Equal(t, "name: test", fm) + assert.Equal(t, "Content", content) + }) + + t.Run("missing opening delimiter returns error", func(t *testing.T) { + data := `name: test +--- +Content` + + fm, content, err := parseFrontmatter(data) + assert.Error(t, err) + assert.Empty(t, fm) + assert.Empty(t, content) + assert.Contains(t, err.Error(), "does not start with frontmatter delimiter") + }) + + t.Run("missing closing delimiter returns error", func(t *testing.T) { + data := `--- +name: test +Content without closing` + + fm, content, err := parseFrontmatter(data) + assert.Error(t, err) + assert.Empty(t, fm) + assert.Empty(t, content) + assert.Contains(t, err.Error(), "closing delimiter not found") + }) + + t.Run("empty frontmatter", func(t *testing.T) { + data := `--- +--- +Content only` + + fm, content, err := parseFrontmatter(data) + assert.NoError(t, err) + assert.Empty(t, fm) + assert.Equal(t, "Content only", content) + }) + + t.Run("empty content", func(t *testing.T) { + data := `--- +name: test +---` + + fm, content, err := parseFrontmatter(data) + assert.NoError(t, err) + assert.Equal(t, "name: test", fm) + assert.Empty(t, content) + }) + + t.Run("content with --- inside", func(t *testing.T) { + data := `--- +name: test +--- +Content with --- in the middle` + + fm, content, err := parseFrontmatter(data) + assert.NoError(t, err) + assert.Equal(t, "name: test", fm) + assert.Equal(t, "Content with --- in the middle", content) + }) +} + +func TestLoadSkillFromFile(t *testing.T) { + ctx := context.Background() + + t.Run("valid skill file", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/SKILL.md", + Content: `--- +name: file-skill +description: Skill from file +--- +File skill content.`, + }) + + backend := &filesystemBackend{backend: fsBackend, baseDir: "/skills"} + skill, err := backend.loadSkillFromFile(ctx, "/skills/SKILL.md") + assert.NoError(t, err) + assert.Equal(t, "file-skill", skill.Name) + assert.Equal(t, "Skill from file", skill.Description) + assert.Equal(t, "File skill content.", skill.Content) + assert.Equal(t, "/skills", skill.BaseDirectory) + }) + + t.Run("non-existent file returns error", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + backend := &filesystemBackend{backend: fsBackend, baseDir: "/tmp"} + skill, err := backend.loadSkillFromFile(ctx, "/path/to/nonexistent/SKILL.md") + assert.Error(t, err) + assert.Empty(t, skill) + assert.Contains(t, err.Error(), "failed to read file") + }) + + t.Run("invalid yaml in frontmatter returns error", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/SKILL.md", + Content: `--- +name: [invalid yaml +--- +Content`, + }) + + backend := &filesystemBackend{backend: fsBackend, baseDir: "/skills"} + skill, err := backend.loadSkillFromFile(ctx, "/skills/SKILL.md") + assert.Error(t, err) + assert.Empty(t, skill) + assert.Contains(t, err.Error(), "failed to unmarshal frontmatter") + }) + + t.Run("content with extra whitespace is trimmed", func(t *testing.T) { + fsBackend := filesystem.NewInMemoryBackend() + _ = fsBackend.Write(ctx, &filesystem.WriteRequest{ + FilePath: "/skills/SKILL.md", + Content: `--- +name: trimmed-skill +description: desc +--- + + Content with whitespace + +`, + }) + + backend := &filesystemBackend{backend: fsBackend, baseDir: "/skills"} + skill, err := backend.loadSkillFromFile(ctx, "/skills/SKILL.md") + assert.NoError(t, err) + assert.Equal(t, "Content with whitespace", skill.Content) + }) +} diff --git a/adk/middlewares/skill/prompt.go b/adk/middlewares/skill/prompt.go new file mode 100644 index 0000000..87ae617 --- /dev/null +++ b/adk/middlewares/skill/prompt.go @@ -0,0 +1,154 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package skill + +const ( + systemPrompt = ` +# Skills System + +**How to Use Skills (Progressive Disclosure):** + +Skills follow a **progressive disclosure** pattern - you see their name and description above, but only read full instructions when needed: + +1. **Recognize when a skill applies**: Check if the user's task matches a skill's description +2. **Read the skill's full instructions**: Use the '{tool_name}' tool to load skill +3. **Follow the skill's instructions**: tool result contains step-by-step workflows, best practices, and examples +4. **Access supporting files**: Skills may include helper scripts, configs, or reference docs - use absolute paths + +**When to Use Skills:** +- User's request matches a skill's domain (e.g., "research X" -> web-research skill) +- You need specialized knowledge or structured workflows +- A skill provides proven patterns for complex tasks + +**Executing Skill Scripts:** +Skills may contain Python scripts or other executable files. Always use absolute paths. + +**Example Workflow:** + +User: "Can you research the latest developments in quantum computing?" + +1. Check available skills -> See "web-research" skill +2. Call '{tool_name}' tool to read the full skill instructions +3. Follow the skill's research workflow (search -> organize -> synthesize) +4. Use any helper scripts with absolute paths + +Remember: Skills make you more capable and consistent. When in doubt, check if a skill exists for the task! +` + + systemPromptChinese = ` +# Skill 系统 + +**如何使用 Skill（技能）（渐进式展示）：** + +Skill 遵循**渐进式展示**模式 - 你可以在上方看到 Skill 的名称和描述，但只在需要时才阅读完整说明： + +1. **识别 Skill 适用场景**：检查用户的任务是否匹配某个 Skill 的描述 +2. **阅读 Skill 的完整说明**：使用 '{tool_name}' 工具加载 Skill +3. **遵循 Skill 说明操作**：工具结果包含逐步工作流程、最佳实践和示例 +4. **访问支持文件**：Skill 可能包含辅助脚本、配置或参考文档 - 使用绝对路径访问 + +**何时使用 Skill：** +- 用户请求匹配某个 Skill 的领域（例如"研究 X" -> web-research Skill） +- 你需要专业知识或结构化工作流程 +- 某个 Skill 为复杂任务提供了经过验证的模式 + +**执行 Skill 脚本：** +Skill 可能包含 Python 脚本或其他可执行文件。始终使用绝对路径。 + +**示例工作流程：** + +用户："你能研究一下量子计算的最新发展吗？" + +1. 检查可用 Skill -> 发现 "web-research" Skill +2. 调用 '{tool_name}' 工具读取完整的 Skill 说明 +3. 遵循 Skill 的研究工作流程（搜索 -> 整理 -> 综合） +4. 使用绝对路径运行任何辅助脚本 + +记住：Skill 让你更加强大和稳定。如有疑问，请检查是否存在适用于该任务的 Skill！ +` + + toolDescriptionBase = `Execute a skill within the main conversation + + +When users ask you to perform tasks, check if any of the available skills below can help complete the task more effectively. Skills provide specialized capabilities and domain knowledge. + +How to invoke: +- Use the exact string inside tag as the skill name (no arguments) +- Examples: + - ` + "`" + `skill: "pdf"` + "`" + ` - invoke the pdf skill + - ` + "`" + `skill: "xlsx"` + "`" + ` - invoke the xlsx skill + - ` + "`" + `skill: "ms-office-suite:pdf"` + "`" + ` - invoke using fully qualified name + +Important: +- When a skill is relevant, you must invoke this tool IMMEDIATELY as your first action +- NEVER just announce or mention a skill in your text response without actually calling this tool +- This is a BLOCKING REQUIREMENT: invoke the relevant Skill tool BEFORE generating any other response about the task +- Only use skills listed in below +- Do not invoke a skill that is already running +- Skill content may contain relative paths. Convert them to absolute paths using the base directory provided in the tool result + + +` + toolDescriptionBaseChinese = `在主对话中执行 Skill（技能） + + +当用户要求你执行任务时，检查下方可用 Skill 列表中是否有 Skill 可以更有效地完成任务。Skill 提供专业能力和领域知识。 + +如何调用： +- 使用 标签内的完整字符串作为 Skill 名称（无需其他参数） +- 示例： + - ` + "`" + `skill: "pdf"` + "`" + ` - 调用 pdf Skill + - ` + "`" + `skill: "xlsx"` + "`" + ` - 调用 xlsx Skill + - ` + "`" + `skill: "ms-office-suite:pdf"` + "`" + ` - 使用完全限定名称调用 + +重要说明： +- 当 Skill 相关时，你必须立即调用此工具作为第一个动作 +- 切勿仅在文本回复中提及 Skill 而不实际调用此工具 +- 这是阻塞性要求：在生成任何关于任务的其他响应之前，先调用相关的 Skill 工具 +- 仅使用 中列出的 Skill +- 不要调用已经运行中的 Skill +- Skill 内容中可能包含相对路径，需使用工具返回的 base directory 将其转换为绝对路径 + + +` + toolDescriptionTemplate = ` + +{{- range .Matters }} + + +{{ .Name }} + + +{{ .Description }} + + +{{- end }} + +` + toolResult = "Launching skill: %s\n" + toolResultChinese = "正在启动 Skill：%s\n" + userContent = `Base directory for this skill: %s + +%s` + userContentChinese = `此 Skill 的目录：%s + +%s` + toolName = "skill" + + subAgentResultFormat = "Skill \"%s\" completed (sub-agent execution).\n\nResult:\n%s" + subAgentResultFormatChinese = "Skill \"%s\" 已完成（子 Agent 执行）。\n\n结果：\n%s" +) diff --git a/adk/middlewares/skill/skill.go b/adk/middlewares/skill/skill.go new file mode 100644 index 0000000..8f8b2ca --- /dev/null +++ b/adk/middlewares/skill/skill.go @@ -0,0 +1,711 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package skill + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "strings" + "text/template" + + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// ContextMode defines the execution mode of a skill. +type ContextMode string + +const ( + // ContextModeFork creates a new sub-agent without parent history + ContextModeFork ContextMode = "fork" + // ContextModeForkWithContext creates a new sub-agent with parent history + ContextModeForkWithContext ContextMode = "fork_with_context" +) + +// FrontMatter defines the YAML frontmatter schema parsed from a SKILL.md file. +type FrontMatter struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + Context ContextMode `yaml:"context"` + Agent string `yaml:"agent"` + Model string `yaml:"model"` +} + +// Skill represents a skill loaded from a backend. +type Skill struct { + FrontMatter + // Content is the markdown body after the frontmatter contains the skill instructions of a SKILL.md file. + Content string + // BaseDirectory is the absolute directory path where the SKILL.md file is located (e.g., "/absolute/path/to/skills/my-skill"). + BaseDirectory string +} + +// Backend loads skills and provides metadata for tool description rendering. +type Backend interface { + List(ctx context.Context) ([]FrontMatter, error) + Get(ctx context.Context, name string) (Skill, error) +} + +// TypedAgentHubOptions contains options passed to TypedAgentHub.Get when creating an agent for skill execution. +type TypedAgentHubOptions[M adk.MessageType] struct { + // Model is the resolved model instance when a skill specifies a "model" field in frontmatter. + // nil means the skill did not specify a model override; implementations should use their default. + Model model.BaseModel[M] +} + +// AgentHubOptions is a backward-compatible alias for TypedAgentHubOptions instantiated with *schema.Message. +type AgentHubOptions = TypedAgentHubOptions[*schema.Message] + +// TypedAgentHub provides agent instances for context mode (fork/fork_with_context) execution. +type TypedAgentHub[M adk.MessageType] interface { + // Get returns an Agent by name. When name is empty, implementations should return a default agent. + // The opts parameter carries skill-level overrides (e.g., model) resolved by the framework. + Get(ctx context.Context, name string, opts *TypedAgentHubOptions[M]) (adk.TypedAgent[M], error) +} + +// AgentHub is a backward-compatible alias for TypedAgentHub instantiated with *schema.Message. +type AgentHub = TypedAgentHub[*schema.Message] + +// TypedModelHub resolves model instances by name for skills that specify a "model" field in frontmatter. +type TypedModelHub[M adk.MessageType] interface { + Get(ctx context.Context, name string) (model.BaseModel[M], error) +} + +// ModelHub is a backward-compatible alias for TypedModelHub instantiated with *schema.Message. +type ModelHub = TypedModelHub[*schema.Message] + +// SystemPromptFunc is a function that returns a custom system prompt. +// The toolName parameter is the name of the skill tool (default: "skill"). +type SystemPromptFunc func(ctx context.Context, toolName string) string + +// ToolDescriptionFunc is a function that returns a custom tool description. +// The skills parameter contains all available skill front matters. +type ToolDescriptionFunc func(ctx context.Context, skills []FrontMatter) string + +// TypedSubAgentInput contains the context available when building the sub-agent's +// initial messages in fork/fork_with_context mode. +type TypedSubAgentInput[M adk.MessageType] struct { + Skill Skill + Mode ContextMode + RawArguments string + SkillContent string + History []M + ToolCallID string +} + +// SubAgentInput is a backward-compatible alias for TypedSubAgentInput instantiated with *schema.Message. +type SubAgentInput = TypedSubAgentInput[*schema.Message] + +// TypedSubAgentOutput contains the sub-agent's execution results, available when +// formatting the final tool response. +type TypedSubAgentOutput[M adk.MessageType] struct { + Skill Skill + Mode ContextMode + RawArguments string + Messages []M + Results []string +} + +// SubAgentOutput is a backward-compatible alias for TypedSubAgentOutput instantiated with *schema.Message. +type SubAgentOutput = TypedSubAgentOutput[*schema.Message] + +// TypedConfig is the configuration for the skill middleware. +type TypedConfig[M adk.MessageType] struct { + // Backend is the backend for retrieving skills. + Backend Backend + // SkillToolName is the custom name for the skill tool. If nil, the default name "skill" is used. + SkillToolName *string + // Deprecated: Use adk.SetLanguage(adk.LanguageChinese) instead to enable Chinese prompts globally. + // This field will be removed in a future version. + UseChinese bool + // AgentHub provides agent instances for context mode (fork/fork_with_context) execution. + // Required when skills use "context: fork" or "context: fork_with_context" in frontmatter. + // The agent factory is retrieved by agent name (skill.Agent) from this hub. + // When skill.Agent is empty, AgentHub.Get is called with an empty string, + // allowing the hub implementation to return a default agent. + AgentHub TypedAgentHub[M] + // ModelHub provides model instances for skills that specify a "model" field in frontmatter. + // Used in two scenarios: + // - With context mode (fork/fork_with_context): The model is passed to the AgentHub + // - Without context mode (inline): The model becomes active for subsequent ChatModel requests + // If nil, skills with model specification will be ignored in inline mode, + // or return an error in context mode. + ModelHub TypedModelHub[M] + + // CustomSystemPrompt allows customizing the system prompt injected into the agent. + // If nil, the default system prompt is used. + // The function receives the skill tool name as a parameter. + CustomSystemPrompt SystemPromptFunc + // CustomToolDescription allows customizing the tool description for the skill tool. + // If nil, the default tool description is used. + // The function receives all available skill front matters as a parameter. + CustomToolDescription ToolDescriptionFunc + + // CustomToolParams customizes tool parameters for the skill tool. + // defaults is the default schema with only the required "skill" field. + // optional + CustomToolParams func(ctx context.Context, defaults map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) + + // BuildContent customizes the skill content generated for this invocation. + // rawArgs contains the original tool call arguments in JSON form. + // optional + BuildContent func(ctx context.Context, skill Skill, rawArgs string) (string, error) + + // BuildForkMessages customizes the messages passed to the forked sub-agent. + // When nil, fork uses [UserMessage(skillContent)] and fork_with_context uses + // [history..., ToolMessage(skillContent, toolCallID)]. + // optional + BuildForkMessages func(ctx context.Context, in TypedSubAgentInput[M]) ([]M, error) + + // FormatForkResult customizes the final text returned from the forked sub-agent results. + // When nil, assistant message contents emitted by the sub-agent are concatenated and returned + // in a default formatted string. + // optional + FormatForkResult func(ctx context.Context, in TypedSubAgentOutput[M]) (string, error) +} + +// Config is a backward-compatible alias for TypedConfig instantiated with *schema.Message. +type Config = TypedConfig[*schema.Message] + +// NewTyped creates a generic skill middleware handler for TypedChatModelAgent. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +// For *schema.AgenticMessage, tool execution is message-type-independent; the model override +// via ModelHub only takes effect when M is *schema.Message (for other types it is a no-op). +// +// See NewMiddleware for full usage documentation. +func NewTyped[M adk.MessageType](ctx context.Context, config *TypedConfig[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + if config == nil { + return nil, fmt.Errorf("config is required") + } + if config.Backend == nil { + return nil, fmt.Errorf("backend is required") + } + + name := toolName + if config.SkillToolName != nil { + name = *config.SkillToolName + } + + var instruction string + if config.CustomSystemPrompt != nil { + instruction = config.CustomSystemPrompt(ctx, name) + } else { + var err error + instruction, err = buildSystemPrompt(name, config.UseChinese) + if err != nil { + return nil, err + } + } + + return &typedSkillHandler[M]{ + instruction: instruction, + tool: &typedSkillTool[M]{ + b: config.Backend, + toolName: name, + useChinese: config.UseChinese, + agentHub: config.AgentHub, + modelHub: config.ModelHub, + customToolDesc: config.CustomToolDescription, + customToolParams: config.CustomToolParams, + buildContent: config.BuildContent, + buildForkMessages: config.BuildForkMessages, + formatForkResult: config.FormatForkResult, + }, + }, nil +} + +// NewMiddleware creates a new skill middleware handler for ChatModelAgent. +// +// The handler provides a skill tool that allows agents to load and execute skills +// defined in SKILL.md files. Skills can run in different modes based on their +// frontmatter configuration: +// +// - Inline mode (default): Skill content is returned directly as tool result +// - Fork mode (context: fork): Forks a new agent with a clean context, discarding message history +// - Fork with context mode (context: fork_with_context): Forks a new agent carrying over message history +// +// Example usage: +// +// handler, err := skill.NewMiddleware(ctx, &skill.Config{ +// Backend: backend, +// AgentHub: myAgentHub, +// ModelHub: myModelHub, +// }) +// if err != nil { +// return err +// } +// +// agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ +// // ... +// Handlers: []adk.ChatModelAgentMiddleware{handler}, +// }) +func NewMiddleware(ctx context.Context, config *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped(ctx, config) +} + +type typedSkillHandler[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + instruction string + tool *typedSkillTool[M] +} + +func (h *typedSkillHandler[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + runCtx.Instruction = runCtx.Instruction + "\n" + h.instruction + runCtx.Tools = append(runCtx.Tools, h.tool) + return ctx, runCtx, nil +} + +func (h *typedSkillHandler[M]) WrapModel(ctx context.Context, m model.BaseModel[M], _ *adk.TypedModelContext[M]) (model.BaseModel[M], error) { + if h.tool.modelHub == nil { + return m, nil + } + modelName, found, err := adk.GetRunLocalValue(ctx, activeModelKey) + if err != nil { + return nil, fmt.Errorf("failed to get active model from run local value: %w", err) + } + if !found { + return m, nil + } + name, ok := modelName.(string) + if !ok || name == "" { + return m, nil + } + newModel, err := h.tool.modelHub.Get(ctx, name) + if err != nil { + return nil, fmt.Errorf("failed to get model '%s' from ModelHub: %w", name, err) + } + return newModel, nil +} + +const activeModelKey = "__skill_active_model__" + +// New creates a new skill middleware. +// It provides a tool for the agent to use skills. +// +// Deprecated: Use NewMiddleware instead. New does not support fork mode execution +// because AgentMiddleware cannot save message history for fork mode. +func New(ctx context.Context, config *Config) (adk.AgentMiddleware, error) { + if config == nil { + return adk.AgentMiddleware{}, fmt.Errorf("config is required") + } + if config.Backend == nil { + return adk.AgentMiddleware{}, fmt.Errorf("backend is required") + } + + name := toolName + if config.SkillToolName != nil { + name = *config.SkillToolName + } + + var sp string + if config.CustomSystemPrompt != nil { + sp = config.CustomSystemPrompt(ctx, name) + } else { + var err error + sp, err = buildSystemPrompt(name, config.UseChinese) + if err != nil { + return adk.AgentMiddleware{}, err + } + } + + return adk.AgentMiddleware{ + AdditionalInstruction: sp, + AdditionalTools: []tool.BaseTool{&typedSkillTool[*schema.Message]{ + b: config.Backend, + toolName: name, + useChinese: config.UseChinese, + customToolDesc: config.CustomToolDescription, + }}, + }, nil +} + +func buildSystemPrompt(skillToolName string, useChinese bool) (string, error) { + var prompt string + if useChinese { + prompt = systemPromptChinese + } else { + prompt = internal.SelectPrompt(internal.I18nPrompts{ + English: systemPrompt, + Chinese: systemPromptChinese, + }) + } + return pyfmt.Fmt(prompt, map[string]string{ + "tool_name": skillToolName, + }) +} + +type typedSkillTool[M adk.MessageType] struct { + b Backend + toolName string + + useChinese bool + agentHub TypedAgentHub[M] + modelHub TypedModelHub[M] + + customToolDesc ToolDescriptionFunc + + customToolParams func(ctx context.Context, defaults map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) + buildContent func(ctx context.Context, skill Skill, rawArgs string) (string, error) + + buildForkMessages func(ctx context.Context, in TypedSubAgentInput[M]) ([]M, error) + formatForkResult func(ctx context.Context, in TypedSubAgentOutput[M]) (string, error) +} + +type descriptionTemplateHelper struct { + Matters []FrontMatter +} + +func (s *typedSkillTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) { + skills, err := s.b.List(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list skills: %w", err) + } + + var fullDesc string + if s.customToolDesc != nil { + fullDesc = s.customToolDesc(ctx, skills) + } else { + desc, renderErr := renderToolDescription(skills) + if renderErr != nil { + return nil, fmt.Errorf("failed to render skill tool description: %w", renderErr) + } + + descBase := internal.SelectPrompt(internal.I18nPrompts{ + English: toolDescriptionBase, + Chinese: toolDescriptionBaseChinese, + }) + fullDesc = descBase + desc + } + + oneOf, err := s.buildParamsOneOf(ctx) + if err != nil { + return nil, fmt.Errorf("failed to build skill tool params: %w", err) + } + + return &schema.ToolInfo{ + Name: s.toolName, + Desc: fullDesc, + ParamsOneOf: oneOf, + }, nil +} + +type inputArguments struct { + Skill string `json:"skill"` +} + +func (s *typedSkillTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + args := &inputArguments{} + err := json.Unmarshal([]byte(argumentsInJSON), args) + if err != nil { + return "", fmt.Errorf("failed to unmarshal arguments: %w", err) + } + skill, err := s.b.Get(ctx, args.Skill) + if err != nil { + return "", fmt.Errorf("failed to get skill: %w", err) + } + + switch skill.Context { + case ContextModeForkWithContext: + return s.runAgentMode(ctx, skill, true, argumentsInJSON) + case ContextModeFork: + return s.runAgentMode(ctx, skill, false, argumentsInJSON) + default: + if skill.Model != "" { + s.setActiveModel(ctx, skill.Model) + } + return s.buildSkillResult(ctx, skill, argumentsInJSON) + } +} + +func (s *typedSkillTool[M]) setActiveModel(ctx context.Context, modelName string) { + _ = adk.SetRunLocalValue(ctx, activeModelKey, modelName) +} + +func defaultToolParams() map[string]*schema.ParameterInfo { + skillParamDesc := internal.SelectPrompt(internal.I18nPrompts{ + English: "The skill name (no arguments). E.g., \"pdf\" or \"xlsx\"", + Chinese: "Skill 名称（无需其他参数）。例如：\"pdf\" 或 \"xlsx\"", + }) + return map[string]*schema.ParameterInfo{ + "skill": { + Type: schema.String, + Desc: skillParamDesc, + Required: true, + }, + } +} + +func (s *typedSkillTool[M]) buildParamsOneOf(ctx context.Context) (*schema.ParamsOneOf, error) { + defaults := defaultToolParams() + if s.customToolParams == nil { + return schema.NewParamsOneOfByParams(defaults), nil + } + + params, err := s.customToolParams(ctx, defaults) + if err != nil { + return nil, err + } + if params == nil { + params = defaults + } + + if _, ok := params["skill"]; !ok { + params["skill"] = defaults["skill"] + } + + if p := params["skill"]; p != nil { + p.Required = true + } + + return schema.NewParamsOneOfByParams(params), nil +} + +func (s *typedSkillTool[M]) buildSkillResult(ctx context.Context, skill Skill, rawArguments string) (string, error) { + if s.buildContent == nil { + return s.defaultSkillContent(skill), nil + } + content, err := s.buildContent(ctx, skill, rawArguments) + if err != nil { + return "", fmt.Errorf("failed to build skill result: %w", err) + } + return content, nil +} + +func (s *typedSkillTool[M]) defaultSkillContent(skill Skill) string { + resultFmt := internal.SelectPrompt(internal.I18nPrompts{ + English: toolResult, + Chinese: toolResultChinese, + }) + contentFmt := internal.SelectPrompt(internal.I18nPrompts{ + English: userContent, + Chinese: userContentChinese, + }) + + return fmt.Sprintf(resultFmt, skill.Name) + fmt.Sprintf(contentFmt, skill.BaseDirectory, skill.Content) +} + +func (s *typedSkillTool[M]) runAgentMode(ctx context.Context, skill Skill, forkHistory bool, rawArguments string) (string, error) { + if s.agentHub == nil { + return "", fmt.Errorf("skill '%s' requires context:%s but AgentHub is not configured", skill.Name, skill.Context) + } + + opts := &TypedAgentHubOptions[M]{} + if skill.Model != "" { + if s.modelHub == nil { + return "", fmt.Errorf("skill '%s' requires model '%s' but ModelHub is not configured", skill.Name, skill.Model) + } + m, err := s.modelHub.Get(ctx, skill.Model) + if err != nil { + return "", fmt.Errorf("failed to get model '%s' from ModelHub: %w", skill.Model, err) + } + opts.Model = m + } + + agent, err := s.agentHub.Get(ctx, skill.Agent, opts) + if err != nil { + return "", fmt.Errorf("failed to get agent '%s' from AgentHub: %w", skill.Agent, err) + } + + var messages []M + skillContent, err := s.buildSkillResult(ctx, skill, rawArguments) + if err != nil { + return "", fmt.Errorf("failed to build skill result: %w", err) + } + + var history []M + var toolCallID string + if forkHistory { + history, err = s.getMessagesFromState(ctx) + if err != nil { + return "", fmt.Errorf("failed to get messages from state: %w", err) + } + toolCallID = compose.GetToolCallID(ctx) + } + + if s.buildForkMessages != nil { + messages, err = s.buildForkMessages(ctx, TypedSubAgentInput[M]{ + Skill: skill, + Mode: skill.Context, + RawArguments: rawArguments, + SkillContent: skillContent, + History: history, + ToolCallID: toolCallID, + }) + if err != nil { + return "", fmt.Errorf("failed to build fork messages: %w", err) + } + } else { + var zero M + if forkHistory { + var toolMsg M + switch any(zero).(type) { + case *schema.Message: + toolMsg = any(schema.ToolMessage(skillContent, toolCallID)).(M) + case *schema.AgenticMessage: + toolMsg = any(&schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: toolCallID, + Name: "", + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: skillContent}}, + }, + }), + }, + }).(M) + } + messages = append(history, toolMsg) + } else { + var userMsg M + switch any(zero).(type) { + case *schema.Message: + userMsg = any(schema.UserMessage(skillContent)).(M) + case *schema.AgenticMessage: + userMsg = any(schema.UserAgenticMessage(skillContent)).(M) + } + messages = []M{userMsg} + } + } + + input := &adk.TypedAgentInput[M]{ + Messages: messages, + EnableStreaming: false, + } + + iter := agent.Run(ctx, input) + + var msgList []M + var results []string + for { + event, ok := iter.Next() + if !ok { + break + } + + if event.Err != nil { + return "", fmt.Errorf("failed to run agent event: %w", event.Err) + } + + if event.Output == nil || event.Output.MessageOutput == nil { + continue + } + + msg, msgErr := event.Output.MessageOutput.GetMessage() + if msgErr != nil { + return "", fmt.Errorf("failed to get message from event: %w", msgErr) + } + + if !isNilMessage(msg) { + msgList = append(msgList, msg) + var content string + switch m := any(msg).(type) { + case *schema.Message: + content = m.Content + case *schema.AgenticMessage: + var parts []string + for _, block := range m.ContentBlocks { + if block == nil { + continue + } + if block.AssistantGenText != nil { + parts = append(parts, block.AssistantGenText.Text) + } + } + content = strings.Join(parts, "\n") + } + if content != "" { + results = append(results, content) + } + } + } + + if s.formatForkResult != nil { + out, err := s.formatForkResult(ctx, TypedSubAgentOutput[M]{ + Skill: skill, + Mode: skill.Context, + RawArguments: rawArguments, + Messages: msgList, + Results: results, + }) + if err != nil { + return "", fmt.Errorf("failed to format fork result: %w", err) + } + return out, nil + } + + resultFmt := internal.SelectPrompt(internal.I18nPrompts{ + English: subAgentResultFormat, + Chinese: subAgentResultFormatChinese, + }) + + return fmt.Sprintf(resultFmt, skill.Name, strings.Join(results, "\n")), nil +} + +func isNilMessage[M adk.MessageType](msg M) bool { + var zero M + return any(msg) == any(zero) +} + +func (s *typedSkillTool[M]) getMessagesFromState(ctx context.Context) ([]M, error) { + var messages []M + var zero M + switch any(zero).(type) { + case *schema.Message: + err := compose.ProcessState(ctx, func(_ context.Context, st *adk.State) error { + messages = make([]M, len(st.Messages)) + for i, m := range st.Messages { + messages[i] = any(m).(M) + } + return nil + }) + if err != nil { + return nil, fmt.Errorf("failed to process state: %w", err) + } + case *schema.AgenticMessage: + // Fork mode is not supported for AgenticMessage because the internal + // agent state type (agenticState) is unexported from the adk package, + // making it inaccessible via compose.ProcessState from middleware packages. + // Agent mode (the default) works normally for AgenticMessage. + return nil, fmt.Errorf("fork mode is not supported for AgenticMessage; use agent mode instead") + } + return messages, nil +} + +func renderToolDescription(matters []FrontMatter) (string, error) { + tpl, err := template.New("skills").Parse(toolDescriptionTemplate) + if err != nil { + return "", err + } + + var buf bytes.Buffer + err = tpl.Execute(&buf, descriptionTemplateHelper{Matters: matters}) + if err != nil { + return "", err + } + + return buf.String(), nil +} diff --git a/adk/middlewares/skill/skill_generic_test.go b/adk/middlewares/skill/skill_generic_test.go new file mode 100644 index 0000000..4e5a8e7 --- /dev/null +++ b/adk/middlewares/skill/skill_generic_test.go @@ -0,0 +1,466 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package skill + +import ( + "context" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// --- Generic mock types --- + +type mockGenericModel[M adk.MessageType] struct { + generateFunc func(ctx context.Context, input []M, opts ...model.Option) (M, error) +} + +func (m *mockGenericModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + if m.generateFunc != nil { + return m.generateFunc(ctx, input, opts...) + } + var zero M + return zero, nil +} + +func (m *mockGenericModel[M]) Stream(_ context.Context, _ []M, _ ...model.Option) (*schema.StreamReader[M], error) { + return nil, nil +} + +type mockGenericModelHub[M adk.MessageType] struct { + models map[string]model.BaseModel[M] +} + +func (h *mockGenericModelHub[M]) Get(_ context.Context, name string) (model.BaseModel[M], error) { + m, ok := h.models[name] + if !ok { + return nil, assert.AnError + } + return m, nil +} + +type mockGenericAgent[M adk.MessageType] struct { + events []*adk.TypedAgentEvent[M] + lastIn *adk.TypedAgentInput[M] +} + +func (a *mockGenericAgent[M]) Name(_ context.Context) string { return "mock-generic-agent" } +func (a *mockGenericAgent[M]) Description(_ context.Context) string { return "mock generic agent" } +func (a *mockGenericAgent[M]) Run(_ context.Context, in *adk.TypedAgentInput[M], _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.TypedAgentEvent[M]] { + a.lastIn = in + iter, gen := adk.NewAsyncIteratorPair[*adk.TypedAgentEvent[M]]() + go func() { + defer gen.Close() + for _, e := range a.events { + gen.Send(e) + } + }() + return iter +} + +type mockGenericAgentHub[M adk.MessageType] struct { + agent adk.TypedAgent[M] + lastOpts *TypedAgentHubOptions[M] +} + +func (h *mockGenericAgentHub[M]) Get(_ context.Context, _ string, opts *TypedAgentHubOptions[M]) (adk.TypedAgent[M], error) { + h.lastOpts = opts + return h.agent, nil +} + +// --- Helper to build an assistant event for each message type --- + +func assistantEvent[M adk.MessageType](text string) *adk.TypedAgentEvent[M] { + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.AssistantMessage(text, nil) + return any(&adk.TypedAgentEvent[*schema.Message]{ + Output: &adk.TypedAgentOutput[*schema.Message]{ + MessageOutput: &adk.TypedMessageVariant[*schema.Message]{ + Message: msg, + }, + }, + }).(*adk.TypedAgentEvent[M]) + case *schema.AgenticMessage: + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: text}), + }, + } + return any(&adk.TypedAgentEvent[*schema.AgenticMessage]{ + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: msg, + }, + }, + }).(*adk.TypedAgentEvent[M]) + } + panic("unreachable") +} + +// --- Part 1: WrapModel test --- + +func testWrapModel[M adk.MessageType](t *testing.T) { + ctx := context.Background() + mockModel := &mockGenericModel[M]{} + hub := &mockGenericModelHub[M]{ + models: map[string]model.BaseModel[M]{ + "test-model": mockModel, + }, + } + + mw, err := NewTyped(ctx, &TypedConfig[M]{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: hub, + }) + require.NoError(t, err) + require.NotNil(t, mw) + + t.Run("nil ModelHub keeps base model", func(t *testing.T) { + handler, err := NewTyped(ctx, &TypedConfig[M]{ + Backend: &inMemoryBackend{m: []Skill{}}, + }) + require.NoError(t, err) + h := handler.(*typedSkillHandler[M]) + base := &mockGenericModel[M]{} + got, err := h.WrapModel(ctx, base, &adk.TypedModelContext[M]{}) + require.NoError(t, err) + assert.Equal(t, base, got) + }) +} + +// --- Part 2: Agent mode tests --- + +func testAgentMode[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + t.Run("successful agent run", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + assistantEvent[M]("agent answer"), + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "test-skill", Context: ContextModeFork}, + Content: "skill content", + BaseDirectory: "/skills/test", + }, + }}, + toolName: "skill", + agentHub: hub, + } + + result, err := st.InvokableRun(ctx, `{"skill": "test-skill"}`) + require.NoError(t, err) + assert.Contains(t, result, "test-skill") + assert.Contains(t, result, "agent answer") + + // Verify agent received a user message (fork mode, no history). + require.NotNil(t, agent.lastIn) + require.Len(t, agent.lastIn.Messages, 1) + }) + + t.Run("fork mode constructs user message", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + assistantEvent[M]("ok"), + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, + Content: "c1", + BaseDirectory: "/d", + }, + }}, + toolName: "skill", + agentHub: hub, + } + + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.NoError(t, err) + require.NotNil(t, agent.lastIn) + require.Len(t, agent.lastIn.Messages, 1) + + // Verify the message is a user-role message. + msg := agent.lastIn.Messages[0] + var zero M + switch any(zero).(type) { + case *schema.Message: + m := any(msg).(*schema.Message) + assert.Equal(t, schema.User, m.Role) + case *schema.AgenticMessage: + m := any(msg).(*schema.AgenticMessage) + assert.Equal(t, schema.AgenticRoleTypeUser, m.Role) + } + }) + + t.Run("custom buildForkMessages called", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + assistantEvent[M]("ok"), + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + var captured TypedSubAgentInput[M] + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, + Content: "c1", + BaseDirectory: "/d", + }, + }}, + toolName: "skill", + agentHub: hub, + buildForkMessages: func(_ context.Context, in TypedSubAgentInput[M]) ([]M, error) { + captured = in + // Build a simple user message. + var zero M + switch any(zero).(type) { + case *schema.Message: + return []M{any(schema.UserMessage("custom")).(M)}, nil + case *schema.AgenticMessage: + return []M{any(schema.UserAgenticMessage("custom")).(M)}, nil + } + return nil, nil + }, + } + + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.NoError(t, err) + assert.Equal(t, "s1", captured.Skill.Name) + assert.Equal(t, ContextModeFork, captured.Mode) + }) + + t.Run("custom formatForkResult called", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + assistantEvent[M]("p1"), + assistantEvent[M]("p2"), + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, + Content: "c1", + BaseDirectory: "/d", + }, + }}, + toolName: "skill", + agentHub: hub, + formatForkResult: func(_ context.Context, in TypedSubAgentOutput[M]) (string, error) { + assert.Equal(t, ContextModeFork, in.Mode) + assert.Equal(t, []string{"p1", "p2"}, in.Results) + assert.Len(t, in.Messages, 2) + return "formatted:" + strings.Join(in.Results, ","), nil + }, + } + + result, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.NoError(t, err) + assert.Equal(t, "formatted:p1,p2", result) + }) + + t.Run("agent error event propagates", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + {Err: assert.AnError}, + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, + Content: "c1", + BaseDirectory: "/d", + }, + }}, + toolName: "skill", + agentHub: hub, + } + + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to run agent event") + }) + + t.Run("multiple events concatenated", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + assistantEvent[M]("part1"), + {Output: nil}, // skipped + assistantEvent[M]("part2"), + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, + Content: "c1", + BaseDirectory: "/d", + }, + }}, + toolName: "skill", + agentHub: hub, + } + + result, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.NoError(t, err) + assert.Contains(t, result, "part1") + assert.Contains(t, result, "part2") + }) + + t.Run("model passed to agent hub", func(t *testing.T) { + mdl := &mockGenericModel[M]{} + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{ + assistantEvent[M]("ok"), + }, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork, Model: "m1"}, + Content: "c1", + BaseDirectory: "/d", + }, + }}, + toolName: "skill", + agentHub: hub, + modelHub: &mockGenericModelHub[M]{ + models: map[string]model.BaseModel[M]{"m1": mdl}, + }, + } + + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.NoError(t, err) + require.NotNil(t, hub.lastOpts) + assert.Equal(t, mdl, hub.lastOpts.Model) + }) + + t.Run("no AgentHub returns error", func(t *testing.T) { + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1"}, + }}, + toolName: "skill", + } + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.Error(t, err) + assert.Contains(t, err.Error(), "AgentHub is not configured") + }) + + t.Run("no ModelHub with model returns error", func(t *testing.T) { + agent := &mockGenericAgent[M]{ + events: []*adk.TypedAgentEvent[M]{assistantEvent[M]("ok")}, + } + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork, Model: "m1"}, Content: "c1"}, + }}, + toolName: "skill", + agentHub: hub, + } + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.Error(t, err) + assert.Contains(t, err.Error(), "ModelHub is not configured") + }) + + t.Run("empty content events still produce result", func(t *testing.T) { + var emptyEvent *adk.TypedAgentEvent[M] + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.AssistantMessage("", nil) + emptyEvent = any(&adk.TypedAgentEvent[*schema.Message]{ + Output: &adk.TypedAgentOutput[*schema.Message]{ + MessageOutput: &adk.TypedMessageVariant[*schema.Message]{ + Message: msg, + }, + }, + }).(*adk.TypedAgentEvent[M]) + case *schema.AgenticMessage: + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{}, + } + emptyEvent = any(&adk.TypedAgentEvent[*schema.AgenticMessage]{ + Output: &adk.TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &adk.TypedMessageVariant[*schema.AgenticMessage]{ + Message: msg, + }, + }, + }).(*adk.TypedAgentEvent[M]) + } + + agent := &mockGenericAgent[M]{events: []*adk.TypedAgentEvent[M]{emptyEvent}} + hub := &mockGenericAgentHub[M]{agent: agent} + + st := &typedSkillTool[M]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + } + + result, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + require.NoError(t, err) + assert.Contains(t, result, "s1") + }) +} + +// --- Top-level test --- + +func TestSkillGeneric(t *testing.T) { + t.Run("Message", func(t *testing.T) { + t.Run("WrapModel", testWrapModel[*schema.Message]) + t.Run("AgentMode", testAgentMode[*schema.Message]) + }) + t.Run("AgenticMessage", func(t *testing.T) { + t.Run("WrapModel", testWrapModel[*schema.AgenticMessage]) + t.Run("AgentMode", testAgentMode[*schema.AgenticMessage]) + }) +} diff --git a/adk/middlewares/skill/skill_test.go b/adk/middlewares/skill/skill_test.go new file mode 100644 index 0000000..3cc536a --- /dev/null +++ b/adk/middlewares/skill/skill_test.go @@ -0,0 +1,1353 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package skill + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +type inMemoryBackend struct { + m []Skill +} + +func (i *inMemoryBackend) List(ctx context.Context) ([]FrontMatter, error) { + matters := make([]FrontMatter, 0, len(i.m)) + for _, skill := range i.m { + matters = append(matters, skill.FrontMatter) + } + return matters, nil +} + +func (i *inMemoryBackend) Get(ctx context.Context, name string) (Skill, error) { + for _, skill := range i.m { + if skill.Name == name { + return skill, nil + } + } + return Skill{}, errors.New("skill not found") +} + +func TestTool(t *testing.T) { + backend := &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{ + Name: "name1", + Description: "desc1", + }, + Content: "content1", + BaseDirectory: "basedir1", + }, + { + FrontMatter: FrontMatter{ + Name: "name2", + Description: "desc2", + }, + Content: "content2", + BaseDirectory: "basedir2", + }, + }} + + ctx := context.Background() + m, err := New(ctx, &Config{Backend: backend}) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 1) + + to := m.AdditionalTools[0].(tool.InvokableTool) + + info, err := to.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "skill", info.Name) + desc := strings.TrimPrefix(info.Desc, toolDescriptionBase) + assert.Equal(t, ` + + + +name1 + + +desc1 + + + + +name2 + + +desc2 + + + +`, desc) + + result, err := to.InvokableRun(ctx, `{"skill": "name1"}`) + assert.NoError(t, err) + assert.Equal(t, `Launching skill: name1 +Base directory for this skill: basedir1 + +content1`, result) + + // chinese + internal.SetLanguage(internal.LanguageChinese) + defer internal.SetLanguage(internal.LanguageEnglish) + m, err = New(ctx, &Config{Backend: backend}) + assert.NoError(t, err) + assert.Len(t, m.AdditionalTools, 1) + + to = m.AdditionalTools[0].(tool.InvokableTool) + + info, err = to.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "skill", info.Name) + desc = strings.TrimPrefix(info.Desc, toolDescriptionBaseChinese) + assert.Equal(t, ` + + + +name1 + + +desc1 + + + + +name2 + + +desc2 + + + +`, desc) + + result, err = to.InvokableRun(ctx, `{"skill": "name1"}`) + assert.NoError(t, err) + assert.Equal(t, `正在启动 Skill：name1 +此 Skill 的目录：basedir1 + +content1`, result) +} + +func TestSkillToolName(t *testing.T) { + ctx := context.Background() + + // default + m, err := New(ctx, &Config{Backend: &inMemoryBackend{m: []Skill{}}}) + assert.NoError(t, err) + // instruction + assert.Contains(t, m.AdditionalInstruction, "'skill'") + // tool name + info, err := m.AdditionalTools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "skill", info.Name) + + // customized + name := "load_skill" + m, err = New(ctx, &Config{Backend: &inMemoryBackend{m: []Skill{}}, SkillToolName: &name}) + assert.NoError(t, err) + assert.Contains(t, m.AdditionalInstruction, "'load_skill'") + info, err = m.AdditionalTools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "load_skill", info.Name) +} + +func TestBuildParamsOneOf_CustomParams(t *testing.T) { + internal.SetLanguage(internal.LanguageEnglish) + ctx := context.Background() + + st := &typedSkillTool[*schema.Message]{ + customToolParams: func(context.Context, map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) { + return map[string]*schema.ParameterInfo{ + "foo": { + Type: schema.String, + Desc: "foo desc", + Required: true, + }, + "bar": { + Type: schema.Integer, + Desc: "bar desc", + Required: false, + }, + "skill": { + Type: schema.String, + Desc: "custom skill desc", + Required: false, + }, + }, nil + }, + } + + oneOf, err := st.buildParamsOneOf(ctx) + require.NoError(t, err) + js, err := oneOf.ToJSONSchema() + require.NoError(t, err) + require.NotNil(t, js) + require.NotNil(t, js.Properties) + + skillSchema, ok := js.Properties.Get("skill") + require.True(t, ok) + require.NotNil(t, skillSchema) + assert.Equal(t, string(schema.String), skillSchema.Type) + assert.Equal(t, "custom skill desc", skillSchema.Description) + + _, ok = js.Properties.Get("foo") + assert.True(t, ok) + _, ok = js.Properties.Get("bar") + assert.True(t, ok) + + assert.Equal(t, []string{"foo", "skill"}, js.Required) +} + +func TestBuildParamsOneOf_CustomParamsNilFallsBackToDefault(t *testing.T) { + ctx := context.Background() + st := &typedSkillTool[*schema.Message]{ + customToolParams: func(context.Context, map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) { + return nil, nil + }, + } + + oneOf, err := st.buildParamsOneOf(ctx) + require.NoError(t, err) + js, err := oneOf.ToJSONSchema() + require.NoError(t, err) + require.NotNil(t, js) + require.NotNil(t, js.Properties) + _, ok := js.Properties.Get("skill") + require.True(t, ok) + assert.Contains(t, js.Required, "skill") +} + +// --- Mock types for NewMiddleware tests --- + +type mockModel struct { + model.BaseModel[*schema.Message] + name string +} + +type mockModelHub struct { + models map[string]model.BaseModel[*schema.Message] +} + +func (h *mockModelHub) Get(_ context.Context, name string) (model.BaseModel[*schema.Message], error) { + m, ok := h.models[name] + if !ok { + return nil, fmt.Errorf("model not found: %s", name) + } + return m, nil +} + +type fakeToolCallingModel struct { + id string + calls int +} + +func (m *fakeToolCallingModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.calls++ + return schema.AssistantMessage(m.id, nil), nil +} + +func (m *fakeToolCallingModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, io.EOF +} + +func (m *fakeToolCallingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +type runLocalSetterHandler struct { + *adk.BaseChatModelAgentMiddleware + key string + val any +} + +func (h *runLocalSetterHandler) BeforeModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, _ *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) { + if err := adk.SetRunLocalValue(ctx, h.key, h.val); err != nil { + return nil, nil, err + } + return ctx, state, nil +} + +type stateMessagesCaptureHandler struct { + *adk.BaseChatModelAgentMiddleware + st *typedSkillTool[*schema.Message] + captured []adk.Message +} + +func (h *stateMessagesCaptureHandler) AfterModelRewriteState(ctx context.Context, state *adk.ChatModelAgentState, _ *adk.ModelContext) (context.Context, *adk.ChatModelAgentState, error) { + msgs, err := h.st.getMessagesFromState(ctx) + if err != nil { + return nil, nil, err + } + h.captured = msgs + return ctx, state, nil +} + +type mockAgent struct { + events []*adk.AgentEvent + lastIn *adk.AgentInput +} + +func (a *mockAgent) Name(_ context.Context) string { return "mock-agent" } +func (a *mockAgent) Description(_ context.Context) string { return "mock agent for testing" } +func (a *mockAgent) Run(_ context.Context, in *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + a.lastIn = in + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + go func() { + defer gen.Close() + for _, e := range a.events { + gen.Send(e) + } + }() + return iter +} + +type mockAgentHub struct { + agents map[string]adk.Agent + lastOpts *AgentHubOptions + defaultAgent adk.Agent +} + +func (h *mockAgentHub) Get(_ context.Context, name string, opts *AgentHubOptions) (adk.Agent, error) { + h.lastOpts = opts + if name == "" && h.defaultAgent != nil { + return h.defaultAgent, nil + } + a, ok := h.agents[name] + if !ok { + return nil, fmt.Errorf("agent not found: %s", name) + } + return a, nil +} + +type errorBackend struct { + listErr error + getErr error +} + +func (b *errorBackend) List(_ context.Context) ([]FrontMatter, error) { + return nil, b.listErr +} +func (b *errorBackend) Get(_ context.Context, _ string) (Skill, error) { + return Skill{}, b.getErr +} + +// --- NewMiddleware tests --- + +func TestNewMiddleware(t *testing.T) { + ctx := context.Background() + + t.Run("nil config returns error", func(t *testing.T) { + handler, err := NewMiddleware(ctx, nil) + assert.Nil(t, handler) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config is required") + }) + + t.Run("nil backend returns error", func(t *testing.T) { + handler, err := NewMiddleware(ctx, &Config{}) + assert.Nil(t, handler) + assert.Error(t, err) + assert.Contains(t, err.Error(), "backend is required") + }) + + t.Run("custom tool params error surfaces in Info", func(t *testing.T) { + backend := &inMemoryBackend{m: []Skill{}} + handler, err := NewMiddleware(ctx, &Config{ + Backend: backend, + CustomToolParams: func(context.Context, map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) { + return nil, errors.New("bad params") + }, + }) + require.NoError(t, err) + h := handler.(*typedSkillHandler[*schema.Message]) + _, err = h.tool.Info(ctx) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to build skill tool params") + assert.Contains(t, err.Error(), "bad params") + }) + + t.Run("valid config succeeds", func(t *testing.T) { + backend := &inMemoryBackend{m: []Skill{}} + handler, err := NewMiddleware(ctx, &Config{Backend: backend}) + assert.NoError(t, err) + assert.NotNil(t, handler) + }) + + t.Run("custom tool name", func(t *testing.T) { + backend := &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Description: "d1"}, Content: "c1"}, + }} + name := "load_skill" + handler, err := NewMiddleware(ctx, &Config{Backend: backend, SkillToolName: &name}) + require.NoError(t, err) + + h := handler.(*typedSkillHandler[*schema.Message]) + assert.Contains(t, h.instruction, "'load_skill'") + assert.Equal(t, "load_skill", h.tool.toolName) + }) + + t.Run("custom system prompt", func(t *testing.T) { + backend := &inMemoryBackend{m: []Skill{}} + handler, err := NewMiddleware(ctx, &Config{ + Backend: backend, + CustomSystemPrompt: func(_ context.Context, toolName string) string { + return "custom prompt for " + toolName + }, + }) + require.NoError(t, err) + + h := handler.(*typedSkillHandler[*schema.Message]) + assert.Equal(t, "custom prompt for skill", h.instruction) + }) + + t.Run("custom tool description", func(t *testing.T) { + backend := &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Description: "d1"}, Content: "c1"}, + }} + handler, err := NewMiddleware(ctx, &Config{ + Backend: backend, + CustomToolDescription: func(_ context.Context, skills []FrontMatter) string { + return fmt.Sprintf("custom desc with %d skills", len(skills)) + }, + }) + require.NoError(t, err) + + h := handler.(*typedSkillHandler[*schema.Message]) + info, err := h.tool.Info(ctx) + require.NoError(t, err) + assert.Equal(t, "custom desc with 1 skills", info.Desc) + }) +} + +func TestBeforeAgent(t *testing.T) { + ctx := context.Background() + backend := &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Description: "d1"}, Content: "c1"}, + }} + handler, err := NewMiddleware(ctx, &Config{Backend: backend}) + require.NoError(t, err) + + runCtx := &adk.ChatModelAgentContext{ + Instruction: "base instruction", + Tools: []tool.BaseTool{}, + } + _, newRunCtx, err := handler.BeforeAgent(ctx, runCtx) + assert.NoError(t, err) + assert.Contains(t, newRunCtx.Instruction, "base instruction") + assert.Contains(t, newRunCtx.Instruction, "Skills System") + assert.Len(t, newRunCtx.Tools, 1) + + // verify the added tool is the skill tool + info, err := newRunCtx.Tools[0].Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "skill", info.Name) +} + +func TestWrapModel_SwitchesModelWhenRunLocalIsSet(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + other := &fakeToolCallingModel{id: "other"} + + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{"other": other}}, + }) + require.NoError(t, err) + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "t", + Description: "t", + Model: base, + Handlers: []adk.ChatModelAgentMiddleware{ + &runLocalSetterHandler{key: activeModelKey, val: "other"}, + handler, + }, + MaxIterations: 1, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}) + var last string + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + msg, err := ev.Output.MessageOutput.GetMessage() + require.NoError(t, err) + if msg != nil { + last = msg.Content + } + } + + assert.Equal(t, "other", last) + assert.Equal(t, 0, base.calls) + assert.Equal(t, 1, other.calls) +} + +func TestWrapModel_OutsideAgentContextReturnsError(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + other := &fakeToolCallingModel{id: "other"} + + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{"other": other}}, + }) + require.NoError(t, err) + + h := handler.(*typedSkillHandler[*schema.Message]) + _, err = h.WrapModel(ctx, base, &adk.ModelContext{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to get active model from run local value") +} + +func TestWrapModel_IgnoresNonStringRunLocalValue(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + other := &fakeToolCallingModel{id: "other"} + + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{"other": other}}, + }) + require.NoError(t, err) + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "t", + Description: "t", + Model: base, + Handlers: []adk.ChatModelAgentMiddleware{ + &runLocalSetterHandler{key: activeModelKey, val: 123}, + handler, + }, + MaxIterations: 1, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}) + var last string + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + msg, err := ev.Output.MessageOutput.GetMessage() + require.NoError(t, err) + if msg != nil { + last = msg.Content + } + } + + assert.Equal(t, "base", last) + assert.Equal(t, 1, base.calls) + assert.Equal(t, 0, other.calls) +} + +func TestWrapModel_ModelHubGetError(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{}}, + }) + require.NoError(t, err) + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "t", + Description: "t", + Model: base, + Handlers: []adk.ChatModelAgentMiddleware{ + &runLocalSetterHandler{key: activeModelKey, val: "missing"}, + handler, + }, + MaxIterations: 1, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}) + var gotErr error + for { + ev, ok := iter.Next() + if !ok { + break + } + if ev.Err != nil { + gotErr = ev.Err + break + } + } + require.Error(t, gotErr) + assert.Contains(t, gotErr.Error(), "failed to get model") + assert.Equal(t, 0, base.calls) +} + +func TestWrapModel_ModelHubNilKeepsBase(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + }) + require.NoError(t, err) + + h := handler.(*typedSkillHandler[*schema.Message]) + m, err := h.WrapModel(ctx, base, &adk.ModelContext{}) + require.NoError(t, err) + assert.Equal(t, base, m) +} + +func TestWrapModel_RunLocalNotFoundKeepsBase(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + other := &fakeToolCallingModel{id: "other"} + + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{"other": other}}, + }) + require.NoError(t, err) + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "t", + Description: "t", + Model: base, + Handlers: []adk.ChatModelAgentMiddleware{ + handler, + }, + MaxIterations: 1, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}) + var last string + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + msg, err := ev.Output.MessageOutput.GetMessage() + require.NoError(t, err) + if msg != nil { + last = msg.Content + } + } + + assert.Equal(t, "base", last) + assert.Equal(t, 1, base.calls) + assert.Equal(t, 0, other.calls) +} + +func TestWrapModel_IgnoresEmptyStringRunLocalValue(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + other := &fakeToolCallingModel{id: "other"} + + handler, err := NewMiddleware(ctx, &Config{ + Backend: &inMemoryBackend{m: []Skill{}}, + ModelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{"other": other}}, + }) + require.NoError(t, err) + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "t", + Description: "t", + Model: base, + Handlers: []adk.ChatModelAgentMiddleware{ + &runLocalSetterHandler{key: activeModelKey, val: ""}, + handler, + }, + MaxIterations: 1, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}) + var last string + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + msg, err := ev.Output.MessageOutput.GetMessage() + require.NoError(t, err) + if msg != nil { + last = msg.Content + } + } + + assert.Equal(t, "base", last) + assert.Equal(t, 1, base.calls) + assert.Equal(t, 0, other.calls) +} + +func TestGetMessagesFromState_InAgentContext(t *testing.T) { + ctx := context.Background() + + base := &fakeToolCallingModel{id: "base"} + st := &typedSkillTool[*schema.Message]{} + capture := &stateMessagesCaptureHandler{st: st} + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "t", + Description: "t", + Model: base, + Handlers: []adk.ChatModelAgentMiddleware{ + capture, + }, + MaxIterations: 1, + }) + require.NoError(t, err) + + iter := agent.Run(ctx, &adk.AgentInput{Messages: []adk.Message{schema.UserMessage("hi")}}) + for { + ev, ok := iter.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + require.NotNil(t, capture.captured) + require.NotEmpty(t, capture.captured) +} + +func TestSkillToolInfo(t *testing.T) { + ctx := context.Background() + + t.Run("list error propagates", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &errorBackend{listErr: errors.New("list failed")}, + toolName: "skill", + } + info, err := st.Info(ctx) + assert.Nil(t, info) + assert.Error(t, err) + assert.Contains(t, err.Error(), "list failed") + }) + + t.Run("description contains all skills", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "alpha", Description: "desc-alpha"}}, + {FrontMatter: FrontMatter{Name: "beta", Description: "desc-beta"}}, + }}, + toolName: "skill", + } + info, err := st.Info(ctx) + require.NoError(t, err) + assert.Contains(t, info.Desc, "alpha") + assert.Contains(t, info.Desc, "desc-alpha") + assert.Contains(t, info.Desc, "beta") + assert.Contains(t, info.Desc, "desc-beta") + }) + + t.Run("custom tool params is used", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "alpha", Description: "desc-alpha"}}, + }}, + toolName: "skill", + customToolParams: func(_ context.Context, _ map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) { + return map[string]*schema.ParameterInfo{ + "foo": {Type: schema.String, Desc: "foo-desc", Required: true}, + "skill": {Type: schema.String, Desc: "custom-skill-desc", Required: false}, + }, nil + }, + } + info, err := st.Info(ctx) + require.NoError(t, err) + js, err := info.ParamsOneOf.ToJSONSchema() + require.NoError(t, err) + _, ok := js.Properties.Get("foo") + require.True(t, ok) + v, ok := js.Properties.Get("skill") + require.True(t, ok) + assert.Equal(t, "custom-skill-desc", v.Description) + assert.Contains(t, js.Required, "skill") + assert.Contains(t, js.Required, "foo") + }) +} + +func TestInvokableRun_InlineMode(t *testing.T) { + ctx := context.Background() + + t.Run("invalid json returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{}}, + toolName: "skill", + } + _, err := st.InvokableRun(ctx, "not json") + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to unmarshal") + }) + + t.Run("skill not found returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{}}, + toolName: "skill", + } + _, err := st.InvokableRun(ctx, `{"skill": "nonexistent"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to get skill") + }) + + t.Run("inline mode returns skill content", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "pdf", Description: "PDF processing"}, + Content: "Process PDF files here", + BaseDirectory: "/skills/pdf", + }, + }}, + toolName: "skill", + } + result, err := st.InvokableRun(ctx, `{"skill": "pdf"}`) + assert.NoError(t, err) + assert.Contains(t, result, "pdf") + assert.Contains(t, result, "/skills/pdf") + assert.Contains(t, result, "Process PDF files here") + }) + + t.Run("inline mode with model triggers setActiveModel", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "pdf", Description: "PDF processing", Model: "m1"}, + Content: "Process PDF files here", + BaseDirectory: "/skills/pdf", + }, + }}, + toolName: "skill", + } + result, err := st.InvokableRun(ctx, `{"skill": "pdf"}`) + assert.NoError(t, err) + assert.Contains(t, result, "pdf") + }) + + t.Run("custom skill content is used", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "pdf", Description: "PDF processing"}, + Content: "Process PDF files here", + BaseDirectory: "/skills/pdf", + }, + }}, + toolName: "skill", + buildContent: func(_ context.Context, _ Skill, rawArgs string) (string, error) { + var raw map[string]any + require.NoError(t, json.Unmarshal([]byte(rawArgs), &raw)) + assert.Equal(t, "pdf", raw["skill"]) + return "custom-content", nil + }, + } + result, err := st.InvokableRun(ctx, `{"skill":"pdf"}`) + assert.NoError(t, err) + assert.Equal(t, "custom-content", result) + }) + + t.Run("custom tool params with decoder is used", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "pdf", Description: "PDF processing"}, + Content: "Process PDF files here", + BaseDirectory: "/skills/pdf", + }, + }}, + toolName: "skill", + customToolParams: func(_ context.Context, _ map[string]*schema.ParameterInfo) (map[string]*schema.ParameterInfo, error) { + return map[string]*schema.ParameterInfo{ + "skill": {Type: schema.String, Desc: "custom", Required: false}, + "task": {Type: schema.String, Desc: "custom", Required: false}, + "x": {Type: schema.Integer, Desc: "custom", Required: false}, + }, nil + }, + buildContent: func(_ context.Context, _ Skill, rawArgs string) (string, error) { + var raw struct { + Skill string `json:"skill"` + Task string `json:"task"` + X int `json:"x"` + } + require.NoError(t, json.Unmarshal([]byte(rawArgs), &raw)) + assert.Equal(t, "pdf", raw.Skill) + assert.Equal(t, "t", raw.Task) + assert.Equal(t, 1, raw.X) + return "decoded", nil + }, + } + result, err := st.InvokableRun(ctx, `{"skill":"pdf","task":"t","x":1}`) + assert.NoError(t, err) + assert.Equal(t, "decoded", result) + }) + + t.Run("custom skill content returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "pdf", Description: "PDF processing"}, + Content: "Process PDF files here", + BaseDirectory: "/skills/pdf", + }, + }}, + toolName: "skill", + buildContent: func(context.Context, Skill, string) (string, error) { + return "", errors.New("boom") + }, + } + _, err := st.InvokableRun(ctx, `{"skill":"pdf"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to build skill result") + }) +} + +func TestInvokableRun_AgentMode(t *testing.T) { + ctx := context.Background() + + t.Run("fork mode without AgentHub returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1"}, + }}, + toolName: "skill", + } + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "AgentHub is not configured") + }) + + t.Run("fork_with_context mode without AgentHub returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeForkWithContext}, Content: "c1"}, + }}, + toolName: "skill", + } + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "AgentHub is not configured") + }) + + t.Run("fork_with_context mode without state returns error", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("agent response", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeForkWithContext}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + } + + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to get messages from state") + }) + + t.Run("model specified without ModelHub returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork, Model: "gpt-4"}, Content: "c1"}, + }}, + toolName: "skill", + agentHub: &mockAgentHub{}, + } + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "ModelHub is not configured") + }) + + t.Run("model not found in ModelHub returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork, Model: "gpt-4"}, Content: "c1"}, + }}, + toolName: "skill", + agentHub: &mockAgentHub{}, + modelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{}}, + } + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to get model") + }) + + t.Run("agent not found in AgentHub returns error", func(t *testing.T) { + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork, Agent: "nonexistent"}, Content: "c1"}, + }}, + toolName: "skill", + agentHub: &mockAgentHub{agents: map[string]adk.Agent{}}, + } + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to get agent") + }) + + t.Run("fork mode runs agent and returns result", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("agent response", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "test-skill", Context: ContextModeFork}, + Content: "skill content", + BaseDirectory: "/skills/test", + }, + }}, + toolName: "skill", + agentHub: hub, + } + + result, err := st.InvokableRun(ctx, `{"skill": "test-skill"}`) + assert.NoError(t, err) + assert.Contains(t, result, "test-skill") + assert.Contains(t, result, "agent response") + assert.Contains(t, result, "completed") + // verify no model was passed in opts + assert.NotNil(t, hub.lastOpts) + assert.Nil(t, hub.lastOpts.Model) + }) + + t.Run("fork mode with model passes model to AgentHub", func(t *testing.T) { + m := &mockModel{name: "test-model"} + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("response", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + { + FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork, Model: "test-model"}, + Content: "c1", + BaseDirectory: "/skills/s1", + }, + }}, + toolName: "skill", + agentHub: hub, + modelHub: &mockModelHub{models: map[string]model.BaseModel[*schema.Message]{"test-model": m}}, + } + + result, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.NoError(t, err) + assert.Contains(t, result, "s1") + // verify model was passed + assert.NotNil(t, hub.lastOpts) + assert.Equal(t, m, hub.lastOpts.Model) + }) + + t.Run("agent returns multiple events", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("part1", nil), + }, + }, + }, + {Output: nil}, // nil output should be skipped + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("part2", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + } + + result, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.NoError(t, err) + assert.Contains(t, result, "part1") + assert.Contains(t, result, "part2") + }) + + t.Run("agent returns empty content events", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + } + + result, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.NoError(t, err) + // result should contain skill name but no extra content + assert.Contains(t, result, "s1") + }) + + t.Run("agent event error returns error", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + {Err: errors.New("boom")}, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + } + + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to run agent event") + }) + + t.Run("custom fork messages is used", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("ok", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + buildForkMessages: func(_ context.Context, in SubAgentInput) ([]adk.Message, error) { + assert.Equal(t, ContextModeFork, in.Mode) + assert.Equal(t, "s1", in.Skill.Name) + return []adk.Message{schema.UserMessage("custom")}, nil + }, + } + + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.NoError(t, err) + require.NotNil(t, agent.lastIn) + require.Len(t, agent.lastIn.Messages, 1) + msg := agent.lastIn.Messages[0] + assert.Equal(t, "custom", msg.Content) + }) + + t.Run("custom fork messages returns error", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("ok", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + buildForkMessages: func(context.Context, SubAgentInput) ([]adk.Message, error) { + return nil, errors.New("build msg fail") + }, + } + + _, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to build fork messages") + }) + + t.Run("custom fork result prompts from results is used", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("p1", nil), + }, + }, + }, + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("p2", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + formatForkResult: func(_ context.Context, in SubAgentOutput) (string, error) { + assert.Equal(t, ContextModeFork, in.Mode) + assert.Equal(t, []string{"p1", "p2"}, in.Results) + return "E:" + strings.Join(in.Results, ","), nil + }, + } + + result, err := st.InvokableRun(ctx, `{"skill": "s1"}`) + assert.NoError(t, err) + assert.Equal(t, "E:p1,p2", result) + }) + + t.Run("custom fork result prompts returns error", func(t *testing.T) { + agent := &mockAgent{ + events: []*adk.AgentEvent{ + { + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Message: schema.AssistantMessage("p1", nil), + }, + }, + }, + }, + } + hub := &mockAgentHub{defaultAgent: agent} + + st := &typedSkillTool[*schema.Message]{ + b: &inMemoryBackend{m: []Skill{ + {FrontMatter: FrontMatter{Name: "s1", Context: ContextModeFork}, Content: "c1", BaseDirectory: "/d"}, + }}, + toolName: "skill", + agentHub: hub, + formatForkResult: func(context.Context, SubAgentOutput) (string, error) { + return "", errors.New("format fail") + }, + } + + _, err := st.InvokableRun(ctx, `{"skill":"s1"}`) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to format fork result") + assert.Contains(t, err.Error(), "format fail") + }) +} + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + mw, err := NewTyped(ctx, &TypedConfig[*schema.AgenticMessage]{ + Backend: &inMemoryBackend{m: []Skill{}}, + }) + assert.NoError(t, err) + assert.NotNil(t, mw) + + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} diff --git a/adk/middlewares/summarization/consts.go b/adk/middlewares/summarization/consts.go new file mode 100644 index 0000000..ae9e117 --- /dev/null +++ b/adk/middlewares/summarization/consts.go @@ -0,0 +1,39 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +type ActionType string + +const ( + ActionTypeBeforeSummarize ActionType = "before_summarize" + ActionTypeAfterSummarize ActionType = "after_summarize" + ActionTypeGenerateSummary ActionType = "generate_summary" +) + +const ( + extraKeyContentType = "_eino_summarization_content_type" + preserveUserMsgsMaxTokens = 30000 +) + +type summarizationContentType string + +const ( + contentTypeSummary summarizationContentType = "summary" + contentTypeSkills summarizationContentType = "skills" +) + +type ctxKeyModelInput struct{} diff --git a/adk/middlewares/summarization/customized_action.go b/adk/middlewares/summarization/customized_action.go new file mode 100644 index 0000000..61000f5 --- /dev/null +++ b/adk/middlewares/summarization/customized_action.go @@ -0,0 +1,101 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +import ( + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// TypedCustomizedAction is the generic customized action for summarization events. +type TypedCustomizedAction[M adk.MessageType] struct { + // Type is the action type. + Type ActionType `json:"type"` + + // Before is set when Type is ActionTypeBeforeSummarize. + // Emitted after trigger condition is met, before calling model to generate summary. + Before *TypedBeforeSummarizeAction[M] `json:"before,omitempty"` + + // After is set when Type is ActionTypeAfterSummarize. + // Emitted after summarization. + After *TypedAfterSummarizeAction[M] `json:"after,omitempty"` + + // GenerateSummary is set when Type is ActionTypeGenerateSummary. + // Emitted on each summary generation attempt, including retries and failovers. + GenerateSummary *TypedGenerateSummaryAction[M] `json:"generate_summary,omitempty"` +} + +// CustomizedAction is the default action type using *schema.Message. +type CustomizedAction = TypedCustomizedAction[*schema.Message] + +// TypedBeforeSummarizeAction contains the state messages before summarization. +type TypedBeforeSummarizeAction[M adk.MessageType] struct { + // Messages is the original state messages before summarization. + Messages []M `json:"messages,omitempty"` +} + +// BeforeSummarizeAction is the default type using *schema.Message. +type BeforeSummarizeAction = TypedBeforeSummarizeAction[*schema.Message] + +// TypedAfterSummarizeAction contains the state messages after summarization. +type TypedAfterSummarizeAction[M adk.MessageType] struct { + // Messages is the final state messages after summarization. + Messages []M `json:"messages,omitempty"` +} + +// AfterSummarizeAction is the default type using *schema.Message. +type AfterSummarizeAction = TypedAfterSummarizeAction[*schema.Message] + +// GenerateSummaryPhase indicates which phase a model generate attempt belongs to during summarization. +type GenerateSummaryPhase string + +const ( + // GenerateSummaryPhasePrimary indicates an attempt using the primary model. + // Attempt=1 is the initial call; Attempt>1 indicates a retry. + GenerateSummaryPhasePrimary GenerateSummaryPhase = "primary" + + // GenerateSummaryPhaseFailover indicates an attempt using a failover model + // after the primary model exhausted all retries or was deemed unrecoverable. + GenerateSummaryPhaseFailover GenerateSummaryPhase = "failover" +) + +// TypedGenerateSummaryAction contains details of a single model generate attempt during summarization. +// Emitted on every attempt, whether it succeeds or fails. +type TypedGenerateSummaryAction[M adk.MessageType] struct { + // Attempt is the 1-based attempt number within the current phase. + // For primary phase, Attempt=1 is the initial call and Attempt>1 indicates retries. + // For failover phase, Attempt counts the failover rounds (1, 2, 3, ...). + Attempt int `json:"attempt"` + + // Phase indicates which phase this generate attempt belongs to. + Phase GenerateSummaryPhase `json:"phase"` + + // ModelResponse is the raw response returned by the model. + // It may be nil when the model call fails without returning a response. + ModelResponse M `json:"model_response,omitempty"` + + // err is the error returned by the model call, if any. Use GetError to access it. + err error +} + +// GenerateSummaryAction is the default type using *schema.Message. +type GenerateSummaryAction = TypedGenerateSummaryAction[*schema.Message] + +// GetError returns the error from the model call, if any. +func (a *TypedGenerateSummaryAction[M]) GetError() error { + return a.err +} diff --git a/adk/middlewares/summarization/finalizer_builder.go b/adk/middlewares/summarization/finalizer_builder.go new file mode 100644 index 0000000..7e79b2b --- /dev/null +++ b/adk/middlewares/summarization/finalizer_builder.go @@ -0,0 +1,442 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +import ( + "context" + "fmt" + "strings" + "unicode/utf8" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +// DefaultFinalize is the default TypedFinalizeFunc implementation, providing the same +// summary post-processing as the middleware does internally: +// 1. Replaces the ... section in the model-generated +// summary with recent original user messages from the conversation (up to ~30k tokens). +// 2. Adds a preamble and a postamble around the summary content. +// 3. Converts the summary into a user message, prepended with the original system messages. +func DefaultFinalize[M adk.MessageType](ctx context.Context, originalMessages []M, summary M) ([]M, error) { + systemMsgs, contextMsgs := splitSystemAndContextMsgs(originalMessages) + + processed, err := postProcessSummary(ctx, &postProcessSummaryParams[M]{ + contextMsgs: contextMsgs, + summaryContent: getAssistantTextContent(summary), + }) + if err != nil { + return nil, err + } + + result := make([]M, 0, len(systemMsgs)+1) + result = append(result, systemMsgs...) + result = append(result, processed) + return result, nil +} + +// TypedFinalizerBuilder builds a TypedFinalizeFunc by chaining handlers, +// generic over message type M. +type TypedFinalizerBuilder[M adk.MessageType] struct { + handlers []TypedFinalizeFunc[M] + errs []error +} + +// FinalizerBuilder is a backward-compatible alias for TypedFinalizerBuilder +// specialized with *schema.Message. +type FinalizerBuilder = TypedFinalizerBuilder[*schema.Message] + +// NewTypedFinalizer creates a new TypedFinalizerBuilder. +// +// Handlers run in registration order, and the summary message is post-processed +// as DefaultFinalize does after all handlers have run. For example, with +// PreserveSkills and a system message in originalMessages, the final output is: +// +// [system message, preserved skill message, processed summary] +// +// Example: +// +// finalizer, err := NewTypedFinalizer[*schema.Message](). +// PreserveSkills(&PreserveSkillsConfig{}). +// Build() +// +// cfg := &Config{ +// Finalize: finalizer, +// // ... +// } +func NewTypedFinalizer[M adk.MessageType]() *TypedFinalizerBuilder[M] { + return &TypedFinalizerBuilder[M]{} +} + +// NewFinalizer creates a new FinalizerBuilder that builds a FinalizeFunc +// by chaining handlers. +func NewFinalizer() *FinalizerBuilder { + return &FinalizerBuilder{} +} + +// Build constructs the final TypedFinalizeFunc by chaining all registered handlers +// and post-processing the summary message as DefaultFinalize does. +// For example, with PreserveSkills and a system message in +// originalMessages, the final output is: +// +// [system message, preserved skill message, processed summary] +func (b *TypedFinalizerBuilder[M]) Build() (TypedFinalizeFunc[M], error) { + if len(b.errs) > 0 { + msgs := make([]string, len(b.errs)) + for i, e := range b.errs { + msgs[i] = e.Error() + } + return nil, fmt.Errorf("failed to build finalizer:\n%s", strings.Join(msgs, "\n")) + } + + if len(b.handlers) == 0 { + return nil, fmt.Errorf("at least one handler is required") + } + + handlers := make([]TypedFinalizeFunc[M], len(b.handlers)) + copy(handlers, b.handlers) + + return func(ctx context.Context, originalMessages []M, summary M) ([]M, error) { + var extraMessages []M + for _, fn := range handlers { + result, err := fn(ctx, originalMessages, summary) + if err != nil { + return nil, err + } + if len(result) == 0 { + return nil, fmt.Errorf("finalizer handler returned no messages") + } + extraMessages = append(extraMessages, result[:len(result)-1]...) + summary = result[len(result)-1] + } + + systemMsgs, contextMsgs := splitSystemAndContextMsgs(originalMessages) + + processed, err := postProcessSummary(ctx, &postProcessSummaryParams[M]{ + contextMsgs: contextMsgs, + summaryContent: getAssistantTextContent(summary), + }) + if err != nil { + return nil, err + } + + result := make([]M, 0, len(systemMsgs)+len(extraMessages)+1) + result = append(result, systemMsgs...) + result = append(result, extraMessages...) + result = append(result, processed) + return result, nil + }, nil +} + +type PreserveSkillsConfig struct { + // SkillToolName is the tool name used for loading skills. + // Must match the tool name configured in the ADK skill middleware. + // Optional. Defaults to "skill". + SkillToolName string + + // MaxSkills limits the maximum number of skills to preserve. + // = 0 means do not preserve any skills (disabled). + // > 0 means preserve up to this many most recent skills. + // Optional. Defaults to 5. + MaxSkills *int + + // MaxTokensPerSkill limits the maximum token count for a single preserved skill. + // Skills exceeding this limit are truncated, with the truncated portion replaced + // by a short marker text (e.g. "[... skill content truncated ...]"). + // Note: if this value is set smaller than the token count of the marker text itself, + // the skill will contain only the marker text with no original content preserved. + // Optional. Defaults to 5000. + MaxTokensPerSkill *int + + // SkillsTokenBudget limits the total token count for all preserved skills combined. + // Skills are preserved from most recent to oldest; once the budget is exhausted, + // remaining skills are dropped. + // Optional. Defaults to 25000. + SkillsTokenBudget *int +} + +// PreserveSkills preserves skill contents loaded by the ADK skill middleware. +// It scans the conversation for matching skill tool calls and returns the preserved +// skill content as a user message before the summary. +// +// Example: +// +// messages: [assistant(tool_call: skill "foo"), tool(content: "bar")] +// summary: S +// +// When skill content is found, PreserveSkills returns: +// +// []M{user(""), S} +func (b *TypedFinalizerBuilder[M]) PreserveSkills(config *PreserveSkillsConfig) *TypedFinalizerBuilder[M] { + if err := config.check(); err != nil { + b.errs = append(b.errs, fmt.Errorf("PreserveSkills: %w", err)) + return b + } + b.handlers = append(b.handlers, func(ctx context.Context, originalMessages []M, summary M) ([]M, error) { + messages := originalMessages + + modelInput, ok := ctx.Value(ctxKeyModelInput{}).([]M) + if ok && len(modelInput) > 0 { + messages = modelInput + } + + if len(messages) == 0 { + return []M{summary}, nil + } + + skillText, err := buildPreservedSkillsText(ctx, messages, config) + if err != nil { + return nil, err + } + + if skillText == "" { + return []M{summary}, nil + } + + preserved := makeUserMsg[M](skillText) + setMsgExtra(preserved, extraKeyContentType, string(contentTypeSkills)) + return []M{preserved, summary}, nil + }) + return b +} + +func (c *PreserveSkillsConfig) check() error { + if c == nil { + return fmt.Errorf("PreserveSkillsConfig is required") + } + if c.MaxSkills != nil && *c.MaxSkills < 0 { + return fmt.Errorf("MaxSkills must be non-negative") + } + if c.MaxTokensPerSkill != nil && *c.MaxTokensPerSkill < 0 { + return fmt.Errorf("MaxTokensPerSkill must be non-negative") + } + if c.SkillsTokenBudget != nil && *c.SkillsTokenBudget < 0 { + return fmt.Errorf("SkillsTokenBudget must be non-negative") + } + return nil +} + +type skillInfo struct { + Name string + Content string +} + +func extractSkillInfos[M adk.MessageType](messages []M, skillTool string) ([]*skillInfo, error) { + var skills []*skillInfo + argsMap := make(map[string]string) + + for _, msg := range messages { + switch m := any(msg).(type) { + case *schema.Message: + if m.Role == schema.Assistant { + for _, tc := range m.ToolCalls { + if tc.Function.Name == skillTool { + argsMap[tc.ID] = tc.Function.Arguments + } + } + } else if m.Role == schema.Tool { + arguments, ok := argsMap[m.ToolCallID] + if !ok { + continue + } + var arg struct { + Skill string `json:"skill"` + } + if err := sonic.UnmarshalString(arguments, &arg); err != nil { + return nil, fmt.Errorf("failed to parse skill arguments from tool call %s: %w", m.ToolCallID, err) + } + skills = append(skills, &skillInfo{ + Name: arg.Skill, + Content: messageUserTextContent(m), + }) + } + + case *schema.AgenticMessage: + for _, block := range m.ContentBlocks { + if block == nil { + continue + } + if block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil { + if block.FunctionToolCall.Name == skillTool { + argsMap[block.FunctionToolCall.CallID] = block.FunctionToolCall.Arguments + } + } + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + arguments, ok := argsMap[block.FunctionToolResult.CallID] + if !ok { + continue + } + var arg struct { + Skill string `json:"skill"` + } + if err := sonic.UnmarshalString(arguments, &arg); err != nil { + return nil, fmt.Errorf("failed to parse skill arguments from tool call %s: %w", block.FunctionToolResult.CallID, err) + } + var contentParts []string + for _, cb := range block.FunctionToolResult.Content { + if cb != nil && cb.Type == schema.FunctionToolResultContentBlockTypeText && cb.Text != nil { + contentParts = append(contentParts, cb.Text.Text) + } + } + skills = append(skills, &skillInfo{ + Name: arg.Skill, + Content: strings.Join(contentParts, "\n"), + }) + } + } + } + } + + return skills, nil +} + +func buildPreservedSkillsText[M adk.MessageType](_ context.Context, messages []M, config *PreserveSkillsConfig) (string, error) { + const ( + defaultSkillTool = "skill" + defaultMaxTokensPerSkill = 5000 + defaultSkillsTokenBudget = 25000 + ) + + if config == nil { + config = &PreserveSkillsConfig{} + } + + maxSkills := 5 + if config.MaxSkills != nil { + maxSkills = *config.MaxSkills + } + if maxSkills <= 0 { + return "", nil + } + + skillTool := defaultSkillTool + if config.SkillToolName != "" { + skillTool = config.SkillToolName + } + + maxTokensPerSkill := defaultMaxTokensPerSkill + if config.MaxTokensPerSkill != nil { + maxTokensPerSkill = *config.MaxTokensPerSkill + } + + skillsTokenBudget := defaultSkillsTokenBudget + if config.SkillsTokenBudget != nil { + skillsTokenBudget = *config.SkillsTokenBudget + } + + skills, err := extractSkillInfos(messages, skillTool) + if err != nil { + return "", err + } + + if len(skills) == 0 { + return "", nil + } + + uniqueSkills := make([]*skillInfo, 0, len(skills)) + seenNames := make(map[string]bool) + for i := len(skills) - 1; i >= 0; i-- { + skill := skills[i] + if !seenNames[skill.Name] { + seenNames[skill.Name] = true + uniqueSkills = append(uniqueSkills, skill) + } + } + + for i, j := 0, len(uniqueSkills)-1; i < j; i, j = i+1, j-1 { + uniqueSkills[i], uniqueSkills[j] = uniqueSkills[j], uniqueSkills[i] + } + skills = uniqueSkills + + if len(skills) > maxSkills { + skills = skills[len(skills)-maxSkills:] + } + + totalTokens := 0 + var budgetedSkills []*skillInfo + for i := len(skills) - 1; i >= 0; i-- { + skill := skills[i] + tokens := estimateTokenCount(len(skill.Content)) + + if tokens > maxTokensPerSkill { + skill = &skillInfo{ + Name: skill.Name, + Content: truncateSkillContent(skill.Content, maxTokensPerSkill), + } + tokens = maxTokensPerSkill + } + + if totalTokens+tokens > skillsTokenBudget { + break + } + + totalTokens += tokens + budgetedSkills = append(budgetedSkills, skill) + } + + if len(budgetedSkills) == 0 { + return "", nil + } + + // Reverse to restore chronological order. + for i, j := 0, len(budgetedSkills)-1; i < j; i, j = i+1, j-1 { + budgetedSkills[i], budgetedSkills[j] = budgetedSkills[j], budgetedSkills[i] + } + + var parts []string + for _, skill := range budgetedSkills { + parts = append(parts, fmt.Sprintf(skillSectionFormat, skill.Name, skill.Content)) + } + + skillsText := strings.Join(parts, "\n\n---\n\n") + skillsText = fmt.Sprintf(getSkillPreamble(), skillsText) + skillsText = fmt.Sprintf("\n%s"+"\n", skillsText) + + return skillsText, nil +} + +// truncateSkillContent truncates skill content to fit within maxTokens. +// It keeps the first portion of the content and appends a truncation marker +// (e.g. "[... skill content truncated ...]") to indicate the omission. +// If maxTokens is smaller than the marker itself, only the marker is returned. +func truncateSkillContent(content string, maxTokens int) string { + if len(content) == 0 { + return content + } + + if estimateTokenCount(len(content)) <= maxTokens { + return content + } + + marker := getSkillTruncationMarker() + targetBytes := estimateTokenBytes(maxTokens) - len(marker) + if targetBytes < 0 { + targetBytes = 0 + } + if targetBytes > len(content) { + targetBytes = len(content) + } + + // Back up to a valid UTF-8 rune boundary. + for targetBytes > 0 && targetBytes < len(content) && !utf8.RuneStart(content[targetBytes]) { + targetBytes-- + } + + return content[:targetBytes] + marker +} diff --git a/adk/middlewares/summarization/finalizer_builder_test.go b/adk/middlewares/summarization/finalizer_builder_test.go new file mode 100644 index 0000000..a9290a0 --- /dev/null +++ b/adk/middlewares/summarization/finalizer_builder_test.go @@ -0,0 +1,528 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +import ( + "context" + "errors" + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestNewFinalizer(t *testing.T) { + b := NewFinalizer() + assert.NotNil(t, b) + assert.Empty(t, b.handlers) + + tb := NewTypedFinalizer[*schema.Message]() + assert.NotNil(t, tb) + assert.Empty(t, tb.handlers) +} + +func TestBuildEmpty(t *testing.T) { + finalizer, err := NewFinalizer().Build() + assert.Error(t, err) + assert.Nil(t, finalizer) + assert.Contains(t, err.Error(), "at least one handler is required") +} + +func TestBuildConfigError(t *testing.T) { + ptr := func(i int) *int { return &i } + + t.Run("nil config", func(t *testing.T) { + finalizer, err := NewFinalizer(). + PreserveSkills(nil). + Build() + assert.Error(t, err) + assert.Nil(t, finalizer) + assert.Contains(t, err.Error(), "PreserveSkills:") + assert.Contains(t, err.Error(), "PreserveSkillsConfig is required") + }) + + t.Run("negative max skills", func(t *testing.T) { + finalizer, err := NewFinalizer(). + PreserveSkills(&PreserveSkillsConfig{MaxSkills: ptr(-1)}). + Build() + assert.Error(t, err) + assert.Nil(t, finalizer) + assert.Contains(t, err.Error(), "PreserveSkills:") + assert.Contains(t, err.Error(), "MaxSkills must be non-negative") + }) +} + +func TestDefaultFinalizeBasic(t *testing.T) { + ctx := context.Background() + + result, err := DefaultFinalize(ctx, []adk.Message{ + schema.SystemMessage("system prompt"), + schema.UserMessage("original user"), + }, schema.AssistantMessage("raw summary", nil)) + assert.NoError(t, err) + assert.Len(t, result, 2) + + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, schema.User, result[1].Role) + assert.Equal(t, contentTypeSummary, typedGetContentType(result[1])) + assert.Contains(t, result[1].Content, "raw summary") +} + +func TestBuildStepChaining(t *testing.T) { + ctx := context.Background() + + b := NewFinalizer() + b.handlers = append(b.handlers, func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + summary.Content = summary.Content + " | step1" + return []adk.Message{summary}, nil + }) + b.handlers = append(b.handlers, func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + summary.Content = summary.Content + " | step2" + return []adk.Message{summary}, nil + }) + + finalizer, err := b.Build() + assert.NoError(t, err) + + summary := schema.AssistantMessage("start", nil) + result, err := finalizer(ctx, []adk.Message{}, summary) + assert.NoError(t, err) + assert.Len(t, result, 1) + assert.Equal(t, schema.User, result[0].Role) + assert.Contains(t, result[0].Content, "start | step1 | step2") +} + +func TestBuildStepError(t *testing.T) { + ctx := context.Background() + + b := NewFinalizer() + b.handlers = append(b.handlers, func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return nil, errors.New("step failed") + }) + + finalizer, err := b.Build() + assert.NoError(t, err) + + summary := schema.UserMessage("test") + _, err = finalizer(ctx, []adk.Message{}, summary) + assert.Error(t, err) + assert.Contains(t, err.Error(), "step failed") +} + +func TestBuildHandlerReturnsEmpty(t *testing.T) { + ctx := context.Background() + + b := NewFinalizer() + b.handlers = append(b.handlers, func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return []adk.Message{}, nil + }) + + finalizer, err := b.Build() + assert.NoError(t, err) + + _, err = finalizer(ctx, []adk.Message{}, schema.AssistantMessage("test", nil)) + assert.Error(t, err) + assert.Contains(t, err.Error(), "finalizer handler returned no messages") +} + +func TestBuildPostProcessError(t *testing.T) { + ctx := context.Background() + + b := NewFinalizer() + b.handlers = append(b.handlers, func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return []adk.Message{schema.UserMessage("not assistant")}, nil + }) + + finalizer, err := b.Build() + assert.NoError(t, err) + + _, err = finalizer(ctx, []adk.Message{}, schema.AssistantMessage("test", nil)) + assert.Error(t, err) + assert.Contains(t, err.Error(), "summary content is empty") +} + +func TestDefaultFinalizeError(t *testing.T) { + ctx := context.Background() + + _, err := DefaultFinalize(ctx, []adk.Message{ + schema.UserMessage("original"), + }, schema.UserMessage("not an assistant message")) + assert.Error(t, err) + assert.Contains(t, err.Error(), "summary content is empty") +} + +func TestPreserveSkillsConfigCheck(t *testing.T) { + ptr := func(i int) *int { return &i } + + t.Run("nil config", func(t *testing.T) { + var c *PreserveSkillsConfig + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "PreserveSkillsConfig is required") + }) + + t.Run("valid config", func(t *testing.T) { + c := &PreserveSkillsConfig{ + MaxSkills: ptr(5), + SkillToolName: "load_skill", + } + assert.NoError(t, c.check()) + }) + + t.Run("zero max skills", func(t *testing.T) { + c := &PreserveSkillsConfig{ + MaxSkills: ptr(0), + } + assert.NoError(t, c.check()) + }) + + t.Run("negative max skills", func(t *testing.T) { + c := &PreserveSkillsConfig{ + MaxSkills: ptr(-1), + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "MaxSkills must be non-negative") + }) + + t.Run("nil max skills", func(t *testing.T) { + c := &PreserveSkillsConfig{} + err := c.check() + assert.NoError(t, err) + }) + + t.Run("negative max tokens per skill", func(t *testing.T) { + c := &PreserveSkillsConfig{ + MaxTokensPerSkill: ptr(-1), + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "MaxTokensPerSkill must be non-negative") + }) + + t.Run("negative skills token budget", func(t *testing.T) { + c := &PreserveSkillsConfig{ + SkillsTokenBudget: ptr(-1), + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "SkillsTokenBudget must be non-negative") + }) +} + +func TestPreserveSkillsViaBuilder(t *testing.T) { + ptr := func(i int) *int { return &i } + ctx := context.Background() + + finalizer, err := NewFinalizer(). + PreserveSkills(&PreserveSkillsConfig{ + MaxSkills: ptr(2), + SkillToolName: "load_skill", + }). + Build() + assert.NoError(t, err) + + originalMessages := []adk.Message{ + schema.SystemMessage("system prompt"), + schema.UserMessage("original"), + } + modelInput := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "load_skill", + Arguments: `{"skill": "test-skill"}`, + }, + }, + }, + }, + { + Role: schema.Tool, + ToolCallID: "call_1", + Content: "skill content 1", + }, + } + ctx = context.WithValue(ctx, ctxKeyModelInput{}, modelInput) + + summary := schema.AssistantMessage("test summary", nil) + + result, err := finalizer(ctx, originalMessages, summary) + assert.NoError(t, err) + assert.Len(t, result, 3) + + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, "system prompt", result[0].Content) + + assert.Equal(t, schema.User, result[1].Role) + assert.Equal(t, contentTypeSkills, typedGetContentType(result[1])) + assert.Contains(t, result[1].Content, "test-skill") + assert.Contains(t, result[1].Content, "skill content 1") + + assert.Equal(t, schema.User, result[2].Role) + assert.Equal(t, contentTypeSummary, typedGetContentType(result[2])) + assert.Contains(t, result[2].Content, "test summary") +} + +func TestBuildPreservedSkillsText(t *testing.T) { + ptr := func(i int) *int { return &i } + ctx := context.Background() + + t.Run("nil config", func(t *testing.T) { + text, err := buildPreservedSkillsText[*schema.Message](ctx, nil, nil) + assert.NoError(t, err) + assert.Empty(t, text) + }) + + t.Run("zero max skills", func(t *testing.T) { + text, err := buildPreservedSkillsText[*schema.Message](ctx, nil, &PreserveSkillsConfig{MaxSkills: ptr(0)}) + assert.NoError(t, err) + assert.Empty(t, text) + }) + + t.Run("no matching skills", func(t *testing.T) { + text, err := buildPreservedSkillsText(ctx, []adk.Message{ + schema.UserMessage("hi"), + }, &PreserveSkillsConfig{ + MaxSkills: ptr(5), + SkillToolName: "load_skill", + }) + assert.NoError(t, err) + assert.Empty(t, text) + }) + + t.Run("with default skill tool name", func(t *testing.T) { + messages := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "skill", + Arguments: `{"skill": "test-skill"}`, + }, + }, + }, + }, + { + Role: schema.Tool, + ToolCallID: "call_1", + Content: "skill content 1", + }, + } + + config := &PreserveSkillsConfig{ + MaxSkills: ptr(2), + } + + text, err := buildPreservedSkillsText(ctx, messages, config) + assert.NoError(t, err) + assert.Contains(t, text, "test-skill") + assert.Contains(t, text, "skill content 1") + }) + + t.Run("parse error", func(t *testing.T) { + messages := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "load_skill", + Arguments: `invalid json`, + }, + }, + }, + }, + { + Role: schema.Tool, + ToolCallID: "call_1", + Content: "content", + }, + } + + _, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{ + MaxSkills: ptr(2), + SkillToolName: "load_skill", + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to parse skill arguments") + }) + + t.Run("max skills truncation and deduplication", func(t *testing.T) { + messages := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill1"}`}}, + {ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill2"}`}}, + {ID: "call_3", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill1"}`}}, + {ID: "call_4", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill3"}`}}, + }, + }, + {Role: schema.Tool, ToolCallID: "call_1", Content: "c1"}, + {Role: schema.Tool, ToolCallID: "call_2", Content: "c2"}, + {Role: schema.Tool, ToolCallID: "call_3", Content: "c3"}, + {Role: schema.Tool, ToolCallID: "call_4", Content: "c4"}, + } + + text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{ + MaxSkills: ptr(2), + SkillToolName: "load_skill", + }) + assert.NoError(t, err) + assert.Contains(t, text, "skill1") + assert.Contains(t, text, "c3") + assert.Contains(t, text, "skill3") + assert.Contains(t, text, "c4") + assert.NotContains(t, text, "c1") + assert.NotContains(t, text, "skill2") + assert.NotContains(t, text, "c2") + }) + + t.Run("per skill token limit truncates large skills", func(t *testing.T) { + // estimateTokenCount = (len+3)/4 + // "short" = 5 chars → 2 tokens + // strings.Repeat("x", 100) = 100 chars → 25 tokens + largeContent := strings.Repeat("x", 100) + messages := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "small"}`}}, + {ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "large"}`}}, + }, + }, + {Role: schema.Tool, ToolCallID: "call_1", Content: "short"}, + {Role: schema.Tool, ToolCallID: "call_2", Content: largeContent}, + } + + // MaxTokensPerSkill=10: "short"→2 tokens (ok), largeContent→25 tokens (truncated) + text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{ + MaxSkills: ptr(10), + MaxTokensPerSkill: ptr(10), + SkillToolName: "load_skill", + }) + assert.NoError(t, err) + // small skill preserved as-is + assert.Contains(t, text, "small") + assert.Contains(t, text, "short") + // large skill is truncated, not dropped — name still present, full content gone + assert.Contains(t, text, "large") + assert.NotContains(t, text, largeContent) + assert.Contains(t, text, "skill content truncated for compaction") + }) + + t.Run("total token budget drops excess skills", func(t *testing.T) { + // Each content is 40 chars → (40+3)/4 = 10 tokens + content := strings.Repeat("a", 40) + messages := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill1"}`}}, + {ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill2"}`}}, + {ID: "call_3", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "skill3"}`}}, + }, + }, + {Role: schema.Tool, ToolCallID: "call_1", Content: content}, + {Role: schema.Tool, ToolCallID: "call_2", Content: content}, + {Role: schema.Tool, ToolCallID: "call_3", Content: content}, + } + + // Budget=15: skill3=10 tokens fits, skill2=10 tokens → 10+10=20 > 15, stop. + text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{ + MaxSkills: ptr(10), + SkillsTokenBudget: ptr(15), + SkillToolName: "load_skill", + }) + assert.NoError(t, err) + assert.Contains(t, text, "skill3") + assert.NotContains(t, text, "skill1") + assert.NotContains(t, text, "skill2") + }) + + t.Run("token budget and per-skill limit combined", func(t *testing.T) { + // s1: 16 chars → 4 tokens + // s2: 200 chars → 50 tokens (exceeds per-skill limit of 20, gets truncated) + // s3: 24 chars → 6 tokens + // s4: 24 chars → 6 tokens + messages := []adk.Message{ + { + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s1"}`}}, + {ID: "call_2", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s2"}`}}, + {ID: "call_3", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s3"}`}}, + {ID: "call_4", Function: schema.FunctionCall{Name: "load_skill", Arguments: `{"skill": "s4"}`}}, + }, + }, + {Role: schema.Tool, ToolCallID: "call_1", Content: strings.Repeat("a", 16)}, + {Role: schema.Tool, ToolCallID: "call_2", Content: strings.Repeat("b", 200)}, + {Role: schema.Tool, ToolCallID: "call_3", Content: strings.Repeat("c", 24)}, + {Role: schema.Tool, ToolCallID: "call_4", Content: strings.Repeat("d", 24)}, + } + + // Per-skill limit: 20 (s2 with 50 tokens is truncated to 20) + // Budget: 30 (from most recent: s4=6, s3=6, s2=20, total=32 > 30, so s2 cannot fit) + // Result: s4 and s3 preserved + text, err := buildPreservedSkillsText(ctx, messages, &PreserveSkillsConfig{ + MaxSkills: ptr(10), + MaxTokensPerSkill: ptr(20), + SkillsTokenBudget: ptr(30), + SkillToolName: "load_skill", + }) + assert.NoError(t, err) + assert.Contains(t, text, "s3") + assert.Contains(t, text, "s4") + assert.NotContains(t, text, "\"s1\"") + assert.NotContains(t, text, "\"s2\"") + }) + + t.Run("truncated skill content preserves only prefix", func(t *testing.T) { + // Use a long content and generous maxTokens so the prefix is clearly visible. + content := strings.Repeat("abcdefghij", 100) // 1000 bytes → 250 tokens + // maxTokens=125 → targetBytes = 500, minus ~101 marker bytes → ~399 prefix bytes + truncated := truncateSkillContent(content, 125) + assert.True(t, strings.HasPrefix(truncated, "abcdefghij")) // prefix preserved + assert.Contains(t, truncated, "skill content truncated for compaction") + assert.NotEqual(t, content, truncated) + // Ends with marker, not with original content suffix + assert.True(t, strings.HasSuffix(truncated, "]")) + // No suffix from original content + assert.False(t, strings.HasSuffix(truncated, "abcdefghij]")) + }) + + t.Run("truncated multibyte content does not produce invalid utf8", func(t *testing.T) { + // Each Chinese char is 3 bytes. 334 chars = 1002 bytes → 251 tokens + content := strings.Repeat("中", 334) + // maxTokens=125 → targetBytes=500, minus marker ~101 bytes → ~399 bytes + // 399 / 3 = 133 full Chinese chars, no partial rune + truncated := truncateSkillContent(content, 125) + assert.True(t, utf8.ValidString(truncated)) + assert.True(t, strings.HasPrefix(truncated, "中中中")) + assert.Contains(t, truncated, "skill content truncated for compaction") + }) +} diff --git a/adk/middlewares/summarization/prompt.go b/adk/middlewares/summarization/prompt.go new file mode 100644 index 0000000..086017e --- /dev/null +++ b/adk/middlewares/summarization/prompt.go @@ -0,0 +1,340 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +import ( + "regexp" + + "github.com/cloudwego/eino/adk/internal" +) + +var allUserMessagesTagRegex = regexp.MustCompile(`(?s).*`) + +func getSystemInstruction() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: systemInstruction, + Chinese: systemInstructionZh, + }) +} + +func getUserSummaryInstruction() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: userSummaryInstruction, + Chinese: userSummaryInstructionZh, + }) +} + +func getSummaryPreamble() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: summaryPreamble, + Chinese: summaryPreambleZh, + }) +} + +func getContinueInstruction() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: continueInstruction, + Chinese: continueInstructionZh, + }) +} + +func getTranscriptPathInstruction() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: transcriptPathInstruction, + Chinese: transcriptPathInstructionZh, + }) +} + +func getTruncatedMarkerFormat() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: truncatedMarkerFormat, + Chinese: truncatedMarkerFormatZh, + }) +} + +func getUserMessagesReplacedNote() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: userMessagesReplacedNote, + Chinese: userMessagesReplacedNoteZh, + }) +} + +const systemInstruction = `You are a helpful AI assistant tasked with summarizing conversations.` + +const systemInstructionZh = `你是一个负责总结对话的 AI 助手。` + +const userSummaryInstruction = `CRITICAL: Respond with TEXT ONLY. Do NOT call any tools. + +- Do NOT use Read, Bash, Grep, Glob, Edit, Write, or ANY other tool. +- You already have all the context you need in the conversation above. +- Tool calls will be REJECTED and will waste your only turn — you will fail the task. +- Your entire response must be plain text: an block followed by a block. + +Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions. +This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context. + +Before providing your final summary, wrap your analysis in tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process: + +1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify: + - The user's explicit requests and intents + - Your approach to addressing the user's requests + - Key decisions, technical concepts and code patterns + - Specific details like: + - file names + - full code snippets + - function signatures + - file edits + - Errors that you ran into and how you fixed them + - Pay special attention to specific user feedback that you received, especially if the user told you to do something differently. +2. Double-check for technical accuracy and completeness, addressing each required element thoroughly. + +Your summary should include the following sections: + +1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail +2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed. +3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important. +4. Errors and fixes: List all errors that you ran into, and how you fixed them. Pay special attention to specific user feedback that you received, especially if the user told you to do something differently. +5. Problem Solving: Document problems solved and any ongoing troubleshooting efforts. +6. All user messages: List ALL user messages that are not tool results, and wrap them in the ... block. These are critical for understanding the users' feedback and changing intent. +7. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on. +8. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable. +9. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's most recent explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests or really old requests that were already completed without confirming with the user first. + If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation. + +Here's an example of how your output should be structured: + + + +[Your thought process, ensuring all points are covered thoroughly and accurately] + + + +1. Primary Request and Intent: + [Detailed description] + +2. Key Technical Concepts: + - [Concept 1] + - [Concept 2] + - [...] + +3. Files and Code Sections: + - [File Name 1] + - [Summary of why this file is important] + - [Summary of the changes made to this file, if any] + - [Important Code Snippet] + - [File Name 2] + - [Important Code Snippet] + - [...] + +4. Errors and fixes: + - [Detailed description of error 1]: + - [How you fixed the error] + - [User feedback on the error if any] + - [...] + +5. Problem Solving: + [Description of solved problems and ongoing troubleshooting] + +6. All user messages: + + - [Detailed non tool use user message] + - [...] + + +7. Pending Tasks: + - [Task 1] + - [Task 2] + - [...] + +8. Current Work: + [Precise description of current work] + +9. Optional Next Step: + [Optional Next step to take] + + + + +Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response. + +There may be additional summarization instructions provided in the included context. If so, remember to follow these instructions when creating the above summary. Examples of instructions include: + +## Compact Instructions +When summarizing the conversation focus on typescript code changes and also remember the mistakes you made and how you fixed them. + + + +# Summary instructions +When you are using compact - please focus on test output and code changes. Include file reads verbatim. + + + +REMINDER: Do NOT call any tools. Respond with plain text only — an block followed by a block. Tool calls will be rejected and you will fail the task. +` + +const userSummaryInstructionZh = `关键：仅以文本响应。不要调用任何工具。 + +- 不要使用 Read、Bash、Grep、Glob、Edit、Write 或任何其他工具。 +- 你已经拥有上述对话中所需的全部上下文。 +- 工具调用将被拒绝，并且会浪费你唯一的一次回复机会——你将无法完成任务。 +- 你的整个回复必须是纯文本：先是一个 代码块，后面紧跟一个 代码块。 + +你的任务是对目前为止的对话创建一份详细的总结，需要密切关注用户的明确请求和你之前的操作。 +这份总结应该全面捕捉技术细节、代码模式和架构决策，以确保继续开发工作时不丢失上下文。 + +在提供最终总结之前，请将你的分析过程包裹在 标签中，以组织思路并确保涵盖所有必要的要点。在分析过程中： + +1. 按时间顺序分析对话中的每条消息和每个部分。对于每个部分，需要全面识别： + - 用户的明确请求和意图 + - 你处理用户请求的方法 + - 关键决策、技术概念和代码模式 + - 具体细节，例如： + - 文件名 + - 完整代码片段 + - 函数签名 + - 文件编辑 + - 你遇到的错误以及如何修复它们 + - 特别注意你收到的具体用户反馈，尤其是用户要求你以不同方式处理的情况 +2. 仔细检查技术准确性和完整性，彻底处理每个必需的元素。 + +你的总结应包含以下部分： + +1. 主要请求和意图：详细捕捉用户所有的明确请求和意图 +2. 关键技术概念：列出讨论过的所有重要技术概念、技术和框架 +3. 文件和代码部分：列举检查、修改或创建的具体文件和代码部分。特别注意最近的消息，在适用的地方包含完整的代码片段，并总结为什么这个文件的读取或编辑很重要 +4. 错误和修复：列出你遇到的所有错误以及如何修复它们。特别注意你收到的具体用户反馈，尤其是用户要求你以不同方式处理的情况 +5. 问题解决：记录已解决的问题和任何正在进行的故障排除工作 +6. 所有用户消息：列出所有非工具结果的用户消息，并将它们包裹在 ... 块中。这些对于理解用户的反馈和变化的意图至关重要 +7. 待处理任务：列出明确要求你处理的任何待处理任务 +8. 当前工作：详细描述在此总结请求之前正在进行的工作，特别注意用户和助手的最近消息。在适用的地方包含文件名和代码片段 +9. 可选的下一步：列出与你最近工作相关的下一步操作。重要提示：确保这一步与用户最近的明确请求以及你在此总结请求之前正在处理的任务直接相关。如果你的上一个任务已经完成，则只有在与用户请求明确相关时才列出下一步。不要在未与用户确认的情况下开始处理无关的请求或已经完成的旧请求。 + 如果有下一步，请包含最近对话中的直接引用，准确显示你正在处理的任务以及你停止的位置。这应该是逐字引用，以确保任务理解不会偏离。 + +以下是输出结构的示例： + + + +[你的思考过程，确保全面准确地涵盖所有要点] + + + +1. 主要请求和意图： + [详细描述] + +2. 关键技术概念： + - [概念 1] + - [概念 2] + - [...] + +3. 文件和代码部分： + - [文件名 1] + - [该文件为何重要的总结] + - [对该文件所做更改的总结（如有）] + - [重要代码片段] + - [文件名 2] + - [重要代码片段] + - [...] + +4. 错误与修复： + - [错误 1 的详细描述]： + - [你如何修复该错误] + - [与该错误相关的用户反馈（如有）] + - [...] + +5. 问题解决： + [对已解决问题和正在进行中的排查工作的描述] + +6. 所有用户消息： + + - [详细的非工具调用的用户消息] + - [...] + + +7. 待处理任务： + - [任务 1] + - [任务 2] + - [...] + +8. 当前工作： + [当前工作的精确描述] + +9. 可选的下一步： + [可选的下一步操作] + + + + +请根据目前为止的对话提供你的总结，遵循此结构并确保回复的精确性和全面性。 + +上下文中可能包含额外的总结指令。如果有，请在创建上述总结时记得遵循这些指令。指令示例包括： + +## 压缩指令 +在总结对话时，重点关注 typescript 代码更改，并记住你犯的错误以及如何修复它们。 + + + +# 总结指令 +当你使用压缩时，请重点关注测试输出和代码更改。逐字包含文件读取内容。 + + + +提醒：不要调用任何工具。仅以纯文本响应——一个 代码块后面跟一个 代码块。工具调用将被拒绝，你将无法完成任务。 +` + +const summaryPreamble = `This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.` + +const summaryPreambleZh = `此会话延续自此前一段因上下文耗尽而终止的对话。以下总结概述了此前对话的内容。` + +const continueInstruction = `Continue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, do not preface with "I'll continue" or similar. Pick up the last task as if the break never happened.` + +const continueInstructionZh = `从对话中断的地方继续，不要再问用户任何问题。直接继续——不要提及总结，不要复述之前的经过，不要以"我会继续"之类的话作为开场白。就像中断从未发生过一样，继续上次的任务。` + +const transcriptPathInstruction = `If you need specific details from before compaction (like exact code snippets, error messages, or content you generated), read the full transcript at: %s` + +const transcriptPathInstructionZh = `如果你需要压缩之前的具体细节（如精确的代码片段、错误消息或你生成的内容），完整的对话记录位于：%s` + +const truncatedMarkerFormat = "…%d characters truncated…" + +const truncatedMarkerFormatZh = "…已截断 %d 个字符…" + +const userMessagesReplacedNote = "Some earlier user messages have been cleared. Below are the most recent user messages:" + +const userMessagesReplacedNoteZh = "部分较早的用户消息已被清除，以下是保留的最近用户消息：" + +const skillSectionFormat = "### Skill: %s\n\n%s" + +const skillPreamble = "The following skills were invoked in this session. Continue to follow these guidelines:\n\n%s" + +const skillPreambleZh = "以下 Skill 已在本会话中被调用，请继续遵循这些指导原则：\n\n%s" + +func getSkillPreamble() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: skillPreamble, + Chinese: skillPreambleZh, + }) +} + +const skillTruncationMarker = "\n\n[... skill content truncated for compaction; use Read on the skill path if you need the full text]" + +const skillTruncationMarkerZh = "\n\n[... skill 内容已在压缩时截断，如需完整内容请通过 Read 读取 skill 对应的文件路径]" + +func getSkillTruncationMarker() string { + return internal.SelectPrompt(internal.I18nPrompts{ + English: skillTruncationMarker, + Chinese: skillTruncationMarkerZh, + }) +} diff --git a/adk/middlewares/summarization/summarization.go b/adk/middlewares/summarization/summarization.go new file mode 100644 index 0000000..6110107 --- /dev/null +++ b/adk/middlewares/summarization/summarization.go @@ -0,0 +1,1289 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package summarization provides a middleware that automatically summarizes +// conversation history when token count exceeds the configured threshold. +package summarization + +import ( + "context" + "fmt" + "math/rand" + "regexp" + "strings" + "time" + "unicode/utf8" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*TypedCustomizedAction[*schema.Message]]("_eino_adk_summarization_mw_customized_action") + schema.RegisterName[*TypedCustomizedAction[*schema.AgenticMessage]]("_eino_adk_summarization_mw_customized_action_agentic") +} + +type TypedTokenCounterFunc[M adk.MessageType] func(ctx context.Context, input *TypedTokenCounterInput[M]) (int, error) +type TypedGenModelInputFunc[M adk.MessageType] func(ctx context.Context, sysInstruction, userInstruction M, originalMsgs []M) ([]M, error) +type TypedGetFailoverModelFunc[M adk.MessageType] func(ctx context.Context, failoverCtx *TypedFailoverContext[M]) (failoverModel model.BaseModel[M], failoverModelInputMsgs []M, failoverErr error) +type TypedFinalizeFunc[M adk.MessageType] func(ctx context.Context, originalMessages []M, summary M) ([]M, error) +type TypedCallbackFunc[M adk.MessageType] func(ctx context.Context, before, after adk.TypedChatModelAgentState[M]) error + +type TokenCounterFunc = TypedTokenCounterFunc[*schema.Message] +type GenModelInputFunc = TypedGenModelInputFunc[*schema.Message] +type GetFailoverModelFunc = TypedGetFailoverModelFunc[*schema.Message] +type FinalizeFunc = TypedFinalizeFunc[*schema.Message] +type CallbackFunc = TypedCallbackFunc[*schema.Message] + +// TypedConfig defines the configuration for the summarization middleware, +// generic over message type M. +type TypedConfig[M adk.MessageType] struct { + // Model is the chat model used to generate summaries. + Model model.BaseModel[M] + + // ModelOptions specifies options passed to the model when generating summaries. + // Optional. + ModelOptions []model.Option + + // TokenCounter calculates the token count for given messages and tools. + // + // Parameters: + // - input: contains the messages and tools to count tokens for. + // + // Returns: + // - int: the total token count. + // + // Optional. Defaults to using the total tokens reported in the last assistant + // message as baseline, with incremental messages estimated at ~4 chars/token. + TokenCounter TypedTokenCounterFunc[M] + + // Trigger specifies the conditions that activate summarization. + // Optional. Defaults to triggering when total tokens exceed 160k. + Trigger *TriggerCondition + + // EmitInternalEvents indicates whether internal events should be emitted during summarization, + // allowing external observers to track the summarization process. + // + // Event Scoping: + // - ActionTypeBeforeSummarize: emitted before calling model to generate summary + // - ActionTypeGenerateSummary: emitted after each model generate attempt + // - ActionTypeAfterSummarize: emitted after summary generation completes + // + // Optional. Defaults to false. + EmitInternalEvents bool + + // UserInstruction serves as the user-level instruction to guide the model on how to summarize the context. + // It is appended to the message history as a User message. + // If provided, it overrides the default user summarization instruction. + // Optional. + UserInstruction string + + // TranscriptFilePath is the path to the file containing the full conversation history. + // It is appended to the summary to remind the model where to read the original context. + // This field takes effect only when Finalize is not set. + // Optional but strongly recommended. + TranscriptFilePath string + + // GenModelInput allows full control over the summarization model input construction. + // + // Parameters: + // - sysInstruction: System message defining the model's role. It is set + // internally by the middleware and is not configurable. + // - userInstruction: User message with the task instruction. + // - originalMsgs: original complete message list. + // + // Returns: + // - []M: the constructed model input messages. + // + // Typical model input order: systemInstruction -> contextMessages -> userInstruction. + // + // Optional. + GenModelInput TypedGenModelInputFunc[M] + + // Finalize is called after summary generation. + // The returned messages are used as the final conversation history. + // When set, the middleware does not perform any post-processing on the summary. + // Use DefaultFinalize to apply the same post-processing as the default path. + // + // Parameters: + // - originalMessages: the original conversation messages before summarization. + // - summary: the model-generated summary message. + // + // Returns: + // - []M: the new conversation history to replace the original messages. + // + // Optional. + Finalize TypedFinalizeFunc[M] + + // Callback is called after Finalize, before exiting the middleware. + // Read-only, do not modify state. + // + // Parameters: + // - before: the agent state before summarization. + // - after: the agent state after summarization. + // + // Optional. + Callback TypedCallbackFunc[M] + + // Retry configures retry behavior for summary generation on the primary model. + // Optional. Defaults to no retries. + Retry *TypedRetryConfig[M] + + // Failover configures fallback behavior when summary generation on the primary model fails. + // Optional. + Failover *TypedFailoverConfig[M] +} + +// Config is a backward-compatible alias for TypedConfig specialized with *schema.Message. +type Config = TypedConfig[*schema.Message] + +// TypedTokenCounterInput is the input for TypedTokenCounterFunc. +type TypedTokenCounterInput[M adk.MessageType] struct { + // Messages is the list of messages to count tokens for. + Messages []M + // Tools is the list of tools to count tokens for. + Tools []*schema.ToolInfo +} + +// TokenCounterInput is a backward-compatible alias for TypedTokenCounterInput specialized with *schema.Message. +type TokenCounterInput = TypedTokenCounterInput[*schema.Message] + +// TriggerCondition specifies when summarization should be activated. +// Summarization triggers if ANY of the set conditions is met. +type TriggerCondition struct { + // ContextTokens triggers summarization when total token count exceeds this threshold. + ContextTokens int + // ContextMessages triggers summarization when total messages count exceeds this threshold. + ContextMessages int +} + +type TypedRetryConfig[M adk.MessageType] struct { + // MaxRetries specifies the maximum number of retry attempts. + // Optional. Defaults to 3. + MaxRetries *int + + // ShouldRetry determines whether a failed summary generation attempt should be retried. + // It is called after each failed attempt with the model response and error. + // Optional. Defaults to retrying when err is non-nil. + ShouldRetry func(ctx context.Context, resp M, err error) bool + + // BackoffFunc calculates the delay before the next retry attempt. + // The attempt parameter starts at 1 for the first retry. + // Optional. Defaults to a default exponential backoff with jitter. + BackoffFunc func(ctx context.Context, attempt int, resp M, err error) time.Duration +} + +// RetryConfig is a backward-compatible alias for TypedRetryConfig specialized with *schema.Message. +type RetryConfig = TypedRetryConfig[*schema.Message] + +type TypedFailoverConfig[M adk.MessageType] struct { + // MaxRetries specifies the maximum number of failover attempts. + // Optional. Defaults to 3. + MaxRetries *int + + // ShouldFailover determines whether another failover attempt should be made. + // It is called after each failover attempt with the model response and error. + // Optional. Defaults to failing over when err is non-nil. + ShouldFailover func(ctx context.Context, resp M, err error) bool + + // BackoffFunc calculates the delay before the next failover attempt. + // The attempt parameter starts at 1 for the first failover attempt. + // Optional. Defaults to a default exponential backoff with jitter. + BackoffFunc func(ctx context.Context, attempt int, resp M, err error) time.Duration + + // GetFailoverModel selects the model and input messages for the current failover attempt. + // + // Parameters: + // - failoverCtx: contains the context for the current failover attempt. + // + // Returns: + // - failoverModel: the model to use for this failover attempt. + // - failoverModelInputMsgs: the input messages to send to failoverModel. + // - failoverErr: an error encountered while preparing the failover model or input. + // + // Constraints: + // - When provided, it must return a non-nil model and a non-empty input message list. + // + // Optional. Defaults to reusing the primary model with the default input messages. + GetFailoverModel TypedGetFailoverModelFunc[M] +} + +// FailoverConfig is a backward-compatible alias for TypedFailoverConfig specialized with *schema.Message. +type FailoverConfig = TypedFailoverConfig[*schema.Message] + +// TypedFailoverContext contains the state for a failover attempt. +type TypedFailoverContext[M adk.MessageType] struct { + // Attempt is the current failover attempt number, starting at 1. + Attempt int + + // SystemInstruction is the system instruction used for summary generation. + // It is set internally by the middleware and is not configurable. + SystemInstruction M + + // UserInstruction is the user instruction used for summary generation. + UserInstruction M + + // OriginalMessages is the full original conversation before summarization. + OriginalMessages []M + + // LastModelResponse is the response returned by the previous attempt, if any. + LastModelResponse M + + // LastErr is the error returned by the previous attempt, if any. + LastErr error +} + +// FailoverContext is a backward-compatible alias for TypedFailoverContext specialized with *schema.Message. +type FailoverContext = TypedFailoverContext[*schema.Message] + +// NewTyped creates a generic summarization middleware that automatically summarizes +// conversation history when trigger conditions are met. +// +// This is the generic constructor that supports both *schema.Message and *schema.AgenticMessage. +func NewTyped[M adk.MessageType](_ context.Context, cfg *TypedConfig[M]) (adk.TypedChatModelAgentMiddleware[M], error) { + if err := cfg.check(); err != nil { + return nil, err + } + mw := &TypedMiddleware[M]{ + cfg: cfg, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[M]{}, + } + return mw, nil +} + +// New creates a summarization middleware that automatically summarizes conversation history +// when trigger conditions are met. +func New(ctx context.Context, cfg *Config) (adk.ChatModelAgentMiddleware, error) { + return NewTyped(ctx, cfg) +} + +type TypedMiddleware[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + cfg *TypedConfig[M] +} + +func (m *TypedMiddleware[M]) Summarize(ctx context.Context, state *adk.TypedChatModelAgentState[M]) ([]M, error) { + beforeState := *state + + if m.cfg.EmitInternalEvents { + err := m.emitEvent(ctx, &TypedCustomizedAction[M]{ + Type: ActionTypeBeforeSummarize, + Before: &TypedBeforeSummarizeAction[M]{ + Messages: beforeState.Messages, + }, + }) + if err != nil { + return nil, err + } + } + + rawSummary, modelInput, err := m.summarize(ctx, beforeState.Messages) + if err != nil { + return nil, err + } + + finalizeCtx := context.WithValue(ctx, ctxKeyModelInput{}, modelInput) + + finalizer := m.cfg.Finalize + if finalizer == nil { + finalizer = buildInternalFinalizer(m.cfg) + } + finalMsgs, err := finalizer(finalizeCtx, beforeState.Messages, rawSummary) + if err != nil { + return nil, err + } + + if m.cfg.Callback != nil { + afterState := beforeState + afterState.Messages = finalMsgs + if err = m.cfg.Callback(ctx, beforeState, afterState); err != nil { + return nil, err + } + } + + if m.cfg.EmitInternalEvents { + err = m.emitEvent(ctx, &TypedCustomizedAction[M]{ + Type: ActionTypeAfterSummarize, + After: &TypedAfterSummarizeAction[M]{ + Messages: finalMsgs, + }, + }) + if err != nil { + return nil, err + } + } + + return finalMsgs, nil +} + +func (m *TypedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M], + _ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) { + + triggered, err := m.shouldSummarize(ctx, &TypedTokenCounterInput[M]{ + Messages: state.Messages, + Tools: state.ToolInfos, + }) + if err != nil { + return nil, nil, err + } + if !triggered { + return ctx, state, nil + } + + finalMsgs, err := m.Summarize(ctx, state) + if err != nil { + return nil, nil, err + } + + afterState := *state + afterState.Messages = finalMsgs + + return ctx, &afterState, nil +} + +func (m *TypedMiddleware[M]) shouldSummarize(ctx context.Context, input *TypedTokenCounterInput[M]) (bool, error) { + if m.cfg.Trigger != nil && m.cfg.Trigger.ContextMessages > 0 { + if len(input.Messages) > m.cfg.Trigger.ContextMessages { + return true, nil + } + } + tokens, err := m.countTokens(ctx, input) + if err != nil { + return false, fmt.Errorf("failed to count tokens: %w", err) + } + return tokens > m.getTriggerContextTokens(), nil +} + +func (m *TypedMiddleware[M]) getTriggerContextTokens() int { + const defaultTriggerContextTokens = 160000 + if m.cfg.Trigger != nil { + return m.cfg.Trigger.ContextTokens + } + return defaultTriggerContextTokens +} + +func (m *TypedMiddleware[M]) emitEvent(ctx context.Context, action *TypedCustomizedAction[M]) error { + err := adk.TypedSendEvent(ctx, &adk.TypedAgentEvent[M]{ + Action: &adk.AgentAction{ + CustomizedAction: action, + }, + }) + if err != nil { + return fmt.Errorf("failed to send internal event: %w", err) + } + return nil +} + +func (m *TypedMiddleware[M]) emitGenerateSummaryEvent(ctx context.Context, attempt int, phase GenerateSummaryPhase, + resp M, err error) error { + + if !m.cfg.EmitInternalEvents { + return nil + } + + action := &TypedGenerateSummaryAction[M]{ + Attempt: attempt, + Phase: phase, + ModelResponse: resp, + err: err, + } + + return m.emitEvent(ctx, &TypedCustomizedAction[M]{ + Type: ActionTypeGenerateSummary, + GenerateSummary: action, + }) +} + +func (m *TypedMiddleware[M]) countTokens(ctx context.Context, input *TypedTokenCounterInput[M]) (int, error) { + if m.cfg.TokenCounter != nil { + return m.cfg.TokenCounter(ctx, input) + } + return defaultTypedTokenCounter(ctx, input) +} + +func defaultTypedTokenCounter[M adk.MessageType](_ context.Context, input *TypedTokenCounterInput[M]) (int, error) { + var ( + baseTokens int + incrementStart int + ) + + for i := len(input.Messages) - 1; i >= 0; i-- { + if tokens := getAssistantTotalTokens(input.Messages[i]); tokens > 0 { + baseTokens = tokens + incrementStart = i + 1 + break + } + } + + var incrementTokens int + for _, msg := range input.Messages[incrementStart:] { + switch m := any(msg).(type) { + case *schema.Message: + incrementTokens += estimateMessageTokens(m) + case *schema.AgenticMessage: + incrementTokens += estimateAgenticMessageTokens(m) + } + } + + for _, tl := range input.Tools { + tl_ := *tl + tl_.Extra = nil + text, err := sonic.MarshalString(tl_) + if err != nil { + return 0, fmt.Errorf("failed to marshal tool info: %w", err) + } + incrementTokens += estimateTokenCount(len(text)) + } + + return baseTokens + incrementTokens, nil +} + +func getAssistantTotalTokens[M adk.MessageType](msg M) int { + if msg == nil { + return 0 + } + switch m := any(msg).(type) { + case *schema.Message: + if m.Role != schema.Assistant || m.ResponseMeta == nil || m.ResponseMeta.Usage == nil { + return 0 + } + return m.ResponseMeta.Usage.TotalTokens + case *schema.AgenticMessage: + if m.Role != schema.AgenticRoleTypeAssistant || m.ResponseMeta == nil || m.ResponseMeta.TokenUsage == nil { + return 0 + } + return m.ResponseMeta.TokenUsage.TotalTokens + } + return 0 +} + +func estimateTokenCount(charLen int) int { + return charLen / 4 +} + +func estimateTokenBytes(tokens int) int { + return tokens * 4 +} + +func (m *TypedMiddleware[M]) summarize(ctx context.Context, originalMsgs []M) (M, []M, error) { + var zero M + _, contextMsgs := splitSystemAndContextMsgs(originalMsgs) + + modelInput, err := m.buildSummarizationModelInput(ctx, originalMsgs, contextMsgs) + if err != nil { + return zero, nil, err + } + + rawSummary, err := m.generateWithRetry(ctx, m.cfg.Model, modelInput, m.cfg.ModelOptions, m.cfg.Retry) + if typedShouldFailover(ctx, m.cfg.Failover, rawSummary, err) { + rawSummary, modelInput, err = m.runFailover(ctx, originalMsgs, modelInput, rawSummary, err) + if err != nil { + return zero, nil, err + } + } else if err != nil { + return zero, nil, fmt.Errorf("failed to generate summary: %w", err) + } + + return rawSummary, modelInput, nil +} + +func splitSystemAndContextMsgs[M adk.MessageType](msgs []M) ([]M, []M) { + var systemMsgs []M + for _, msg := range msgs { + if isSystemRole(msg) { + systemMsgs = append(systemMsgs, msg) + } else { + break + } + } + contextMsgs := msgs[len(systemMsgs):] + return systemMsgs, contextMsgs +} + +func (m *TypedMiddleware[M]) runFailover(ctx context.Context, originalMsgs, defaultInput []M, lastResp M, + lastErr error) (M, []M, error) { + + var zero M + const defaultMaxRetries = 3 + + sysInstruction, userInstruction := m.getModelInstructions() + + maxRetries := defaultMaxRetries + if m.cfg.Failover.MaxRetries != nil { + maxRetries = *m.cfg.Failover.MaxRetries + } + + backoff := m.cfg.Failover.BackoffFunc + if backoff == nil { + backoff = defaultTypedBackoffFunc[M] + } + + modelInput := defaultInput + + if maxRetries <= 0 { + return lastResp, modelInput, lastErr + } + + for attempt := 1; ; attempt++ { + fctx := &TypedFailoverContext[M]{ + Attempt: attempt, + SystemInstruction: sysInstruction, + UserInstruction: userInstruction, + OriginalMessages: originalMsgs, + LastModelResponse: lastResp, + LastErr: lastErr, + } + + failoverModel, nextInput, failoverErr := m.getFailoverModel(ctx, fctx, defaultInput) + if failoverErr != nil { + lastResp = zero + lastErr = failoverErr + if emitErr := m.emitGenerateSummaryEvent(ctx, attempt, GenerateSummaryPhaseFailover, zero, failoverErr); emitErr != nil { + return zero, nil, emitErr + } + } else { + modelInput = nextInput + lastResp, lastErr = m.generateAndEmit(ctx, failoverModel, modelInput, m.cfg.ModelOptions, attempt, GenerateSummaryPhaseFailover) + } + + if !typedShouldFailover(ctx, m.cfg.Failover, lastResp, lastErr) { + return lastResp, modelInput, lastErr + } + if attempt == maxRetries { + if lastErr != nil { + return zero, nil, fmt.Errorf("exceeds max failover attempts: %w", lastErr) + } + return zero, nil, fmt.Errorf("exceeds max failover attempts") + } + + select { + case <-time.After(backoff(ctx, attempt, lastResp, lastErr)): + case <-ctx.Done(): + return zero, nil, ctx.Err() + } + } +} + +func (m *TypedMiddleware[M]) getFailoverModel(ctx context.Context, failoverCtx *TypedFailoverContext[M], defaultInput []M) (model.BaseModel[M], []M, error) { + if m.cfg.Failover == nil { + return nil, nil, fmt.Errorf("failover config is required") + } + if m.cfg.Failover.GetFailoverModel == nil { + return m.cfg.Model, defaultInput, nil + } + + failoverModel, nextModelInput, err := m.cfg.Failover.GetFailoverModel(ctx, failoverCtx) + if err != nil { + return nil, nil, fmt.Errorf("failed to get failover model: %w", err) + } + + if failoverModel == nil { + return nil, nil, fmt.Errorf("failover model is required") + } + if len(nextModelInput) == 0 { + return nil, nil, fmt.Errorf("failover model input messages are required") + } + + return failoverModel, nextModelInput, nil +} + +func (m *TypedMiddleware[M]) buildSummarizationModelInput(ctx context.Context, originMsgs, contextMsgs []M) ([]M, error) { + sysInstruction, userInstruction := m.getModelInstructions() + + if m.cfg.GenModelInput != nil { + input, err := m.cfg.GenModelInput(ctx, sysInstruction, userInstruction, originMsgs) + if err != nil { + return nil, fmt.Errorf("failed to generate model input: %w", err) + } + return input, nil + } + + input := make([]M, 0, len(contextMsgs)+2) + input = append(input, sysInstruction) + input = append(input, contextMsgs...) + input = append(input, userInstruction) + + return input, nil +} + +func (m *TypedMiddleware[M]) getModelInstructions() (M, M) { + userInstruction := m.cfg.UserInstruction + if userInstruction == "" { + userInstruction = getUserSummaryInstruction() + } + + return makeSystemMsg[M](getSystemInstruction()), makeUserMsg[M](userInstruction) +} + +func buildInternalFinalizer[M adk.MessageType](cfg *TypedConfig[M]) TypedFinalizeFunc[M] { + return func(ctx context.Context, originalMessages []M, rawSummary M) ([]M, error) { + systemMsgs, contextMsgs := splitSystemAndContextMsgs(originalMessages) + + processed, err := postProcessSummary(ctx, &postProcessSummaryParams[M]{ + contextMsgs: contextMsgs, + summaryContent: getAssistantTextContent(rawSummary), + transcriptPath: cfg.TranscriptFilePath, + }) + if err != nil { + return nil, err + } + + return append(systemMsgs, processed), nil + } +} + +func getAssistantTextContent[M adk.MessageType](msg M) string { + switch r := any(msg).(type) { + case *schema.Message: + if r.Role != schema.Assistant { + return "" + } + var parts []string + for _, part := range r.AssistantGenMultiContent { + if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { + parts = append(parts, part.Text) + } + } + if len(parts) > 0 { + return strings.Join(parts, "\n") + } + return r.Content + case *schema.AgenticMessage: + if r.Role != schema.AgenticRoleTypeAssistant { + return "" + } + var parts []string + for _, block := range r.ContentBlocks { + if block != nil && block.AssistantGenText != nil { + parts = append(parts, block.AssistantGenText.Text) + } + } + return strings.Join(parts, "\n") + default: + return "" + } +} + +type postProcessSummaryParams[M adk.MessageType] struct { + contextMsgs []M + summaryContent string + transcriptPath string +} + +func postProcessSummary[M adk.MessageType](ctx context.Context, p *postProcessSummaryParams[M]) (processed M, err error) { + content := p.summaryContent + if content == "" { + return processed, fmt.Errorf("summary content is empty") + } + + if len(p.contextMsgs) > 0 { + newContent, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[M]{ + contextMsgs: p.contextMsgs, + summaryText: content, + }) + if err != nil { + return processed, fmt.Errorf("failed to populate user messages in summary: %w", err) + } + content = newContent + } + + if p.transcriptPath != "" { + content = appendSection(content, fmt.Sprintf(getTranscriptPathInstruction(), p.transcriptPath)) + } + + content = appendSection(getSummaryPreamble(), content) + content = appendSection(content, getContinueInstruction()) + + processed = newTypedSummaryMessage[M](content) + + return processed, nil +} + +type replaceUserMessagesInSummaryParams[M adk.MessageType] struct { + contextMsgs []M + summaryText string +} + +func replaceUserMessagesInSummary[M adk.MessageType](ctx context.Context, p *replaceUserMessagesInSummaryParams[M]) (string, error) { + var userMsgs []M + var hasUserMsgs bool + for _, msg := range p.contextMsgs { + if isInternalUserMessage(msg) { + continue + } + if isUserRole(msg) { + hasUserMsgs = true + userMsgs = append(userMsgs, msg) + } + } + + if !hasUserMsgs { + return p.summaryText, nil + } + + var selected []M + if len(userMsgs) == 1 { + selected = userMsgs + } else { + var totalTokens int + for i := len(userMsgs) - 1; i >= 0; i-- { + msg := userMsgs[i] + + tokens, err := defaultTypedTokenCounter(ctx, &TypedTokenCounterInput[M]{ + Messages: []M{msg}, + }) + if err != nil { + return "", fmt.Errorf("failed to count tokens: %w", err) + } + + remaining := preserveUserMsgsMaxTokens - totalTokens + if tokens <= remaining { + totalTokens += tokens + selected = append(selected, msg) + continue + } + + trimmedMsg := defaultTypedTrimUserMessage(msg, remaining) + var zero M + if any(trimmedMsg) != any(zero) { + selected = append(selected, trimmedMsg) + } + + break + } + + for i, j := 0, len(selected)-1; i < j; i, j = i+1, j-1 { + selected[i], selected[j] = selected[j], selected[i] + } + } + + var msgLines []string + for _, msg := range selected { + text := getUserMsgTextContent(msg) + if text != "" { + msgLines = append(msgLines, " - "+text) + } + } + userMsgsText := strings.Join(msgLines, "\n") + + lastMatch := findLastMatch(allUserMessagesTagRegex, p.summaryText) + if lastMatch == nil { + return p.summaryText, nil + } + + var replacement string + if len(selected) < len(userMsgs) { + replacement = "\n" + getUserMessagesReplacedNote() + "\n" + userMsgsText + "\n" + } else { + replacement = "\n" + userMsgsText + "\n" + } + + newSummaryText := p.summaryText[:lastMatch[0]] + replacement + p.summaryText[lastMatch[1]:] + + return newSummaryText, nil +} + +func findLastMatch(re *regexp.Regexp, s string) []int { + matches := re.FindAllStringIndex(s, -1) + if len(matches) == 0 { + return nil + } + return matches[len(matches)-1] +} + +func appendSection(base, section string) string { + if base == "" { + return section + } + if section == "" { + return base + } + return base + "\n\n" + section +} + +func (m *TypedMiddleware[M]) generateAndEmit(ctx context.Context, chatModel model.BaseModel[M], input []M, + opts []model.Option, attempt int, phase GenerateSummaryPhase) (M, error) { + + resp, err := chatModel.Generate(ctx, input, opts...) + if emitErr := m.emitGenerateSummaryEvent(ctx, attempt, phase, resp, err); emitErr != nil { + var zero M + return zero, emitErr + } + return resp, err +} + +func (m *TypedMiddleware[M]) generateWithRetry(ctx context.Context, chatModel model.BaseModel[M], input []M, + opts []model.Option, retryCfg *TypedRetryConfig[M]) (M, error) { + + const defaultMaxRetries = 3 + + if retryCfg == nil { + return m.generateAndEmit(ctx, chatModel, input, opts, 1, GenerateSummaryPhasePrimary) + } + + shouldRetry := retryCfg.ShouldRetry + if shouldRetry == nil { + shouldRetry = defaultTypedShouldRetry[M] + } + backoffFunc := retryCfg.BackoffFunc + if backoffFunc == nil { + backoffFunc = defaultTypedBackoffFunc[M] + } + + maxRetries := defaultMaxRetries + if retryCfg.MaxRetries != nil { + maxRetries = *retryCfg.MaxRetries + } + totalAttempts := maxRetries + 1 + + var ( + lastModelResp M + lastErr error + ) + for attempt := 1; attempt <= totalAttempts; attempt++ { + resp, err := m.generateAndEmit(ctx, chatModel, input, opts, attempt, GenerateSummaryPhasePrimary) + if !shouldRetry(ctx, resp, err) { + return resp, err + } + + lastModelResp = resp + lastErr = err + if attempt < totalAttempts { + select { + case <-time.After(backoffFunc(ctx, attempt, resp, err)): + case <-ctx.Done(): + var zero M + return zero, ctx.Err() + } + } + } + + if maxRetries > 0 { + return lastModelResp, fmt.Errorf("exceeds max retries: %w", lastErr) + } + + return lastModelResp, lastErr +} + +func truncateTextByChars(text string) string { + const maxRunes = 2000 + + if text == "" { + return "" + } + + if utf8.RuneCountInString(text) <= maxRunes { + return text + } + + halfRunes := maxRunes / 2 + runes := []rune(text) + totalRunes := len(runes) + + prefix := string(runes[:halfRunes]) + suffix := string(runes[totalRunes-halfRunes:]) + removedChars := totalRunes - maxRunes + + marker := fmt.Sprintf(getTruncatedMarkerFormat(), removedChars) + + return prefix + marker + suffix +} + +func (c *TypedConfig[M]) check() error { + if c == nil { + return fmt.Errorf("config is required") + } + if c.Model == nil { + return fmt.Errorf("model is required") + } + if c.Trigger != nil { + if err := c.Trigger.check(); err != nil { + return err + } + } + if c.Retry != nil { + if err := c.Retry.check(); err != nil { + return err + } + } + if c.Failover != nil { + if err := c.Failover.check(); err != nil { + return err + } + } + return nil +} + +func (c *TypedRetryConfig[M]) check() error { + if c.MaxRetries != nil && *c.MaxRetries < 0 { + return fmt.Errorf("retry.MaxRetries must be non-negative") + } + return nil +} + +func (c *TypedFailoverConfig[M]) check() error { + if c.MaxRetries != nil && *c.MaxRetries < 0 { + return fmt.Errorf("failover.MaxRetries must be non-negative") + } + return nil +} + +func (c *TriggerCondition) check() error { + if c.ContextTokens < 0 { + return fmt.Errorf("contextTokens must be non-negative") + } + if c.ContextMessages < 0 { + return fmt.Errorf("contextMessages must be non-negative") + } + if c.ContextTokens == 0 && c.ContextMessages == 0 { + return fmt.Errorf("at least one of contextTokens or contextMessages must be non-negative") + } + return nil +} + +// ============================================================================ +// Generic helper functions +// ============================================================================ + +func isSystemRole[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.System + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeSystem + } + panic("unreachable") +} + +func isUserRole[M adk.MessageType](msg M) bool { + switch m := any(msg).(type) { + case *schema.Message: + return m.Role == schema.User + case *schema.AgenticMessage: + return m.Role == schema.AgenticRoleTypeUser + } + panic("unreachable") +} + +func messageUserTextContent(m *schema.Message) string { + if m == nil { + return "" + } + var parts []string + for _, part := range m.UserInputMultiContent { + if part.Type == schema.ChatMessagePartTypeText && part.Text != "" { + parts = append(parts, part.Text) + } + } + if len(parts) > 0 { + return strings.Join(parts, "\n") + } + return m.Content +} + +func getUserMsgTextContent[M adk.MessageType](msg M) string { + switch m := any(msg).(type) { + case *schema.Message: + return messageUserTextContent(m) + + case *schema.AgenticMessage: + if m == nil { + return "" + } + var parts []string + for _, block := range m.ContentBlocks { + if block == nil { + continue + } + if block.UserInputText != nil { + parts = append(parts, block.UserInputText.Text) + } + } + return strings.Join(parts, "\n") + + default: + panic("unreachable") + } +} + +const multimodalTokenEstimate = 2000 + +func estimateMessageTokens(msg *schema.Message) int { + if msg == nil { + return 0 + } + var totalLen int + var multimodalTokens int + + if msg.Role == schema.Assistant { + if len(msg.AssistantGenMultiContent) > 0 { + hasReasoning := false + for _, part := range msg.AssistantGenMultiContent { + switch part.Type { + case schema.ChatMessagePartTypeText: + totalLen += len(part.Text) + case schema.ChatMessagePartTypeReasoning: + hasReasoning = true + if part.Reasoning != nil { + totalLen += len(part.Reasoning.Text) + } + case schema.ChatMessagePartTypeImageURL, schema.ChatMessagePartTypeAudioURL, + schema.ChatMessagePartTypeVideoURL, schema.ChatMessagePartTypeFileURL: + multimodalTokens += multimodalTokenEstimate + } + } + if !hasReasoning { + totalLen += len(msg.ReasoningContent) + } + } else { + totalLen += len(msg.Content) + len(msg.ReasoningContent) + } + for _, tc := range msg.ToolCalls { + totalLen += len(tc.Function.Name) + len(tc.Function.Arguments) + } + } else { + if len(msg.UserInputMultiContent) > 0 { + for _, part := range msg.UserInputMultiContent { + switch part.Type { + case schema.ChatMessagePartTypeText: + totalLen += len(part.Text) + case schema.ChatMessagePartTypeToolSearchResult: + if part.ToolSearchResult != nil { + for _, tl := range part.ToolSearchResult.Tools { + totalLen += len(tl.Name) + len(tl.Desc) + if b, err := sonic.Marshal(tl.ParamsOneOf); err == nil { + totalLen += len(b) + } + } + } + case schema.ChatMessagePartTypeImageURL, schema.ChatMessagePartTypeAudioURL, + schema.ChatMessagePartTypeVideoURL, schema.ChatMessagePartTypeFileURL: + multimodalTokens += multimodalTokenEstimate + } + } + } else { + totalLen += len(msg.Content) + } + } + + return estimateTokenCount(totalLen) + multimodalTokens +} + +func estimateAgenticMessageTokens(msg *schema.AgenticMessage) int { + if msg == nil { + return 0 + } + var totalLen int + var multimodalTokens int + + if msg.Role == schema.AgenticRoleTypeAssistant { + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + switch block.Type { + case schema.ContentBlockTypeAssistantGenText: + totalLen += len(block.AssistantGenText.Text) + case schema.ContentBlockTypeFunctionToolCall: + totalLen += len(block.FunctionToolCall.Name) + len(block.FunctionToolCall.Arguments) + case schema.ContentBlockTypeReasoning: + totalLen += len(block.Reasoning.Text) + case schema.ContentBlockTypeAssistantGenImage, schema.ContentBlockTypeAssistantGenAudio, + schema.ContentBlockTypeAssistantGenVideo: + multimodalTokens += multimodalTokenEstimate + } + } + } else { + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + switch block.Type { + case schema.ContentBlockTypeUserInputText: + totalLen += len(block.UserInputText.Text) + case schema.ContentBlockTypeFunctionToolResult: + for _, cb := range block.FunctionToolResult.Content { + if cb == nil { + continue + } + switch cb.Type { + case schema.FunctionToolResultContentBlockTypeText: + if cb.Text != nil { + totalLen += len(cb.Text.Text) + } + case schema.FunctionToolResultContentBlockTypeImage, schema.FunctionToolResultContentBlockTypeAudio, + schema.FunctionToolResultContentBlockTypeVideo, schema.FunctionToolResultContentBlockTypeFile: + multimodalTokens += multimodalTokenEstimate + } + } + case schema.ContentBlockTypeToolSearchResult: + if block.ToolSearchFunctionToolResult != nil && block.ToolSearchFunctionToolResult.Result != nil { + for _, tl := range block.ToolSearchFunctionToolResult.Result.Tools { + totalLen += len(tl.Name) + len(tl.Desc) + if b, err := sonic.Marshal(tl.ParamsOneOf); err == nil { + totalLen += len(b) + } + } + } + case schema.ContentBlockTypeUserInputImage, schema.ContentBlockTypeUserInputFile, + schema.ContentBlockTypeUserInputAudio, schema.ContentBlockTypeUserInputVideo: + multimodalTokens += multimodalTokenEstimate + } + } + } + + return estimateTokenCount(totalLen) + multimodalTokens +} + +func getMsgExtra[M adk.MessageType](msg M) map[string]any { + switch m := any(msg).(type) { + case *schema.Message: + return m.Extra + case *schema.AgenticMessage: + return m.Extra + default: + panic("unreachable") + } +} + +func setMsgExtra[M adk.MessageType](msg M, key string, value any) { + switch m := any(msg).(type) { + case *schema.Message: + if m.Extra == nil { + m.Extra = map[string]any{} + } + m.Extra[key] = value + case *schema.AgenticMessage: + if m.Extra == nil { + m.Extra = map[string]any{} + } + m.Extra[key] = value + } +} + +func makeSystemMsg[M adk.MessageType](text string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.SystemMessage(text)).(M) + case *schema.AgenticMessage: + return any(schema.SystemAgenticMessage(text)).(M) + default: + panic("unreachable") + } +} + +func makeUserMsg[M adk.MessageType](text string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(text)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(text)).(M) + default: + panic("unreachable") + } +} + +func newTypedSummaryMessage[M adk.MessageType](content string) M { + msg := makeUserMsg[M](content) + setMsgExtra(msg, extraKeyContentType, string(contentTypeSummary)) + return msg +} + +func typedGetContentType[M adk.MessageType](msg M) summarizationContentType { + extra := getMsgExtra(msg) + if extra == nil { + return "" + } + ct, ok := extra[extraKeyContentType].(string) + if !ok { + return "" + } + return summarizationContentType(ct) +} + +func isInternalUserMessage[M adk.MessageType](msg M) bool { + return typedGetContentType(msg) == contentTypeSummary || isPreservedMessage(msg) +} + +func isPreservedMessage[M adk.MessageType](msg M) bool { + return typedGetContentType(msg) == contentTypeSkills +} + +func typedShouldFailover[M adk.MessageType](ctx context.Context, cfg *TypedFailoverConfig[M], resp M, err error) bool { + if cfg == nil { + return false + } + if cfg.ShouldFailover == nil { + return err != nil + } + return cfg.ShouldFailover(ctx, resp, err) +} + +func defaultTypedShouldRetry[M adk.MessageType](_ context.Context, _ M, err error) bool { + return err != nil +} + +func defaultTypedBackoffFunc[M adk.MessageType](_ context.Context, attempt int, _ M, _ error) time.Duration { + return defaultBackoffDuration(attempt) +} + +func defaultBackoffDuration(attempt int) time.Duration { + const ( + baseDelay = time.Second + maxDelay = 10 * time.Second + ) + + if attempt <= 0 { + return baseDelay + } + + if attempt > 7 { + return maxDelay + time.Duration(rand.Int63n(int64(maxDelay/2))) + } + + delay := baseDelay * time.Duration(1< maxDelay { + delay = maxDelay + } + + jitter := time.Duration(rand.Int63n(int64(delay / 2))) + + return delay + jitter +} + +func defaultTypedTrimUserMessage[M adk.MessageType](msg M, remainingTokens int) M { + var zero M + if remainingTokens <= 0 { + return zero + } + + textContent := getUserMsgTextContent(msg) + if len(textContent) == 0 { + return zero + } + + trimmed := truncateTextByChars(textContent) + if trimmed == "" { + return zero + } + + return makeUserMsg[M](trimmed) +} diff --git a/adk/middlewares/summarization/summarization_test.go b/adk/middlewares/summarization/summarization_test.go new file mode 100644 index 0000000..d70f396 --- /dev/null +++ b/adk/middlewares/summarization/summarization_test.go @@ -0,0 +1,2245 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package summarization + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +func intPtr(v int) *int { + return &v +} + +func TestNew(t *testing.T) { + ctx := context.Background() + + t.Run("valid config", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + cfg := &Config{ + Model: cm, + } + + mw, err := New(ctx, cfg) + assert.NoError(t, err) + assert.NotNil(t, mw) + }) + + t.Run("nil config returns error", func(t *testing.T) { + mw, err := New(ctx, nil) + assert.Error(t, err) + assert.Nil(t, mw) + }) + + t.Run("nil model returns error", func(t *testing.T) { + mw, err := New(ctx, &Config{}) + assert.Error(t, err) + assert.Nil(t, mw) + }) +} + +func TestMiddlewareBeforeModelRewriteState(t *testing.T) { + ctx := context.Background() + mtx := &adk.ModelContext{} + + t.Run("no summarization when under threshold", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 1000}, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage("hello"), + schema.AssistantMessage("hi", nil), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 2) + assert.Equal(t, "hello", newState.Messages[0].Content) + }) + + t.Run("summarization triggered when over threshold", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary content", + }, nil).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + schema.AssistantMessage(strings.Repeat("b", 100), nil), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 1) + assert.Equal(t, schema.User, newState.Messages[0].Role) + }) + + t.Run("preserves system messages after summarization", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + for i, msg := range msgs { + if i == 0 { + assert.Equal(t, schema.System, msg.Role) + } else { + assert.NotEqual(t, schema.System, msg.Role) + } + } + return &schema.Message{ + Role: schema.Assistant, + Content: "Summary content", + }, nil + }).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("You are a helpful assistant"), + schema.UserMessage(strings.Repeat("a", 100)), + schema.AssistantMessage(strings.Repeat("b", 100), nil), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 2) + assert.Equal(t, schema.System, newState.Messages[0].Role) + assert.Equal(t, "You are a helpful assistant", newState.Messages[0].Content) + assert.Equal(t, schema.User, newState.Messages[1].Role) + }) + + t.Run("preserves multiple system messages", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("System 1"), + schema.SystemMessage("System 2"), + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 3) + assert.Equal(t, schema.System, newState.Messages[0].Role) + assert.Equal(t, "System 1", newState.Messages[0].Content) + assert.Equal(t, schema.System, newState.Messages[1].Role) + assert.Equal(t, "System 2", newState.Messages[1].Content) + assert.Equal(t, schema.User, newState.Messages[2].Role) + }) + + t.Run("custom finalize function", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return []adk.Message{ + schema.SystemMessage("system prompt"), + summary, + }, nil + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 2) + assert.Equal(t, schema.System, newState.Messages[0].Role) + assert.Equal(t, "system prompt", newState.Messages[0].Content) + }) + + t.Run("retry succeeds after transient error", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + callCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + callCount++ + if callCount == 1 { + return nil, fmt.Errorf("transient error") + } + return &schema.Message{ + Role: schema.Assistant, + Content: "Summary after retry", + }, nil + }).Times(2) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + Retry: &RetryConfig{ + MaxRetries: intPtr(2), + BackoffFunc: func(_ context.Context, _ int, _ adk.Message, _ error) time.Duration { return 0 }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 1) + assert.Equal(t, 2, callCount) + }) + + t.Run("retry uses default max retries when MaxRetries is nil", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + callCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + callCount++ + return nil, fmt.Errorf("transient error") + }).Times(4) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 10}, + Retry: &RetryConfig{ + BackoffFunc: func(_ context.Context, _ int, _ adk.Message, _ error) time.Duration { return 0 }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, _, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to generate summary") + assert.Equal(t, 4, callCount) + }) + + t.Run("failover succeeds after primary failure", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + failover := mockModel.NewMockBaseChatModel(ctrl) + + primary.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("primary error")).Times(1) + failover.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...interface{}) (*schema.Message, error) { + assert.Len(t, msgs, 1) + assert.Equal(t, "failover input", msgs[0].Content) + return &schema.Message{ + Role: schema.Assistant, + Content: "Summary from failover", + }, nil + }).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: primary, + Trigger: &TriggerCondition{ContextTokens: 10}, + Failover: &FailoverConfig{ + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + assert.Equal(t, 1, failoverCtx.Attempt) + assert.Equal(t, schema.System, failoverCtx.SystemInstruction.Role) + assert.Equal(t, schema.User, failoverCtx.UserInstruction.Role) + assert.Len(t, failoverCtx.OriginalMessages, 1) + assert.Nil(t, failoverCtx.LastModelResponse) + assert.EqualError(t, failoverCtx.LastErr, "primary error") + return failover, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 1) + assert.Equal(t, schema.User, newState.Messages[0].Role) + }) + + t.Run("failover context last err is retry exhausted error when retries exhausted", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + failover := mockModel.NewMockBaseChatModel(ctrl) + + primary.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("primary error")).Times(2) + failover.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary from failover", + }, nil).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: primary, + Trigger: &TriggerCondition{ContextTokens: 10}, + Retry: &RetryConfig{ + MaxRetries: intPtr(1), + BackoffFunc: func(_ context.Context, _ int, _ adk.Message, _ error) time.Duration { return 0 }, + }, + Failover: &FailoverConfig{ + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + assert.ErrorContains(t, failoverCtx.LastErr, "exceeds max retries") + return failover, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, _, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + }) + + t.Run("returns failover exhausted error when failover model fails", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + failover := mockModel.NewMockBaseChatModel(ctrl) + + primary.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("primary error")).Times(1) + failover.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("failover error")).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: primary, + Trigger: &TriggerCondition{ContextTokens: 10}, + Failover: &FailoverConfig{ + MaxRetries: intPtr(1), + BackoffFunc: func(_ context.Context, _ int, _ adk.Message, _ error) time.Duration { return 0 }, + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return failover, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, _, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.Error(t, err) + assert.ErrorContains(t, err, "exceeds max failover attempts") + }) + + t.Run("failover retries with max retries and succeeds on second attempt", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + failover1 := mockModel.NewMockBaseChatModel(ctrl) + failover2 := mockModel.NewMockBaseChatModel(ctrl) + + primary.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("primary error")).Times(1) + failover1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("failover error 1")).Times(1) + failover2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary from second failover", + }, nil).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: primary, + Trigger: &TriggerCondition{ContextTokens: 10}, + Failover: &FailoverConfig{ + MaxRetries: intPtr(2), + BackoffFunc: func(_ context.Context, _ int, _ adk.Message, _ error) time.Duration { return 0 }, + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + if failoverCtx.Attempt == 1 { + assert.EqualError(t, failoverCtx.LastErr, "primary error") + return failover1, []*schema.Message{schema.UserMessage("failover input 1")}, nil + } + assert.Equal(t, 2, failoverCtx.Attempt) + assert.EqualError(t, failoverCtx.LastErr, "failover error 1") + return failover2, []*schema.Message{schema.UserMessage("failover input 2")}, nil + }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + assert.Len(t, newState.Messages, 1) + }) + + t.Run("failover context carries generate resp as last output message", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + failover := mockModel.NewMockBaseChatModel(ctrl) + + primary.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "partial output", + }, fmt.Errorf("primary error")).Times(1) + failover.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary from failover", + }, nil).Times(1) + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: primary, + Trigger: &TriggerCondition{ContextTokens: 10}, + Failover: &FailoverConfig{ + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + if assert.NotNil(t, failoverCtx.LastModelResponse) { + assert.Equal(t, "partial output", failoverCtx.LastModelResponse.Content) + } + return failover, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + }, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + _, _, err := mw.BeforeModelRewriteState(ctx, state, mtx) + assert.NoError(t, err) + }) + +} + +func TestMiddlewareShouldSummarize(t *testing.T) { + ctx := context.Background() + + t.Run("returns true when over messages threshold", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Trigger: &TriggerCondition{ContextMessages: 1}, + }, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{ + schema.UserMessage("msg1"), + schema.UserMessage("msg2"), + }, + } + + triggered, err := mw.shouldSummarize(ctx, input) + assert.NoError(t, err) + assert.True(t, triggered) + }) + + t.Run("returns false when under messages threshold", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Trigger: &TriggerCondition{ + ContextMessages: 3, + ContextTokens: 1000, + }, + }, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{ + schema.UserMessage("msg1"), + schema.UserMessage("msg2"), + }, + } + + triggered, err := mw.shouldSummarize(ctx, input) + assert.NoError(t, err) + assert.False(t, triggered) + }) + + t.Run("returns true when over threshold", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Trigger: &TriggerCondition{ContextTokens: 10}, + }, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{ + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + triggered, err := mw.shouldSummarize(ctx, input) + assert.NoError(t, err) + assert.True(t, triggered) + }) + + t.Run("returns false when under threshold", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Trigger: &TriggerCondition{ContextTokens: 1000}, + }, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{ + schema.UserMessage("short message"), + }, + } + + triggered, err := mw.shouldSummarize(ctx, input) + assert.NoError(t, err) + assert.False(t, triggered) + }) + + t.Run("uses default threshold when trigger is nil", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{}, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{ + schema.UserMessage("short message"), + }, + } + + triggered, err := mw.shouldSummarize(ctx, input) + assert.NoError(t, err) + assert.False(t, triggered) + }) +} + +func TestMiddlewareCountTokens(t *testing.T) { + ctx := context.Background() + + t.Run("uses custom token counter", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + TokenCounter: func(ctx context.Context, input *TokenCounterInput) (int, error) { + return 42, nil + }, + }, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{schema.UserMessage("test")}, + } + tokens, err := mw.countTokens(ctx, input) + assert.NoError(t, err) + assert.Equal(t, 42, tokens) + }) + + t.Run("uses default token counter when nil", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{}, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{schema.UserMessage("test")}, + } + tokens, err := mw.countTokens(ctx, input) + assert.NoError(t, err) + assert.Greater(t, tokens, 0) + }) + + t.Run("custom token counter error", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + TokenCounter: func(ctx context.Context, input *TokenCounterInput) (int, error) { + return 0, errors.New("token count error") + }, + }, + } + + input := &TokenCounterInput{ + Messages: []adk.Message{schema.UserMessage("test")}, + } + _, err := mw.countTokens(ctx, input) + assert.Error(t, err) + }) +} + +func TestGetUserMsgTextContent(t *testing.T) { + t.Run("Message extracts from Content field", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + Content: "hello world", + } + assert.Equal(t, "hello world", getUserMsgTextContent(msg)) + }) + + t.Run("Message extracts from UserInputMultiContent", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "part1"}, + {Type: schema.ChatMessagePartTypeText, Text: "part2"}, + }, + } + assert.Equal(t, "part1\npart2", getUserMsgTextContent(msg)) + }) + + t.Run("Message prefers UserInputMultiContent over Content", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + Content: "content field", + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "multi content"}, + }, + } + assert.Equal(t, "multi content", getUserMsgTextContent(msg)) + }) + + t.Run("Message nil returns empty", func(t *testing.T) { + assert.Equal(t, "", getUserMsgTextContent[*schema.Message](nil)) + }) + + t.Run("AgenticMessage extracts UserInputText", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + {UserInputText: &schema.UserInputText{Text: "user input"}}, + }, + } + assert.Equal(t, "user input", getUserMsgTextContent(msg)) + }) + + t.Run("AgenticMessage nil returns empty", func(t *testing.T) { + assert.Equal(t, "", getUserMsgTextContent[*schema.AgenticMessage](nil)) + }) +} + +func TestTruncateTextByChars(t *testing.T) { + t.Run("returns empty for empty string", func(t *testing.T) { + result := truncateTextByChars("") + assert.Equal(t, "", result) + }) + + t.Run("returns original if under limit", func(t *testing.T) { + result := truncateTextByChars("short") + assert.Equal(t, "short", result) + }) + + t.Run("truncates long text", func(t *testing.T) { + longText := strings.Repeat("a", 3000) + result := truncateTextByChars(longText) + assert.Less(t, len(result), len(longText)) + assert.Contains(t, result, "truncated") + }) + + t.Run("preserves prefix and suffix", func(t *testing.T) { + longText := strings.Repeat("a", 1000) + strings.Repeat("b", 1000) + strings.Repeat("c", 1000) + result := truncateTextByChars(longText) + assert.True(t, strings.HasPrefix(result, strings.Repeat("a", 1000))) + assert.True(t, strings.HasSuffix(result, strings.Repeat("c", 1000))) + }) +} + +func TestAppendSection(t *testing.T) { + tests := []struct { + name string + base string + section string + expected string + }{ + { + name: "both empty", + base: "", + section: "", + expected: "", + }, + { + name: "base empty", + base: "", + section: "section", + expected: "section", + }, + { + name: "section empty", + base: "base", + section: "", + expected: "base", + }, + { + name: "both non-empty", + base: "base", + section: "section", + expected: "base\n\nsection", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := appendSection(tt.base, tt.section) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestAllUserMessagesTagRegex(t *testing.T) { + t.Run("matches tag", func(t *testing.T) { + text := ` + - msg1 + - msg2 +` + assert.True(t, allUserMessagesTagRegex.MatchString(text)) + }) + + t.Run("replaces tag content", func(t *testing.T) { + text := `before + + - old msg + +after` + replacement := "\n - new msg\n" + result := allUserMessagesTagRegex.ReplaceAllString(text, replacement) + assert.Contains(t, result, "new msg") + assert.NotContains(t, result, "old msg") + assert.Contains(t, result, "before") + assert.Contains(t, result, "after") + }) +} + +func TestConfigCheck(t *testing.T) { + t.Run("nil config", func(t *testing.T) { + var c *Config + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "config is required") + }) + + t.Run("nil model", func(t *testing.T) { + c := &Config{} + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "model is required") + }) + + t.Run("valid config", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + } + err := c.check() + assert.NoError(t, err) + }) + + t.Run("invalid trigger max tokens", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: -1}, + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "must be non-negative") + }) + + t.Run("invalid trigger max messages", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextMessages: -1}, + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "must be non-negative") + }) + + t.Run("both trigger conditions are zero", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Trigger: &TriggerCondition{ContextTokens: 0, ContextMessages: 0}, + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "must be non-negative") + }) + + t.Run("negative retry max retries", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Retry: &RetryConfig{MaxRetries: intPtr(-1)}, + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "retry.MaxRetries must be non-negative") + }) + + t.Run("failover getFailoverModel is optional", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Failover: &FailoverConfig{}, + } + err := c.check() + assert.NoError(t, err) + }) + + t.Run("failover max retries accepts int value", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + failover := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Failover: &FailoverConfig{ + MaxRetries: intPtr(1), + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return failover, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + } + err := c.check() + assert.NoError(t, err) + }) + + t.Run("failover max retries must be non-negative", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + failover := mockModel.NewMockBaseChatModel(ctrl) + + c := &Config{ + Model: cm, + Failover: &FailoverConfig{ + MaxRetries: intPtr(-1), + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return failover, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + } + err := c.check() + assert.Error(t, err) + assert.Contains(t, err.Error(), "failover.MaxRetries must be non-negative") + }) +} + +func TestSetGetContentType(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + Content: "test", + } + + setMsgExtra(msg, extraKeyContentType, string(contentTypeSummary)) + + ct := typedGetContentType(msg) + assert.Equal(t, contentTypeSummary, ct) +} + +func TestSetGetExtra(t *testing.T) { + t.Run("set and get", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + Content: "test", + } + + setMsgExtra(msg, "key", "value") + + extra := getMsgExtra(msg) + v, ok := extra["key"].(string) + assert.True(t, ok) + assert.Equal(t, "value", v) + }) + + t.Run("get non-existent key", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + Content: "test", + } + + extra := getMsgExtra(msg) + assert.Nil(t, extra) + }) +} + +func TestMiddlewareBuildSummarizationModelInput(t *testing.T) { + ctx := context.Background() + + t.Run("message structure", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{}, + } + + testMsg := []adk.Message{schema.UserMessage("test")} + input, err := mw.buildSummarizationModelInput(ctx, testMsg, testMsg) + assert.NoError(t, err) + assert.GreaterOrEqual(t, len(input), 3) + assert.Equal(t, schema.System, input[0].Role) + assert.Equal(t, schema.User, input[len(input)-1].Role) + }) + + t.Run("uses context messages", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{}, + } + + contextMsgs := []adk.Message{ + schema.UserMessage("context message"), + } + input, err := mw.buildSummarizationModelInput(ctx, contextMsgs, contextMsgs) + assert.NoError(t, err) + + found := false + for _, msg := range input { + if msg.Content == "context message" { + found = true + break + } + } + assert.True(t, found, "should contain context message") + }) + + t.Run("uses GenModelInput", func(t *testing.T) { + expectedInput := []adk.Message{ + schema.UserMessage("custom input"), + } + + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + GenModelInput: func(ctx context.Context, defaultSystemInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) { + return expectedInput, nil + }, + }, + } + + testMsg := []adk.Message{schema.UserMessage("test")} + input, err := mw.buildSummarizationModelInput(ctx, testMsg, testMsg) + assert.NoError(t, err) + assert.Len(t, input, 1) + assert.Equal(t, "custom input", input[0].Content) + }) + + t.Run("GenModelInput error", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + GenModelInput: func(ctx context.Context, defaultSystemInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) { + return nil, errors.New("gen input error") + }, + }, + } + + testMsg := []adk.Message{schema.UserMessage("test")} + _, err := mw.buildSummarizationModelInput(ctx, testMsg, testMsg) + assert.Error(t, err) + assert.Contains(t, err.Error(), "gen input error") + }) + + t.Run("uses custom instruction", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + UserInstruction: "custom instruction", + }, + } + + testMsg := []adk.Message{schema.UserMessage("test")} + input, err := mw.buildSummarizationModelInput(ctx, testMsg, testMsg) + assert.NoError(t, err) + + lastMsg := input[len(input)-1] + assert.Equal(t, schema.User, lastMsg.Role) + assert.Contains(t, lastMsg.Content, "custom instruction") + }) +} + +func TestMiddlewareSummarize(t *testing.T) { + ctx := context.Background() + + t.Run("generates summary", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "summary", + }, nil).Times(1) + + input := []adk.Message{schema.UserMessage("test")} + resp, err := cm.Generate(ctx, input) + assert.NoError(t, err) + assert.NotNil(t, resp) + summary := newTypedSummaryMessage[*schema.Message](resp.Content) + assert.NotNil(t, summary) + assert.Equal(t, "summary", summary.Content) + }) + + t.Run("model generate error", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, errors.New("generate error")).Times(1) + + input := []adk.Message{schema.UserMessage("test")} + _, err := cm.Generate(ctx, input) + assert.Error(t, err) + }) +} + +func TestMiddlewareGenerateWithRetry(t *testing.T) { + ctx := context.Background() + + t.Run("retries until success", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{}, + } + + callCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(context.Context, []*schema.Message, ...any) (*schema.Message, error) { + callCount++ + if callCount == 1 { + return schema.AssistantMessage("partial output", nil), errors.New("transient error") + } + return schema.AssistantMessage("final summary", nil), nil + }).Times(2) + + resp, err := mw.generateWithRetry(ctx, cm, []adk.Message{schema.UserMessage("test")}, nil, &RetryConfig{}) + + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.Equal(t, "final summary", resp.Content) + } + }) + + t.Run("delegates to generateAndEmit without retry config", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{}, + } + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("partial output", nil), errors.New("generate error")).Times(1) + + resp, err := mw.generateWithRetry(ctx, cm, []adk.Message{schema.UserMessage("test")}, nil, nil) + + assert.EqualError(t, err, "generate error") + if assert.NotNil(t, resp) { + assert.Equal(t, "partial output", resp.Content) + } + }) +} + +func TestPopulateUserMessagesInternal(t *testing.T) { + ctx := context.Background() + + t.Run("replaces user messages section", func(t *testing.T) { + msgs := []adk.Message{ + schema.UserMessage("msg1"), + schema.AssistantMessage("response1", nil), + schema.UserMessage("msg2"), + } + + summary := `1. Primary Request: + test + +6. All user messages: + + - [old message] + + +7. Pending Tasks: + - task1` + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: msgs, + summaryText: summary, + }) + assert.NoError(t, err) + assert.Contains(t, result, "msg1") + assert.Contains(t, result, "msg2") + assert.NotContains(t, result, "old message") + assert.Contains(t, result, "7. Pending Tasks:") + }) + + t.Run("returns original if no matching sections", func(t *testing.T) { + msgs := []adk.Message{ + schema.UserMessage("test"), + } + + summary := "summary without sections" + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: msgs, + summaryText: summary, + }) + assert.NoError(t, err) + assert.Equal(t, summary, result) + }) + + t.Run("skips summary messages", func(t *testing.T) { + summaryMsg := &schema.Message{ + Role: schema.User, + Content: "summary", + } + setMsgExtra(summaryMsg, extraKeyContentType, string(contentTypeSummary)) + + msgs := []adk.Message{ + summaryMsg, + schema.UserMessage("regular message"), + } + + summary := `6. All user messages: + + - [old] + + +7. Pending Tasks: + - task` + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: msgs, + summaryText: summary, + }) + assert.NoError(t, err) + assert.Contains(t, result, "regular message") + assert.NotContains(t, result, " - summary") + }) + + t.Run("returns original if empty user messages", func(t *testing.T) { + msgs := []adk.Message{ + schema.AssistantMessage("response", nil), + } + + summary := `6. All user messages: + - [old] + +7. Pending Tasks: + - task` + + result, err := replaceUserMessagesInSummary(ctx, &replaceUserMessagesInSummaryParams[*schema.Message]{ + contextMsgs: msgs, + summaryText: summary, + }) + assert.NoError(t, err) + assert.Equal(t, summary, result) + }) +} + +func TestAllUserMessagesTagRegexMatch(t *testing.T) { + t.Run("matches xml tag", func(t *testing.T) { + text := "\n - msg\n" + assert.True(t, allUserMessagesTagRegex.MatchString(text)) + }) + + t.Run("does not match without tag", func(t *testing.T) { + text := "6. All user messages:\n - msg" + assert.False(t, allUserMessagesTagRegex.MatchString(text)) + }) +} + +func TestDefaultTrimUserMessage(t *testing.T) { + t.Run("returns nil for zero remaining tokens", func(t *testing.T) { + msg := schema.UserMessage("test") + result := defaultTypedTrimUserMessage(msg, 0) + assert.Nil(t, result) + }) + + t.Run("returns nil for empty content", func(t *testing.T) { + msg := schema.UserMessage("") + result := defaultTypedTrimUserMessage(msg, 100) + assert.Nil(t, result) + }) + + t.Run("trims long message", func(t *testing.T) { + longText := strings.Repeat("a", 3000) + msg := schema.UserMessage(longText) + result := defaultTypedTrimUserMessage(msg, 100) + assert.NotNil(t, result) + assert.Less(t, len(result.Content), len(longText)) + }) +} + +func TestDefaultTokenCounter(t *testing.T) { + ctx := context.Background() + + t.Run("counts tool tokens", func(t *testing.T) { + input := &TokenCounterInput{ + Messages: []adk.Message{}, + Tools: []*schema.ToolInfo{ + {Name: "test_tool", Desc: "description"}, + }, + } + count, err := defaultTypedTokenCounter(ctx, input) + assert.NoError(t, err) + assert.Greater(t, count, 0) + }) + + t.Run("reuses latest assistant total tokens as baseline", func(t *testing.T) { + input := &TokenCounterInput{ + Messages: []adk.Message{ + schema.UserMessage("earlier context"), + { + Role: schema.Assistant, + Content: "baseline", + ResponseMeta: &schema.ResponseMeta{ + Usage: &schema.TokenUsage{TotalTokens: 100}, + }, + }, + schema.UserMessage("later context"), + }, + } + + count, err := defaultTypedTokenCounter(ctx, input) + require.NoError(t, err) + assert.Equal(t, 100+estimateMessageTokens(schema.UserMessage("later context")), count) + }) +} + +func TestGetAssistantTotalTokens(t *testing.T) { + t.Run("returns zero for nil message", func(t *testing.T) { + assert.Zero(t, getAssistantTotalTokens[*schema.Message](nil)) + assert.Zero(t, getAssistantTotalTokens[*schema.AgenticMessage](nil)) + }) + + t.Run("reads total tokens from assistant messages only", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + ResponseMeta: &schema.ResponseMeta{ + Usage: &schema.TokenUsage{TotalTokens: 42}, + }, + } + assert.Equal(t, 42, getAssistantTotalTokens(msg)) + assert.Zero(t, getAssistantTotalTokens(schema.UserMessage("ignored"))) + }) + + t.Run("reads total tokens from agentic assistant messages only", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ResponseMeta: &schema.AgenticResponseMeta{ + TokenUsage: &schema.TokenUsage{TotalTokens: 64}, + }, + } + assert.Equal(t, 64, getAssistantTotalTokens(msg)) + assert.Zero(t, getAssistantTotalTokens(schema.UserAgenticMessage("ignored"))) + }) +} + +func TestEstimateMessageTokens(t *testing.T) { + t.Run("returns zero for nil message", func(t *testing.T) { + assert.Zero(t, estimateMessageTokens(nil)) + }) + + t.Run("counts assistant text reasoning and tool calls", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + ReasoningContent: "reason", + ToolCalls: []schema.ToolCall{ + { + Function: schema.FunctionCall{ + Name: "tool", + Arguments: `{"k":"v"}`, + }, + }, + }, + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "answer"}, + }, + } + + expectedLen := len("answer") + len("reason") + len("tool") + len(`{"k":"v"}`) + assert.Equal(t, estimateTokenCount(expectedLen), estimateMessageTokens(msg)) + }) + + t.Run("adds multimodal estimate for user content", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.User, + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "hello"}, + {Type: schema.ChatMessagePartTypeImageURL}, + }, + } + + assert.Equal(t, estimateTokenCount(len("hello"))+multimodalTokenEstimate, estimateMessageTokens(msg)) + }) +} + +func TestEstimateAgenticMessageTokens(t *testing.T) { + t.Run("returns zero for nil message", func(t *testing.T) { + assert.Zero(t, estimateAgenticMessageTokens(nil)) + }) + + t.Run("counts assistant blocks and multimodal outputs", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "answer"}), + schema.NewContentBlock(&schema.Reasoning{Text: "reason"}), + schema.NewContentBlock(&schema.FunctionToolCall{Name: "tool", Arguments: `{"k":"v"}`}), + schema.NewContentBlock(&schema.AssistantGenImage{}), + }, + } + + expectedLen := len("answer") + len("reason") + len("tool") + len(`{"k":"v"}`) + assert.Equal(t, estimateTokenCount(expectedLen)+multimodalTokenEstimate, estimateAgenticMessageTokens(msg)) + }) +} + +func TestPostProcessSummary(t *testing.T) { + ctx := context.Background() + + t.Run("with transcript path", func(t *testing.T) { + result, err := postProcessSummary(ctx, &postProcessSummaryParams[*schema.Message]{ + contextMsgs: []adk.Message{}, + summaryContent: "summary content", + transcriptPath: "/path/to/transcript.txt", + }) + assert.NoError(t, err) + assert.Contains(t, result.Content, "/path/to/transcript.txt") + assert.Contains(t, result.Content, getContinueInstruction()) + }) +} + +func TestEventHelpers(t *testing.T) { + ctx := context.Background() + + t.Run("emitEvent returns wrapped error outside execution context", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{cfg: &Config{}} + err := mw.emitEvent(ctx, &CustomizedAction{Type: ActionTypeBeforeSummarize}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to send internal event") + }) + + t.Run("emitGenerateSummaryEvent is skipped when internal events are disabled", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{cfg: &Config{EmitInternalEvents: false}} + err := mw.emitGenerateSummaryEvent(ctx, 1, GenerateSummaryPhasePrimary, schema.AssistantMessage("ok", nil), nil) + assert.NoError(t, err) + }) + + t.Run("emitGenerateSummaryEvent returns wrapped error when enabled outside execution context", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{cfg: &Config{EmitInternalEvents: true}} + err := mw.emitGenerateSummaryEvent(ctx, 1, GenerateSummaryPhasePrimary, schema.AssistantMessage("ok", nil), nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to send internal event") + }) +} + +func TestGetFailoverModel(t *testing.T) { + ctx := context.Background() + defaultInput := []adk.Message{schema.UserMessage("default")} + fctx := &FailoverContext{Attempt: 1} + + t.Run("requires failover config", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{cfg: &Config{}} + mdl, input, err := mw.getFailoverModel(ctx, fctx, defaultInput) + assert.Nil(t, mdl) + assert.Nil(t, input) + assert.ErrorContains(t, err, "failover config is required") + }) + + t.Run("uses primary model and default input when callback is not provided", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Model: primary, + Failover: &FailoverConfig{}, + }, + } + + mdl, input, err := mw.getFailoverModel(ctx, fctx, defaultInput) + assert.NoError(t, err) + assert.Same(t, primary, mdl) + assert.Equal(t, defaultInput, input) + }) + + t.Run("wraps callback error", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Failover: &FailoverConfig{ + GetFailoverModel: func(context.Context, *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, errors.New("boom") + }, + }, + }, + } + + mdl, input, err := mw.getFailoverModel(ctx, fctx, defaultInput) + assert.Nil(t, mdl) + assert.Nil(t, input) + assert.ErrorContains(t, err, "failed to get failover model") + }) + + t.Run("requires non nil failover model", func(t *testing.T) { + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Failover: &FailoverConfig{ + GetFailoverModel: func(context.Context, *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return nil, []*schema.Message{schema.UserMessage("input")}, nil + }, + }, + }, + } + + mdl, input, err := mw.getFailoverModel(ctx, fctx, defaultInput) + assert.Nil(t, mdl) + assert.Nil(t, input) + assert.ErrorContains(t, err, "failover model is required") + }) + + t.Run("requires non empty failover input", func(t *testing.T) { + ctrl := gomock.NewController(t) + failoverModel := mockModel.NewMockBaseChatModel(ctrl) + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Failover: &FailoverConfig{ + GetFailoverModel: func(context.Context, *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return failoverModel, nil, nil + }, + }, + }, + } + + mdl, input, err := mw.getFailoverModel(ctx, fctx, defaultInput) + assert.Nil(t, mdl) + assert.Nil(t, input) + assert.ErrorContains(t, err, "failover model input messages are required") + }) + + t.Run("returns custom failover model and input", func(t *testing.T) { + ctrl := gomock.NewController(t) + failoverModel := mockModel.NewMockBaseChatModel(ctrl) + customInput := []*schema.Message{schema.UserMessage("custom")} + mw := &TypedMiddleware[*schema.Message]{ + cfg: &Config{ + Failover: &FailoverConfig{ + GetFailoverModel: func(context.Context, *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return failoverModel, customInput, nil + }, + }, + }, + } + + mdl, input, err := mw.getFailoverModel(ctx, fctx, defaultInput) + assert.NoError(t, err) + assert.Same(t, failoverModel, mdl) + if assert.Len(t, input, 1) { + assert.Equal(t, "custom", input[0].Content) + } + }) +} + +func TestHelperBranches(t *testing.T) { + t.Run("should failover branches", func(t *testing.T) { + assert.False(t, typedShouldFailover(context.Background(), (*FailoverConfig)(nil), nil, errors.New("x"))) + assert.False(t, typedShouldFailover(context.Background(), &FailoverConfig{}, nil, nil)) + assert.True(t, typedShouldFailover(context.Background(), &FailoverConfig{}, nil, errors.New("x"))) + + cfg := &FailoverConfig{ + ShouldFailover: func(ctx context.Context, resp adk.Message, err error) bool { + return resp != nil && err == nil + }, + } + assert.True(t, typedShouldFailover(context.Background(), cfg, schema.AssistantMessage("ok", nil), nil)) + }) + + t.Run("config check branches", func(t *testing.T) { + assert.ErrorContains(t, (&RetryConfig{MaxRetries: intPtr(-1)}).check(), "retry.MaxRetries must be non-negative") + assert.ErrorContains(t, (&FailoverConfig{MaxRetries: intPtr(-1)}).check(), "failover.MaxRetries must be non-negative") + assert.ErrorContains(t, (&TriggerCondition{}).check(), "at least one of contextTokens or contextMessages") + assert.ErrorContains(t, (&TriggerCondition{ContextTokens: -1}).check(), "contextTokens must be non-negative") + assert.ErrorContains(t, (&TriggerCondition{ContextMessages: -1, ContextTokens: 1}).check(), "contextMessages must be non-negative") + }) + + t.Run("default backoff branches", func(t *testing.T) { + assert.Equal(t, time.Second, defaultBackoffDuration(0)) + + delay := defaultBackoffDuration(8) + assert.GreaterOrEqual(t, delay, 10*time.Second) + assert.Less(t, delay, 15*time.Second) + }) + + t.Run("user messages replaced note is present", func(t *testing.T) { + note := getUserMessagesReplacedNote() + assert.NotEmpty(t, note) + assert.Contains(t, []string{userMessagesReplacedNote, userMessagesReplacedNoteZh}, note) + }) +} + +func TestSummarize(t *testing.T) { + ctx := context.Background() + + newMW := func(cfg *Config) *TypedMiddleware[*schema.Message] { + return &TypedMiddleware[*schema.Message]{ + cfg: cfg, + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } + } + + t.Run("basic summarization", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary content", + }, nil).Times(1) + + mw := newMW(&Config{Model: cm}) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("You are a helpful assistant"), + schema.UserMessage(strings.Repeat("a", 100)), + schema.AssistantMessage(strings.Repeat("b", 100), nil), + }, + } + + result, err := mw.Summarize(ctx, state) + assert.NoError(t, err) + assert.NotEmpty(t, result) + assert.Equal(t, schema.System, result[0].Role) + }) + + t.Run("model error propagates", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("model error")).Times(1) + + mw := newMW(&Config{Model: cm}) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.Error(t, err) + assert.Nil(t, result) + }) + + t.Run("retry works", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + + callCount := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...any) (*schema.Message, error) { + callCount++ + if callCount == 1 { + return nil, fmt.Errorf("transient error") + } + return &schema.Message{ + Role: schema.Assistant, + Content: "Summary after retry", + }, nil + }).Times(2) + + mw := newMW(&Config{ + Model: cm, + Retry: &RetryConfig{ + MaxRetries: intPtr(2), + BackoffFunc: func(_ context.Context, _ int, _ adk.Message, _ error) time.Duration { return 0 }, + }, + }) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.NoError(t, err) + assert.NotEmpty(t, result) + assert.Equal(t, 2, callCount) + }) + + t.Run("failover works", func(t *testing.T) { + ctrl := gomock.NewController(t) + primary := mockModel.NewMockBaseChatModel(ctrl) + failoverModel := mockModel.NewMockBaseChatModel(ctrl) + + primary.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("primary error")).Times(1) + failoverModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary from failover", + }, nil).Times(1) + + mw := newMW(&Config{ + Model: primary, + Failover: &FailoverConfig{ + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext) (model.BaseChatModel, []*schema.Message, error) { + return failoverModel, []*schema.Message{schema.UserMessage("failover input")}, nil + }, + }, + }) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.NoError(t, err) + assert.NotEmpty(t, result) + }) + + t.Run("callback is invoked", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + callbackCalled := false + mw := newMW(&Config{ + Model: cm, + Callback: func(ctx context.Context, before, after adk.ChatModelAgentState) error { + callbackCalled = true + assert.Len(t, before.Messages, 1) + assert.NotEmpty(t, after.Messages) + return nil + }, + }) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.True(t, callbackCalled) + }) + + t.Run("custom finalize is used", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + mw := newMW(&Config{ + Model: cm, + Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return []adk.Message{ + schema.SystemMessage("custom system"), + summary, + }, nil + }, + }) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.NoError(t, err) + assert.Len(t, result, 2) + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, "custom system", result[0].Content) + }) + + t.Run("callback error propagates", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + mw := newMW(&Config{ + Model: cm, + Callback: func(ctx context.Context, before, after adk.ChatModelAgentState) error { + return fmt.Errorf("callback error") + }, + }) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "callback error") + }) + + t.Run("finalize error propagates", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + mw := newMW(&Config{ + Model: cm, + Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return nil, fmt.Errorf("finalize error") + }, + }) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{schema.UserMessage("hello")}, + } + + result, err := mw.Summarize(ctx, state) + assert.Error(t, err) + assert.Nil(t, result) + assert.Contains(t, err.Error(), "finalize error") + }) + + t.Run("preserves system messages", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockBaseChatModel(ctrl) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&schema.Message{ + Role: schema.Assistant, + Content: "Summary", + }, nil).Times(1) + + mw := newMW(&Config{Model: cm}) + + state := &adk.ChatModelAgentState{ + Messages: []adk.Message{ + schema.SystemMessage("System 1"), + schema.SystemMessage("System 2"), + schema.UserMessage(strings.Repeat("a", 100)), + }, + } + + result, err := mw.Summarize(ctx, state) + assert.NoError(t, err) + assert.Len(t, result, 3) + assert.Equal(t, schema.System, result[0].Role) + assert.Equal(t, "System 1", result[0].Content) + assert.Equal(t, schema.System, result[1].Role) + assert.Equal(t, "System 2", result[1].Content) + }) +} + +func TestNewTypedAgenticMessage(t *testing.T) { + ctx := context.Background() + + // TypedConfig requires a Model, so passing an empty config will return an error. + // This test verifies that NewTyped[*schema.AgenticMessage] compiles correctly. + mw, err := NewTyped(ctx, &TypedConfig[*schema.AgenticMessage]{}) + assert.Error(t, err) + assert.Nil(t, mw) + + // Verify the return type is correct at compile time. + var _ adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] = mw +} + +// ============================================================================ +// Generic message helpers (prefixed with 's' to avoid conflicts) +// ============================================================================ + +func smakeUserMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(content)).(M) + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(content)).(M) + } + panic("unreachable") +} + +func smakeSystemMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.SystemMessage(content)).(M) + case *schema.AgenticMessage: + return any(schema.SystemAgenticMessage(content)).(M) + } + panic("unreachable") +} + +func smakeAssistantMsg[M adk.MessageType](content string) M { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.AssistantMessage(content, nil)).(M) + case *schema.AgenticMessage: + am := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: content}), + }, + } + return any(am).(M) + } + panic("unreachable") +} + +// ============================================================================ +// Generic mock model +// ============================================================================ + +type genericMockModel[M adk.MessageType] struct { + response M + err error +} + +func (m *genericMockModel[M]) Generate(_ context.Context, _ []M, _ ...model.Option) (M, error) { + return m.response, m.err +} + +func (m *genericMockModel[M]) Stream(_ context.Context, _ []M, _ ...model.Option) (*schema.StreamReader[M], error) { + return nil, fmt.Errorf("not implemented") +} + +// ============================================================================ +// Generic tests +// ============================================================================ + +func TestSummarizationGeneric(t *testing.T) { + t.Run("Message", func(t *testing.T) { + t.Run("Helpers", testSummarizationHelpers[*schema.Message]) + t.Run("Flow", testSummarizationFlow[*schema.Message]) + t.Run("TokenCounterUsesStateToolInfos", testTokenCounterReceivesStateToolInfos[*schema.Message]) + }) + t.Run("AgenticMessage", func(t *testing.T) { + t.Run("Helpers", testSummarizationHelpers[*schema.AgenticMessage]) + t.Run("Flow", testSummarizationFlow[*schema.AgenticMessage]) + t.Run("TokenCounterUsesStateToolInfos", testTokenCounterReceivesStateToolInfos[*schema.AgenticMessage]) + }) +} + +func TestEmitInternalEvents_AgenticMessage_RequiresExecContext(t *testing.T) { + ctx := context.Background() + + longContent := strings.Repeat("x", 800000) + msgs := []*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeSystem, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "system"}), + }, + }, + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.UserInputText{Text: longContent}), + }, + }, + } + + mockResp := smakeAssistantMsg[*schema.AgenticMessage]("This is the summary.") + mw, err := NewTyped(ctx, &TypedConfig[*schema.AgenticMessage]{ + Model: &genericMockModel[*schema.AgenticMessage]{response: mockResp}, + EmitInternalEvents: true, + Trigger: &TriggerCondition{ + ContextTokens: 1, + }, + }) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{Messages: msgs} + _, _, err = mw.BeforeModelRewriteState(ctx, state, nil) + assert.Error(t, err, "should error without exec context when EmitInternalEvents is true") + assert.Contains(t, err.Error(), "send internal event") +} + +func testSummarizationHelpers[M adk.MessageType](t *testing.T) { + t.Run("isSystemRole", func(t *testing.T) { + sys := smakeSystemMsg[M]("hello") + usr := smakeUserMsg[M]("hello") + assert.True(t, isSystemRole(sys)) + assert.False(t, isSystemRole(usr)) + }) + + t.Run("isUserRole", func(t *testing.T) { + usr := smakeUserMsg[M]("hello") + sys := smakeSystemMsg[M]("hello") + assert.True(t, isUserRole(usr)) + assert.False(t, isUserRole(sys)) + }) + + t.Run("getUserMsgTextContent", func(t *testing.T) { + usr := smakeUserMsg[M]("hello world") + assert.Equal(t, "hello world", getUserMsgTextContent(usr)) + }) + + t.Run("getMsgExtra_setMsgExtra", func(t *testing.T) { + msg := smakeUserMsg[M]("test") + extra := getMsgExtra(msg) + assert.Nil(t, extra) + + setMsgExtra(msg, "key1", "value1") + extra = getMsgExtra(msg) + assert.Equal(t, "value1", extra["key1"]) + }) + + t.Run("makeSystemMsg", func(t *testing.T) { + msg := makeSystemMsg[M]("system prompt") + assert.True(t, isSystemRole(msg)) + switch m := any(msg).(type) { + case *schema.Message: + assert.Equal(t, "system prompt", m.Content) + case *schema.AgenticMessage: + require.Len(t, m.ContentBlocks, 1) + assert.Equal(t, "system prompt", m.ContentBlocks[0].UserInputText.Text) + } + }) + + t.Run("makeUserMsg", func(t *testing.T) { + msg := makeUserMsg[M]("user input") + assert.True(t, isUserRole(msg)) + assert.Equal(t, "user input", getUserMsgTextContent(msg)) + }) + + t.Run("newTypedSummaryMessage", func(t *testing.T) { + msg := newTypedSummaryMessage[M]("summary content") + + assert.True(t, isUserRole(msg)) + + switch m := any(msg).(type) { + case *schema.Message: + assert.Equal(t, schema.User, m.Role) + assert.Equal(t, "summary content", m.Content) + case *schema.AgenticMessage: + assert.Equal(t, schema.AgenticRoleTypeUser, m.Role) + require.Len(t, m.ContentBlocks, 1) + assert.Equal(t, "summary content", m.ContentBlocks[0].UserInputText.Text) + } + }) + + t.Run("isInnerMessage", func(t *testing.T) { + summaryMsg := newTypedSummaryMessage[M]("summary content") + assert.True(t, isInternalUserMessage(summaryMsg)) + assert.False(t, isPreservedMessage(summaryMsg)) + + skillsMsg := makeUserMsg[M]("skills content") + setMsgExtra(skillsMsg, extraKeyContentType, string(contentTypeSkills)) + assert.True(t, isInternalUserMessage(skillsMsg)) + assert.True(t, isPreservedMessage(skillsMsg)) + + normalMsg := makeUserMsg[M]("normal content") + assert.False(t, isInternalUserMessage(normalMsg)) + assert.False(t, isPreservedMessage(normalMsg)) + }) +} + +func testSummarizationFlow[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + summaryText := "This is a summary of the conversation." + mockModel := &genericMockModel[M]{ + response: smakeAssistantMsg[M](summaryText), + } + + tokenCounter := func(_ context.Context, input *TypedTokenCounterInput[M]) (int, error) { + total := 0 + for _, msg := range input.Messages { + total += len(getUserMsgTextContent(msg)) + } + return total, nil + } + + cfg := &TypedConfig[M]{ + Model: mockModel, + TokenCounter: tokenCounter, + Trigger: &TriggerCondition{ + ContextTokens: 20, + }, + } + + mw, err := NewTyped(ctx, cfg) + require.NoError(t, err) + + msgs := []M{ + smakeSystemMsg[M]("You are a helpful assistant."), + smakeUserMsg[M]("Tell me a very long story about dragons and castles"), + smakeAssistantMsg[M]("Once upon a time there was a magnificent dragon"), + smakeUserMsg[M]("What happened next?"), + } + + state := &adk.TypedChatModelAgentState[M]{Messages: msgs} + mtx := &adk.TypedModelContext[M]{} + + _, newState, err := mw.BeforeModelRewriteState(ctx, state, mtx) + require.NoError(t, err) + + require.GreaterOrEqual(t, len(newState.Messages), 2, + "should have at least system + summary messages") + + assert.True(t, isSystemRole(newState.Messages[0]), + "first message should be system") + + foundSummary := false + for _, msg := range newState.Messages { + extra := getMsgExtra(msg) + if extra != nil { + if ct, ok := extra[extraKeyContentType]; ok && ct == string(contentTypeSummary) { + foundSummary = true + break + } + } + if strings.Contains(getUserMsgTextContent(msg), summaryText) { + foundSummary = true + break + } + } + assert.True(t, foundSummary, "should have a summary message") +} + +func testTokenCounterReceivesStateToolInfos[M adk.MessageType](t *testing.T) { + ctx := context.Background() + + stateTools := []*schema.ToolInfo{ + {Name: "state_tool_a"}, + {Name: "state_tool_b"}, + } + mcTools := []*schema.ToolInfo{ + {Name: "mc_tool_should_not_appear"}, + } + + var receivedTools []*schema.ToolInfo + tokenCounter := func(_ context.Context, input *TypedTokenCounterInput[M]) (int, error) { + receivedTools = input.Tools + return 0, nil + } + + cfg := &TypedConfig[M]{ + Model: &genericMockModel[M]{ + response: smakeAssistantMsg[M]("unused"), + }, + TokenCounter: tokenCounter, + Trigger: &TriggerCondition{ + ContextTokens: 9999, + }, + } + + mw, err := NewTyped(ctx, cfg) + require.NoError(t, err) + + state := &adk.TypedChatModelAgentState[M]{ + Messages: []M{smakeUserMsg[M]("hello")}, + ToolInfos: stateTools, + } + mc := &adk.TypedModelContext[M]{Tools: mcTools} + + _, _, err = mw.BeforeModelRewriteState(ctx, state, mc) + require.NoError(t, err) + + require.NotNil(t, receivedTools, "token counter should have been called") + require.Len(t, receivedTools, 2) + assert.Equal(t, "state_tool_a", receivedTools[0].Name) + assert.Equal(t, "state_tool_b", receivedTools[1].Name) +} + +func TestGetAssistantTextContent(t *testing.T) { + t.Run("schema.Message with MultiContent", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + Content: "fallback content", + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "hello"}, + {Type: schema.ChatMessagePartTypeText, Text: "world"}, + }, + } + got := getAssistantTextContent(msg) + assert.Equal(t, "hello\nworld", got) + }) + + t.Run("schema.Message with MultiContent skips non-text parts", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + Content: "fallback", + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "text part"}, + {Type: schema.ChatMessagePartTypeImageURL}, + {Type: schema.ChatMessagePartTypeText, Text: ""}, + }, + } + got := getAssistantTextContent(msg) + assert.Equal(t, "text part", got) + }) + + t.Run("schema.Message falls back to Content when MultiContent is empty", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + Content: "plain content", + } + got := getAssistantTextContent(msg) + assert.Equal(t, "plain content", got) + }) + + t.Run("schema.Message falls back to Content when MultiContent has no text", func(t *testing.T) { + msg := &schema.Message{ + Role: schema.Assistant, + Content: "fallback", + AssistantGenMultiContent: []schema.MessageOutputPart{ + {Type: schema.ChatMessagePartTypeImageURL}, + }, + } + got := getAssistantTextContent(msg) + assert.Equal(t, "fallback", got) + }) + + t.Run("schema.AgenticMessage with multiple text blocks", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: "first"}), + schema.NewContentBlock(&schema.AssistantGenText{Text: "second"}), + }, + } + got := getAssistantTextContent(msg) + assert.Equal(t, "first\nsecond", got) + }) + + t.Run("schema.AgenticMessage with nil blocks", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + nil, + schema.NewContentBlock(&schema.AssistantGenText{Text: "only"}), + nil, + }, + } + got := getAssistantTextContent(msg) + assert.Equal(t, "only", got) + }) + + t.Run("schema.AgenticMessage with no text blocks", func(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{}, + } + got := getAssistantTextContent(msg) + assert.Equal(t, "", got) + }) +} diff --git a/adk/prebuilt/deep/checkpoint_compat_resume_test.go b/adk/prebuilt/deep/checkpoint_compat_resume_test.go new file mode 100644 index 0000000..1a4f8ba --- /dev/null +++ b/adk/prebuilt/deep/checkpoint_compat_resume_test.go @@ -0,0 +1,203 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package deep + +import ( + "context" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type compatCheckpointStore struct { + data map[string][]byte +} + +func newCompatCheckpointStore() *compatCheckpointStore { + return &compatCheckpointStore{data: make(map[string][]byte)} +} + +func (s *compatCheckpointStore) Set(_ context.Context, key string, value []byte) error { + s.data[key] = append([]byte(nil), value...) + return nil +} + +func (s *compatCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.data[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), v...), true, nil +} + +type interruptingSubAgentTool struct { + name string +} + +func (t *interruptingSubAgentTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "interrupts on first call and resumes from stored state", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "action": {Type: schema.String}, + }), + }, nil +} + +func (t *interruptingSubAgentTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + _, _, _ = tool.GetInterruptState[string](ctx) + return "resumed", nil +} + +func readTestdataBytes(t *testing.T, filename string) []byte { + t.Helper() + _, file, _, ok := runtime.Caller(0) + assert.True(t, ok) + p := filepath.Join(filepath.Dir(file), "testdata", filename) + b, err := os.ReadFile(p) + assert.NoError(t, err) + assert.NotEmpty(t, b) + return b +} + +func runDeepAgentCheckpointCompat(t *testing.T, checkpointID string, filename string) { + t.Helper() + ctx := context.Background() + + data := readTestdataBytes(t, filename) + + store := newCompatCheckpointStore() + assert.NoError(t, store.Set(ctx, checkpointID, data)) + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + interruptToolName := "interrupt_in_subagent_tool" + subTool := &interruptingSubAgentTool{name: interruptToolName} + + deepModel := mockModel.NewMockBaseChatModel(ctrl) + subModel := mockModel.NewMockBaseChatModel(ctrl) + + deepModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("deep done", nil), nil + }).AnyTimes() + + subModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("sub done", nil), nil + }).AnyTimes() + + subAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "sub_chatmodel_agent", + Description: "sub agent", + Model: subModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{subTool}, + }, + }, + MaxIterations: 4, + }) + assert.NoError(t, err) + + deepAgent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: deepModel, + SubAgents: []adk.Agent{subAgent}, + MaxIteration: 4, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: deepAgent, + CheckPointStore: store, + }) + + it, err := runner.Resume(ctx, checkpointID) + assert.NoError(t, err) + + var sawDeepDone bool + var sawAnyOutput bool + for { + ev, ok := it.Next() + if !ok { + break + } + assert.NoError(t, ev.Err) + if ev.Output != nil && ev.Output.MessageOutput != nil && ev.Output.MessageOutput.Message != nil { + sawAnyOutput = true + msg := ev.Output.MessageOutput.Message + if msg.Role == schema.Assistant && strings.Contains(msg.Content, "deep done") { + sawDeepDone = true + } + } + } + + assert.True(t, sawAnyOutput) + assert.True(t, sawDeepDone) +} + +func TestDeepAgentCheckpointCompat_V0_8_Resume(t *testing.T) { + tests := []struct { + name string + checkpointID string + filename string + }{ + { + name: "v0.7.37", + checkpointID: "checkpoint_compat_v0_7_37", + filename: "checkpoint_data_v0.7.37.bin", + }, + { + name: "v0.8.2", + checkpointID: "checkpoint_compat_v0_8_2", + filename: "checkpoint_data_v0.8.2.bin", + }, + { + name: "v0.8.3", + checkpointID: "checkpoint_compat_v0_8_3", + filename: "checkpoint_data_v0.8.3.bin", + }, + { + name: "v0.8.4", + checkpointID: "checkpoint_compat_v0_8_4", + filename: "checkpoint_data_v0.8.4.bin", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + runDeepAgentCheckpointCompat(t, tc.checkpointID, tc.filename) + }) + } +} diff --git a/adk/prebuilt/deep/deep.go b/adk/prebuilt/deep/deep.go new file mode 100644 index 0000000..5315113 --- /dev/null +++ b/adk/prebuilt/deep/deep.go @@ -0,0 +1,267 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package deep provides a prebuilt agent with deep task orchestration. +package deep + +import ( + "context" + "fmt" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/adk/internal" + filesystem2 "github.com/cloudwego/eino/adk/middlewares/filesystem" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool/utils" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[TODO]("_eino_adk_prebuilt_deep_todo") + schema.RegisterName[[]TODO]("_eino_adk_prebuilt_deep_todo_slice") +} + +// TypedConfig defines the configuration for creating a DeepAgent parameterized by message type. +// An Agentic DeepAgent (M = *schema.AgenticMessage) only supports Agentic sub-agents, +// and a standard DeepAgent (M = *schema.Message) only supports standard sub-agents. +// This is enforced by the type system through the SubAgents field. +type TypedConfig[M adk.MessageType] struct { + // Name is the identifier for the Deep agent. + Name string + // Description provides a brief explanation of the agent's purpose. + Description string + + // ChatModel is the model used by DeepAgent for reasoning and task execution. + // If the agent uses any tools, this model must support the model.WithTools call option, + // as that's how the agent configures the model with tool information. + ChatModel model.BaseModel[M] + // Instruction contains the system prompt that guides the agent's behavior. + // When empty, a built-in default system prompt will be used, which includes general assistant + // behavior guidelines, security policies, coding style guidelines, and tool usage policies. + Instruction string + // SubAgents are specialized agents that can be invoked by the agent. + // For M = *schema.AgenticMessage, only agentic sub-agents are accepted. + SubAgents []adk.TypedAgent[M] + // ToolsConfig provides the tools and tool-calling configurations available for the agent to invoke. + ToolsConfig adk.ToolsConfig + // MaxIteration limits the maximum number of reasoning iterations the agent can perform. + MaxIteration int + + // Backend provides filesystem operations used by tools and offloading. + // If set, filesystem tools (read_file, write_file, edit_file, glob, grep) will be registered. + // Optional. + Backend filesystem.Backend + // Shell provides shell command execution capability. + // If set, an execute tool will be registered to support shell command execution. + // Optional. Mutually exclusive with StreamingShell. + Shell filesystem.Shell + // StreamingShell provides streaming shell command execution capability. + // If set, a streaming execute tool will be registered to support streaming shell command execution. + // Optional. Mutually exclusive with Shell. + StreamingShell filesystem.StreamingShell + + // WithoutWriteTodos disables the built-in write_todos tool when set to true. + WithoutWriteTodos bool + // WithoutGeneralSubAgent disables the general-purpose subagent when set to true. + WithoutGeneralSubAgent bool + // TaskToolDescriptionGenerator allows customizing the description for the task tool. + // If provided, this function generates the tool description based on available subagents. + TaskToolDescriptionGenerator func(ctx context.Context, availableAgents []adk.TypedAgent[M]) (string, error) + + Middlewares []adk.AgentMiddleware + + // Handlers configures interface-based handlers for extending agent behavior. + // Unlike Middlewares (struct-based), Handlers allow users to: + // - Add custom methods to their handler implementations + // - Return modified context from handler methods + // - Centralize configuration in struct fields instead of closures + // + // Handlers are processed after Middlewares, in registration order. + // See adk.ChatModelAgentMiddleware documentation for when to use Handlers vs Middlewares. + Handlers []adk.TypedChatModelAgentMiddleware[M] + + ModelRetryConfig *adk.TypedModelRetryConfig[M] + // ModelFailoverConfig configures failover behavior for the ChatModel. + // When set, the agent will automatically fail over to alternative models on errors. + // This config is also propagated to the general sub-agent. + ModelFailoverConfig *adk.ModelFailoverConfig[M] + + // OutputKey stores the agent's response in the session. + // Optional. When set, stores output via AddSessionValue(ctx, outputKey, msg.Content). + OutputKey string +} + +// Config defines the configuration for creating a standard DeepAgent. +type Config = TypedConfig[*schema.Message] + +// NewTyped creates a new typed Deep agent instance with the provided configuration. +// This function initializes built-in tools, creates a task tool for subagent orchestration, +// and returns a fully configured TypedChatModelAgent ready for execution. +func NewTyped[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M]) (adk.TypedResumableAgent[M], error) { + handlers, err := buildTypedBuiltinAgentMiddlewares(ctx, cfg) + if err != nil { + return nil, err + } + + instruction := cfg.Instruction + if len(instruction) == 0 { + instruction = internal.SelectPrompt(internal.I18nPrompts{ + English: baseAgentInstruction, + Chinese: baseAgentInstructionChinese, + }) + } + + if !cfg.WithoutGeneralSubAgent || len(cfg.SubAgents) > 0 { + tt, err := typedTaskToolMiddleware( + ctx, + cfg.TaskToolDescriptionGenerator, + cfg.SubAgents, + + cfg.WithoutGeneralSubAgent, + cfg.ChatModel, + instruction, + cfg.ToolsConfig, + cfg.MaxIteration, + cfg.Middlewares, + append(handlers, cfg.Handlers...), + cfg.ModelFailoverConfig, + ) + if err != nil { + return nil, fmt.Errorf("failed to new task tool: %w", err) + } + handlers = append(handlers, tt) + } + + return adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[M]{ + Name: cfg.Name, + Description: cfg.Description, + Instruction: instruction, + Model: cfg.ChatModel, + ToolsConfig: cfg.ToolsConfig, + MaxIterations: cfg.MaxIteration, + Middlewares: cfg.Middlewares, + Handlers: append(handlers, cfg.Handlers...), + + GenModelInput: typedGenModelInput[M], + ModelRetryConfig: cfg.ModelRetryConfig, + ModelFailoverConfig: cfg.ModelFailoverConfig, + OutputKey: cfg.OutputKey, + }) +} + +// New creates a new Deep agent instance with the provided configuration. +// This function initializes built-in tools, creates a task tool for subagent orchestration, +// and returns a fully configured ChatModelAgent ready for execution. +func New(ctx context.Context, cfg *Config) (adk.ResumableAgent, error) { + return NewTyped(ctx, cfg) +} + +func typedGenModelInput[M adk.MessageType](_ context.Context, instruction string, input *adk.TypedAgentInput[M]) ([]M, error) { + var zero M + switch any(zero).(type) { + case *schema.Message: + msgs := make([]*schema.Message, 0, len(input.Messages)+1) + if instruction != "" { + msgs = append(msgs, schema.SystemMessage(instruction)) + } + // Type assertion is safe here because M = *schema.Message. + for _, m := range input.Messages { + msgs = append(msgs, any(m).(*schema.Message)) + } + result := make([]M, len(msgs)) + for i, m := range msgs { + result[i] = any(m).(M) + } + return result, nil + case *schema.AgenticMessage: + msgs := make([]*schema.AgenticMessage, 0, len(input.Messages)+1) + if instruction != "" { + msgs = append(msgs, schema.SystemAgenticMessage(instruction)) + } + for _, m := range input.Messages { + msgs = append(msgs, any(m).(*schema.AgenticMessage)) + } + result := make([]M, len(msgs)) + for i, m := range msgs { + result[i] = any(m).(M) + } + return result, nil + } + panic("unreachable") +} + +func buildTypedBuiltinAgentMiddlewares[M adk.MessageType](ctx context.Context, cfg *TypedConfig[M]) ([]adk.TypedChatModelAgentMiddleware[M], error) { + var ms []adk.TypedChatModelAgentMiddleware[M] + if !cfg.WithoutWriteTodos { + t, err := typedNewWriteTodos[M]() + if err != nil { + return nil, err + } + ms = append(ms, t) + } + + if cfg.Backend != nil || cfg.Shell != nil || cfg.StreamingShell != nil { + fm, err := filesystem2.NewTyped[M](ctx, &filesystem2.MiddlewareConfig{ + Backend: cfg.Backend, + Shell: cfg.Shell, + StreamingShell: cfg.StreamingShell, + }) + if err != nil { + return nil, err + } + ms = append(ms, fm) + } + + return ms, nil +} + +type TODO struct { + Content string `json:"content"` + ActiveForm string `json:"activeForm"` + Status string `json:"status" jsonschema:"enum=pending,enum=in_progress,enum=completed"` +} + +type writeTodosArguments struct { + Todos []TODO `json:"todos"` +} + +func typedNewWriteTodos[M adk.MessageType]() (adk.TypedChatModelAgentMiddleware[M], error) { + toolDesc := internal.SelectPrompt(internal.I18nPrompts{ + English: writeTodosToolDescription, + Chinese: writeTodosToolDescriptionChinese, + }) + resultMsg := internal.SelectPrompt(internal.I18nPrompts{ + English: "Updated todo list to %s", + Chinese: "已更新待办列表为 %s", + }) + + t, err := utils.InferTool("write_todos", toolDesc, func(ctx context.Context, input writeTodosArguments) (output string, err error) { + adk.AddSessionValue(ctx, SessionKeyTodos, input.Todos) + todos, err := sonic.MarshalString(input.Todos) + if err != nil { + return "", err + } + return fmt.Sprintf(resultMsg, todos), nil + }) + if err != nil { + return nil, err + } + + return typedBuildAppendPromptTool[M]("", t), nil +} diff --git a/adk/prebuilt/deep/deep_test.go b/adk/prebuilt/deep/deep_test.go new file mode 100644 index 0000000..b39cfe9 --- /dev/null +++ b/adk/prebuilt/deep/deep_test.go @@ -0,0 +1,700 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package deep + +import ( + "context" + "fmt" + "io" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/prebuilt/planexecute" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type sequentialAgenticModel struct { + responses []*schema.AgenticMessage + callCount int32 +} + +func (m *sequentialAgenticModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&m.callCount, 1) - 1 + if int(idx) >= len(m.responses) { + return nil, fmt.Errorf("sequentialAgenticModel: no more responses (call #%d)", idx) + } + return m.responses[idx], nil +} + +func (m *sequentialAgenticModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + result, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.AgenticMessage{result}), nil +} + +func agenticText(text string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.AssistantGenText{Text: text}), + }, + } +} + +func agenticToolCall(toolName, callID, args string) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeAssistant, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolCall{Name: toolName, CallID: callID, Arguments: args}), + }, + } +} + +func readAgenticText(msg *schema.AgenticMessage) string { + if msg == nil { + return "" + } + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + if block.AssistantGenText != nil { + return block.AssistantGenText.Text + } + } + return "" +} + +type mockSearchTool struct{} + +func (m *mockSearchTool) Info(context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "search", + Desc: "search latest news", + }, nil +} + +func (m *mockSearchTool) InvokableRun(context.Context, string, ...tool.Option) (string, error) { + return "latest news search result", nil +} + +func TestGenModelInput(t *testing.T) { + ctx := context.Background() + + t.Run("WithInstruction", func(t *testing.T) { + input := &adk.AgentInput{ + Messages: []*schema.Message{ + schema.UserMessage("hello"), + }, + } + + msgs, err := typedGenModelInput(ctx, "You are a helpful assistant", input) + assert.NoError(t, err) + assert.Len(t, msgs, 2) + assert.Equal(t, schema.System, msgs[0].Role) + assert.Equal(t, "You are a helpful assistant", msgs[0].Content) + assert.Equal(t, schema.User, msgs[1].Role) + assert.Equal(t, "hello", msgs[1].Content) + }) + + t.Run("WithoutInstruction", func(t *testing.T) { + input := &adk.AgentInput{ + Messages: []*schema.Message{ + schema.UserMessage("hello"), + }, + } + + msgs, err := typedGenModelInput(ctx, "", input) + assert.NoError(t, err) + assert.Len(t, msgs, 1) + assert.Equal(t, schema.User, msgs[0].Role) + assert.Equal(t, "hello", msgs[0].Content) + }) +} + +func TestWriteTodos(t *testing.T) { + m, err := buildTypedBuiltinAgentMiddlewares(context.Background(), &Config{WithoutWriteTodos: false}) + assert.NoError(t, err) + + wt := m[0].(*typedAppendPromptTool[*schema.Message]).t.(tool.InvokableTool) + + todos := `[{"content":"content1","activeForm":"","status":"pending"},{"content":"content2","activeForm":"","status":"pending"}]` + args := fmt.Sprintf(`{"todos": %s}`, todos) + + result, err := wt.InvokableRun(context.Background(), args) + assert.NoError(t, err) + assert.Equal(t, fmt.Sprintf("Updated todo list to %s", todos), result) +} + +func TestDeepSubAgentSharesSessionValues(t *testing.T) { + ctx := context.Background() + spy := &spySubAgent{} + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + calls := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + calls++ + if calls == 1 { + c := schema.ToolCall{ID: "id-1", Type: "function"} + c.Function.Name = taskToolName + c.Function.Arguments = fmt.Sprintf(`{"subagent_type":"%s","description":"from_parent"}`, spy.Name(ctx)) + return schema.AssistantMessage("", []schema.ToolCall{c}), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + Instruction: "you are deep agent", + SubAgents: []adk.Agent{spy}, + ToolsConfig: adk.ToolsConfig{}, + MaxIteration: 2, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + + r := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent}) + it := r.Run(ctx, []adk.Message{schema.UserMessage("hi")}, adk.WithSessionValues(map[string]any{"parent_key": "parent_val"})) + for { + if _, ok := it.Next(); !ok { + break + } + } + + assert.Equal(t, "parent_val", spy.seenParentValue) +} + +func TestDeepSubAgentFollowsStreamingMode(t *testing.T) { + ctx := context.Background() + spy := &spyStreamingSubAgent{} + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + subName := spy.Name(ctx) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "id-1", + Type: "function", + Function: schema.FunctionCall{ + Name: taskToolName, + Arguments: fmt.Sprintf(`{"subagent_type":"%s","description":"from_parent"}`, subName), + }, + }, + }), + }), nil). + Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("done", nil), + }), nil). + Times(1) + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + Instruction: "you are deep agent", + SubAgents: []adk.Agent{spy}, + ToolsConfig: adk.ToolsConfig{}, + MaxIteration: 2, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + + r := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []adk.Message{schema.UserMessage("hi")}) + for { + if _, ok := it.Next(); !ok { + break + } + } + + assert.True(t, spy.seenEnableStreaming) +} + +type spySubAgent struct { + seenParentValue any +} + +func (s *spySubAgent) Name(context.Context) string { return "spy-subagent" } +func (s *spySubAgent) Description(context.Context) string { return "spy" } +func (s *spySubAgent) Run(ctx context.Context, _ *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + s.seenParentValue, _ = adk.GetSessionValue(ctx, "parent_key") + it, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + gen.Send(adk.EventFromMessage(schema.AssistantMessage("ok", nil), nil, schema.Assistant, "")) + gen.Close() + return it +} + +type spyStreamingSubAgent struct { + seenEnableStreaming bool +} + +func (s *spyStreamingSubAgent) Name(context.Context) string { return "spy-streaming-subagent" } +func (s *spyStreamingSubAgent) Description(context.Context) string { return "spy" } +func (s *spyStreamingSubAgent) Run(_ context.Context, input *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + if input != nil { + s.seenEnableStreaming = input.EnableStreaming + } + it, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + gen.Send(adk.EventFromMessage(schema.AssistantMessage("ok", nil), nil, schema.Assistant, "")) + gen.Close() + return it +} + +func TestDeepAgentWithPlanExecuteSubAgent_InternalEventsEmitted(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + deepModel := mockModel.NewMockToolCallingChatModel(ctrl) + plannerModel := mockModel.NewMockToolCallingChatModel(ctrl) + executorModel := mockModel.NewMockToolCallingChatModel(ctrl) + replannerModel := mockModel.NewMockToolCallingChatModel(ctrl) + + deepModel.EXPECT().WithTools(gomock.Any()).Return(deepModel, nil).AnyTimes() + plannerModel.EXPECT().WithTools(gomock.Any()).Return(plannerModel, nil).AnyTimes() + executorModel.EXPECT().WithTools(gomock.Any()).Return(executorModel, nil).AnyTimes() + replannerModel.EXPECT().WithTools(gomock.Any()).Return(replannerModel, nil).AnyTimes() + + plannerModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + go func() { + defer sw.Close() + planJSON := `{"steps":["step1"]}` + msg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "plan_call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "plan", + Arguments: planJSON, + }, + }, + }) + sw.Send(msg, nil) + }() + return sr, nil + }, + ).Times(1) + + executorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("executed step1", nil), nil + }, + ).Times(1) + + replannerModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + go func() { + defer sw.Close() + responseJSON := `{"response":"final response"}` + msg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "respond_call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "respond", + Arguments: responseJSON, + }, + }, + }) + sw.Send(msg, nil) + }() + return sr, nil + }, + ).Times(1) + + planner, err := planexecute.NewPlanner(ctx, &planexecute.PlannerConfig{ + ToolCallingChatModel: plannerModel, + }) + assert.NoError(t, err) + + executor, err := planexecute.NewExecutor(ctx, &planexecute.ExecutorConfig{ + Model: executorModel, + }) + assert.NoError(t, err) + + replanner, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{ + ChatModel: replannerModel, + }) + assert.NoError(t, err) + + planExecuteAgent, err := planexecute.New(ctx, &planexecute.Config{ + Planner: planner, + Executor: executor, + Replanner: replanner, + }) + assert.NoError(t, err) + + namedPlanExecuteAgent := &namedPlanExecuteAgent{ + ResumableAgent: planExecuteAgent, + name: "plan_execute_subagent", + description: "a plan execute subagent", + } + + deepModelCalls := 0 + deepModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + deepModelCalls++ + if deepModelCalls == 1 { + c := schema.ToolCall{ID: "id-1", Type: "function"} + c.Function.Name = taskToolName + c.Function.Arguments = fmt.Sprintf(`{"subagent_type":"%s","description":"execute the plan"}`, namedPlanExecuteAgent.name) + return schema.AssistantMessage("", []schema.ToolCall{c}), nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + + deepAgent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: deepModel, + Instruction: "you are deep agent", + SubAgents: []adk.Agent{namedPlanExecuteAgent}, + ToolsConfig: adk.ToolsConfig{EmitInternalEvents: true}, + MaxIteration: 5, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + + r := adk.NewRunner(ctx, adk.RunnerConfig{Agent: deepAgent}) + it := r.Run(ctx, []adk.Message{schema.UserMessage("hi")}) + + var events []*adk.AgentEvent + for { + event, ok := it.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Greater(t, len(events), 0, "should have at least one event") + + var deepAgentEvents []*adk.AgentEvent + var plannerEvents []*adk.AgentEvent + var executorEvents []*adk.AgentEvent + var replannerEvents []*adk.AgentEvent + var planExecuteEvents []*adk.AgentEvent + + for _, event := range events { + switch event.AgentName { + case "deep": + deepAgentEvents = append(deepAgentEvents, event) + case "planner": + plannerEvents = append(plannerEvents, event) + case "executor": + executorEvents = append(executorEvents, event) + case "replanner": + replannerEvents = append(replannerEvents, event) + case "plan_execute_replan", "execute_replan": + planExecuteEvents = append(planExecuteEvents, event) + } + } + + assert.Greater(t, len(deepAgentEvents), 0, "should have events from deep agent") + + assert.Greater(t, len(plannerEvents), 0, "planner internal events should be emitted when EmitInternalEvents is true") + assert.Greater(t, len(executorEvents), 0, "executor internal events should be emitted when EmitInternalEvents is true") + assert.Greater(t, len(replannerEvents), 0, "replanner internal events should be emitted when EmitInternalEvents is true") + + t.Logf("Total events: %d", len(events)) + t.Logf("Deep agent events: %d", len(deepAgentEvents)) + t.Logf("Planner events: %d", len(plannerEvents)) + t.Logf("Executor events: %d", len(executorEvents)) + t.Logf("Replanner events: %d", len(replannerEvents)) + t.Logf("PlanExecute events: %d", len(planExecuteEvents)) +} + +type namedPlanExecuteAgent struct { + adk.ResumableAgent + name string + description string +} + +func (n *namedPlanExecuteAgent) Name(_ context.Context) string { + return n.name +} + +func (n *namedPlanExecuteAgent) Description(_ context.Context) string { + return n.description +} + +func TestAgenticDeepAgentEmitInternalEventsFromSubAgent(t *testing.T) { + ctx := context.Background() + + subAgentModel := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCall("search", "search-call-1", `{"query":"latest news"}`), + agenticText("search result from subagent"), + }, + } + webSearchAgent, err := adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "web_searcher", + Description: "agent dedicated to web search", + Model: subAgentModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&mockSearchTool{}}, + }, + }, + }) + assert.NoError(t, err) + + deepModel := &sequentialAgenticModel{ + responses: []*schema.AgenticMessage{ + agenticToolCall(taskToolName, "task-call-1", `{"subagent_type":"web_searcher","description":"search latest news"}`), + agenticText("deep final answer"), + }, + } + deepAgent, err := NewTyped(ctx, &TypedConfig[*schema.AgenticMessage]{ + Name: "deep", + Description: "deep agent", + ChatModel: deepModel, + SubAgents: []adk.TypedAgent[*schema.AgenticMessage]{webSearchAgent}, + ToolsConfig: adk.ToolsConfig{EmitInternalEvents: true}, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + + runner := adk.NewTypedRunner(adk.TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: deepAgent, + EnableStreaming: true, + }) + iter := runner.Run(ctx, []*schema.AgenticMessage{ + schema.UserAgenticMessage("search latest news"), + }) + + var emittedTexts []string + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output == nil || event.Output.MessageOutput == nil { + continue + } + if event.Output.MessageOutput.IsStreaming { + stream := event.Output.MessageOutput.MessageStream + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + if text := readAgenticText(chunk); text != "" { + emittedTexts = append(emittedTexts, fmt.Sprintf("%s:%s", event.AgentName, text)) + } + } + continue + } + if text := readAgenticText(event.Output.MessageOutput.Message); text != "" { + emittedTexts = append(emittedTexts, fmt.Sprintf("%s:%s", event.AgentName, text)) + } + } + + t.Logf("emitted texts: %v", emittedTexts) + assert.Contains(t, emittedTexts, "deep:deep final answer", "expected DeepAgent final result to be emitted") + assert.Contains(t, emittedTexts, "web_searcher:search result from subagent", "expected DeepAgent to emit Agentic sub-agent events when EmitInternalEvents is true") +} + +func TestDeepAgentOutputKey(t *testing.T) { + t.Run("OutputKeyStoresInSession", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello from DeepAgent", nil), nil). + Times(1) + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + Instruction: "you are deep agent", + MaxIteration: 2, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + OutputKey: "deep_output", + }) + assert.NoError(t, err) + + var capturedSessionValues map[string]any + wrappedAgent := &sessionCaptureAgent{ + Agent: agent, + captureSession: func(values map[string]any) { capturedSessionValues = values }, + } + + r := adk.NewRunner(ctx, adk.RunnerConfig{Agent: wrappedAgent}) + it := r.Run(ctx, []adk.Message{schema.UserMessage("hi")}) + for { + if _, ok := it.Next(); !ok { + break + } + } + + assert.Contains(t, capturedSessionValues, "deep_output") + assert.Equal(t, "Hello from DeepAgent", capturedSessionValues["deep_output"]) + }) + + t.Run("OutputKeyWithStreamingStoresInSession", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("Hello", nil), + schema.AssistantMessage(" from", nil), + schema.AssistantMessage(" DeepAgent", nil), + }), nil). + Times(1) + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + Instruction: "you are deep agent", + MaxIteration: 2, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + OutputKey: "deep_output", + }) + assert.NoError(t, err) + + var capturedSessionValues map[string]any + wrappedAgent := &sessionCaptureAgent{ + Agent: agent, + captureSession: func(values map[string]any) { capturedSessionValues = values }, + } + + r := adk.NewRunner(ctx, adk.RunnerConfig{Agent: wrappedAgent, EnableStreaming: true}) + it := r.Run(ctx, []adk.Message{schema.UserMessage("hi")}) + for { + if _, ok := it.Next(); !ok { + break + } + } + + assert.Contains(t, capturedSessionValues, "deep_output") + assert.Equal(t, "Hello from DeepAgent", capturedSessionValues["deep_output"]) + }) + + t.Run("OutputKeyNotSetWhenEmpty", func(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + cm := mockModel.NewMockToolCallingChatModel(ctrl) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("Hello from DeepAgent", nil), nil). + Times(1) + + agent, err := New(ctx, &Config{ + Name: "deep", + Description: "deep agent", + ChatModel: cm, + Instruction: "you are deep agent", + MaxIteration: 2, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + assert.NoError(t, err) + + var capturedSessionValues map[string]any + wrappedAgent := &sessionCaptureAgent{ + Agent: agent, + captureSession: func(values map[string]any) { capturedSessionValues = values }, + } + + r := adk.NewRunner(ctx, adk.RunnerConfig{Agent: wrappedAgent}) + it := r.Run(ctx, []adk.Message{schema.UserMessage("hi")}) + for { + if _, ok := it.Next(); !ok { + break + } + } + + assert.NotContains(t, capturedSessionValues, "deep_output") + }) +} + +type sessionCaptureAgent struct { + adk.Agent + captureSession func(map[string]any) +} + +func (s *sessionCaptureAgent) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + innerIt := s.Agent.Run(ctx, input, opts...) + it, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + go func() { + defer gen.Close() + for { + event, ok := innerIt.Next() + if !ok { + break + } + gen.Send(event) + } + s.captureSession(adk.GetSessionValues(ctx)) + }() + return it +} diff --git a/adk/prebuilt/deep/prompt.go b/adk/prebuilt/deep/prompt.go new file mode 100644 index 0000000..aba5fd3 --- /dev/null +++ b/adk/prebuilt/deep/prompt.go @@ -0,0 +1,685 @@ +/* + * Copyright (c) 2025 Harrison Chase + * Copyright (c) 2025 CloudWeGo Authors + * SPDX-License-Identifier: MIT + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package deep + +// This file contains prompt templates and tool descriptions adapted from the DeepAgents project and ClaudeCode. +// Original source: https://github.com/langchain-ai/deepagents https://claude.com/product/claude-code +// +// These prompts are used under the terms of the original project's open source license. +// When using this code in your own open source project, ensure compliance with the original license requirements. + +const ( + taskPrompt = ` +# 'task' (subagent spawner) + +You have access to a 'task' tool to launch short-lived subagents that handle isolated tasks. These agents are ephemeral — they live only for the duration of the task and return a single result. +You should proactively use the 'task' tool with specialized agents when the task at hand matches the agent's description. + +When to use the task tool: +- When a task is complex and multi-step, and can be fully delegated in isolation +- When a task is independent of other tasks and can run in parallel +- When a task requires focused reasoning or heavy token/context usage that would bloat the orchestrator thread +- When sandboxing improves reliability (e.g. code execution, structured searches, data formatting) +- When you only care about the output of the subagent, and not the intermediate steps (ex. performing a lot of research and then returned a synthesized report, performing a series of computations or lookups to achieve a concise, relevant answer.) + +Subagent lifecycle: +1. **Spawn** → Provide clear role, instructions, and expected output +2. **Run** → The subagent completes the task autonomously +3. **Return** → The subagent provides a single structured result +4. **Reconcile** → Incorporate or synthesize the result into the main thread + +When NOT to use the task tool: +- If you need to see the intermediate reasoning or steps after the subagent has completed (the task tool hides them) +- If the task is trivial (a few tool calls or simple lookup) +- If delegating does not reduce token usage, complexity, or context switching +- If splitting would add latency without benefit + +## Important Task Tool Usage Notes to Remember +- Whenever possible, parallelize the work that you do. This is true for both tool_calls, and for tasks. Whenever you have independent steps to complete - make tool_calls, or kick off tasks (subagents) in parallel to accomplish them faster. This saves time for the user, which is incredibly important. +- Remember to use the 'task' tool to silo independent tasks within a multi-part objective. +- You should use the 'task' tool whenever you have a complex task that will take multiple steps, and is independent from other tasks that the agent needs to complete. These agents are highly competent and efficient. +` + + baseAgentInstruction = ` +You are a helpful assistant. Use the instructions below and the tools available to you to assist the user. + +IMPORTANT: Assist with authorized security testing, defensive security, CTF challenges, and educational contexts. Refuse requests for destructive techniques, DoS attacks, mass targeting, supply chain compromise, or detection evasion for malicious purposes. Dual-use security tools (C2 frameworks, credential testing, exploit development) require clear authorization context: pentesting engagements, CTF competitions, security research, or defensive use cases. +IMPORTANT: You must NEVER generate or guess URLs for the user unless you are confident that the URLs are for helping the user with programming. You may use URLs provided by the user in their messages or local files. + +# Tone and style +- Only use emojis if the user explicitly requests it. Avoid using emojis in all communication unless asked. +- Your output will be displayed on a command line interface. Your responses should be short and concise. You can use Github-flavored markdown for formatting, and will be rendered in a monospace font using the CommonMark specification. +- Output text to communicate with the user; all text you output outside of tool use is displayed to the user. Only use tools to complete tasks. Never use tools like Bash or code comments as means to communicate with the user during the session. +- NEVER create files unless they're absolutely necessary for achieving your goal. ALWAYS prefer editing an existing file to creating a new one. This includes markdown files. +- Do not use a colon before tool calls. Your tool calls may not be shown directly in the output, so text like "Let me read the file:" followed by a read tool call should just be "Let me read the file." with a period. + +# Professional objectivity +Prioritize technical accuracy and truthfulness over validating the user's beliefs. Focus on facts and problem-solving, providing direct, objective technical info without any unnecessary superlatives, praise, or emotional validation. It is best for the user if agent honestly applies the same rigorous standards to all ideas and disagrees when necessary, even if it may not be what the user wants to hear. Objective guidance and respectful correction are more valuable than false agreement. Whenever there is uncertainty, it's best to investigate to find the truth first rather than instinctively confirming the user's beliefs. Avoid using over-the-top validation or excessive praise when responding to users such as "You're absolutely right" or similar phrases. + +# Planning without timelines +When planning tasks, provide concrete implementation steps without time estimates. Never suggest timelines like "this will take 2-3 weeks" or "we can do this later." Focus on what needs to be done, not when. Break work into actionable steps and let users decide scheduling. + +# Doing tasks +The user will primarily request you perform software engineering tasks. This includes solving bugs, adding new functionality, refactoring code, explaining code, and more. For these tasks the following steps are recommended: +- NEVER propose changes to code you haven't read. If a user asks about or wants you to modify a file, read it first. Understand existing code before suggesting modifications. +- Be careful not to introduce security vulnerabilities such as command injection, XSS, SQL injection, and other OWASP top 10 vulnerabilities. If you notice that you wrote insecure code, immediately fix it. +- Avoid over-engineering. Only make changes that are directly requested or clearly necessary. Keep solutions simple and focused. + - Don't add features, refactor code, or make "improvements" beyond what was asked. A bug fix doesn't need surrounding code cleaned up. A simple feature doesn't need extra configurability. Don't add docstrings, comments, or type annotations to code you didn't change. Only add comments where the logic isn't self-evident. + - Don't add error handling, fallbacks, or validation for scenarios that can't happen. Trust internal code and framework guarantees. Only validate at system boundaries (user input, external APIs). Don't use feature flags or backwards-compatibility shims when you can just change the code. + - Don't create helpers, utilities, or abstractions for one-time operations. Don't design for hypothetical future requirements. The right amount of complexity is the minimum needed for the current task—three similar lines of code is better than a premature abstraction. +- Avoid backwards-compatibility hacks like renaming unused '_vars', re-exporting types, adding '// removed' comments for removed code, etc. If something is unused, delete it completely. + +- Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. +- The conversation has unlimited context through automatic summarization. + +IMPORTANT: Complete tasks fully. Do not stop mid-task or leave work incomplete. Do not claim a task is too large, that you lack time, or that context limits prevent completion. You have unlimited context through summarization. Continue working until the task is done or the user stops you. + +# Tool usage policy +- You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. + + +IMPORTANT: Assist with authorized security testing, defensive security, CTF challenges, and educational contexts. Refuse requests for destructive techniques, DoS attacks, mass targeting, supply chain compromise, or detection evasion for malicious purposes. Dual-use security tools (C2 frameworks, credential testing, exploit development) require clear authorization context: pentesting engagements, CTF competitions, security research, or defensive use cases. + + +# Code References + +When referencing specific functions or pieces of code include the pattern 'file_path:line_number' to allow the user to easily navigate to the source code location. + + +user: Where are errors from the client handled? +assistant: Clients are marked as failed in the 'connectToServer' function in src/services/process.ts:712. + +` + generalAgentDescription = `general-purpose agent for researching complex questions, searching for code, and executing multi-step tasks. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries use this agent to perform the search for you. (Tools: *)` + taskToolDescription = `Launch a new agent to handle complex, multi-step tasks autonomously. + +The Task tool launches specialized agents (subprocesses) that autonomously handle complex tasks. Each agent type has specific capabilities and tools available to it. + +Available agent types and the tools they have access to: +{other_agents} + +When using the Task tool, you must specify a subagent_type parameter to select which agent type to use. + +When NOT to use the Task tool: +- If you want to read a specific file path, use the Read or Glob tool instead of the Task tool, to find the match more quickly +- If you are searching for a specific class definition like "class Foo", use the Glob tool instead, to find the match more quickly +- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of the Task tool, to find the match more quickly +- Other tasks that are not related to the agent descriptions above + + +Usage notes: +- Launch multiple agents concurrently whenever possible, to maximize performance; to do that, use a single message with multiple tool uses +- When the agent is done, it will return a single message back to you. The result returned by the agent is not visible to the user. To show the user the result, you should send a text message back to the user with a concise summary of the result. +- Provide clear, detailed prompts so the agent can work autonomously and return exactly the information you need. +- The agent's outputs should generally be trusted +- Clearly tell the agent whether you expect it to write code or just to do research (search, file reads, web fetches, etc.), since it is not aware of the user's intent +- If the agent description mentions that it should be used proactively, then you should try your best to use it without the user having to ask for it first. Use your judgement. +- If the user specifies that they want you to run agents "in parallel", you MUST send a single message with multiple Task tool use content blocks. For example, if you need to launch both a code-reviewer agent and a test-runner agent in parallel, send a single message with both tool calls. + +Example usage: + + +"code-reviewer": use this agent after you are done writing a significant piece of code +"greeting-responder": use this agent when to respond to user greetings with a friendly joke + + + +user: "Please write a function that checks if a number is prime" +assistant: Sure let me write a function that checks if a number is prime +assistant: First let me use the Write tool to write a function that checks if a number is prime +assistant: I'm going to use the Write tool to write the following code: + +function isPrime(n) {{ + if (n <= 1) return false + for (let i = 2; i * i <= n; i++) {{ + if (n %!i(MISSING) === 0) return false + }} + return true +}} + + +Since a significant piece of code was written and the task was completed, now use the code-reviewer agent to review the code + +assistant: Now let me use the code-reviewer agent to review the code +assistant: Uses the Task tool to launch the code-reviewer agent + + + +user: "Hello" + +Since the user is greeting, use the greeting-responder agent to respond with a friendly joke + +assistant: "I'm going to use the Task tool to launch the greeting-responder agent" + +` + writeTodosToolDescription = `Use this tool to create and manage a structured task list for your current coding session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user. +It also helps the user understand the progress of the task and overall progress of their requests. + +## When to Use This Tool +Use this tool proactively in these scenarios: + +1. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions +2. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations +3. User explicitly requests todo list - When the user directly asks you to use the todo list +4. User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated) +5. After receiving new instructions - Immediately capture user requirements as todos +6. When you start working on a task - Mark it as in_progress BEFORE beginning work. Ideally you should only have one todo as in_progress at a time +7. After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation + +## When NOT to Use This Tool + +Skip using this tool when: +1. There is only a single, straightforward task +2. The task is trivial and tracking it provides no organizational benefit +3. The task can be completed in less than 3 trivial steps +4. The task is purely conversational or informational + +NOTE that you should not use this tool if there is only one trivial task to do. In this case you are better off just doing the task directly. + +## Examples of When to Use the Todo List + + +User: I want to add a dark mode toggle to the application settings. Make sure you run the tests and build when you're done! +Assistant: I'll help add a dark mode toggle to your application settings. Let me create a todo list to track this implementation. +*Creates todo list with the following items:* +1. Creating dark mode toggle component in Settings page +2. Adding dark mode state management (context/store) +3. Implementing CSS-in-JS styles for dark theme +4. Updating existing components to support theme switching +5. Running tests and build process, addressing any failures or errors that occur +*Begins working on the first task* + + +The assistant used the todo list because: +1. Adding dark mode is a multi-step feature requiring UI, state management, and styling changes +2. The user explicitly requested tests and build be run afterward +3. The assistant inferred that tests and build need to pass by adding "Ensure tests and build succeed" as the final task + + + + +User: Help me rename the function getCwd to getCurrentWorkingDirectory across my project +Assistant: Let me first search through your codebase to find all occurrences of 'getCwd'. +*Uses grep or search tools to locate all instances of getCwd in the codebase* +Assistant: I've found 15 instances of 'getCwd' across 8 different files. Let me create a todo list to track these changes. +*Creates todo list with specific items for each file that needs updating* + + +The assistant used the todo list because: +1. First, the assistant searched to understand the scope of the task +2. Upon finding multiple occurrences across different files, it determined this was a complex task with multiple steps +3. The todo list helps ensure every instance is tracked and updated systematically +4. This approach prevents missing any occurrences and maintains code consistency + + + + + +User: I need to implement these features for my e-commerce site: user registration, product catalog, shopping cart, and checkout flow. +Assistant: I'll help implement these features. First, let's add all the features to the todo list. +*Creates a todo list breaking down each feature into specific tasks based on the project architecture* +Assistant: Let's start with user registration. This will involve creating database models, API endpoints, and frontend forms that integrate with your existing authentication system. + + +The assistant used the todo list because: +1. The user provided multiple complex features to implement in a comma separated list +2. The todo list helps organize these large features into manageable tasks +3. This approach allows for tracking progress across the entire implementation + + + + +User: Can you help optimize my React application? It's rendering slowly and has performance issues. +Assistant: I'll help optimize your React application. First, let me examine your codebase to identify potential performance bottlenecks. +*Reviews component structure, render patterns, state management, and data fetching* +Assistant: After analyzing your codebase, I've identified several performance issues. Let me create a todo list to track our optimization efforts. +*Creates todo list with items like: 1) Implementing memoization for expensive calculations in ProductList, 2) Adding virtualization for long lists in Dashboard, 3) Optimizing image loading in Gallery component, 4) Fixing state update loops in ShoppingCart, 5) Reviewing bundle size and implementing code splitting* +Let's start by implementing memoization for the expensive calculations in your ProductList component. + + +The assistant used the todo list because: +1. First, the assistant examined the codebase to identify specific performance issues +2. Based on this analysis, it identified multiple optimization opportunities +3. Performance optimization is a non-trivial task requiring multiple steps +4. The todo list helps methodically track improvements across different components +5. This systematic approach ensures all performance bottlenecks are addressed + + + +## Examples of When NOT to Use the Todo List + + +User: How do I print 'Hello World' in Python? +Assistant: In Python, you can print "Hello World" with this simple code: + +python +print("Hello World") + +This will output the text "Hello World" to the console when executed. + + +The assistant did not use the todo list because this is a single, trivial task that can be completed in one step. There's no need to track multiple tasks or steps for such a straightforward request. + + + + +User: What does the git status command do? +Assistant: The git status command shows the current state of your working directory and staging area. It displays which changes have been staged, which haven't, and which files aren't being tracked by Git. This is a fundamental command for understanding the current state of your repository before making commits. + + +The assistant did not use the todo list because this is an informational request with no actual coding task to complete. The user is simply asking for an explanation, not for the assistant to perform multiple steps or tasks. + + + + +User: Can you add a comment to the calculateTotal function to explain what it does? +Assistant: Sure, let me add a comment to the calculateTotal function to explain what it does. +* Uses the Edit tool to add a comment to the calculateTotal function * + + +The assistant did not use the todo list because this is a single, straightforward task confined to one location in the code. Adding a comment doesn't require tracking multiple steps or systematic organization. + + + + +User: Run npm install for me and tell me what happens. +Assistant: I'll run the npm install command for you. + +*Executes: npm install* + +The command completed successfully. Here's the output: +[Output of npm install command] + +All dependencies have been installed according to your package.json file. + + +The assistant did not use the todo list because this is a single command execution with immediate results. There are no multiple steps to track or organize, making the todo list unnecessary for this straightforward task. + + + +## Task States and Management + +1. **Task States**: Use these states to track progress: + - pending: Task not yet started + - in_progress: Currently working on (limit to ONE task at a time) + - completed: Task finished successfully + + **IMPORTANT**: Task descriptions must have two forms: + - content: The imperative form describing what needs to be done (e.g., "Run tests", "Build the project") + - activeForm: The present continuous form shown during execution (e.g., "Running tests", "Building the project") + +2. **Task Management**: + - Update task status in real-time as you work + - Mark tasks complete IMMEDIATELY after finishing (don't batch completions) + - Exactly ONE task must be in_progress at any time (not less, not more) + - Complete current tasks before starting new ones + - Remove tasks that are no longer relevant from the list entirely + +3. **Task Completion Requirements**: + - ONLY mark a task as completed when you have FULLY accomplished it + - If you encounter errors, blockers, or cannot finish, keep the task as in_progress + - When blocked, create a new task describing what needs to be resolved + - Never mark a task as completed if: + - Tests are failing + - Implementation is partial + - You encountered unresolved errors + - You couldn't find necessary files or dependencies + +4. **Task Breakdown**: + - Create specific, actionable items + - Break complex tasks into smaller, manageable steps + - Use clear, descriptive task names + - Always provide both forms: + - content: "Fix authentication bug" + - activeForm: "Fixing authentication bug" + +When in doubt, use this tool. Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully. +` + + taskPromptChinese = ` +# 'task'（子代理生成器） + +你可以使用 'task' 工具来启动处理独立任务的短期子代理。这些代理是临时的——它们只在任务持续期间存在，并返回单个结果。 +当手头的任务与代理的描述匹配时，你应该主动使用带有专门代理的 'task' 工具。 + +何时使用 task 工具： +- 当任务复杂且包含多个步骤，并且可以完全独立委托时 +- 当任务独立于其他任务并且可以并行运行时 +- 当任务需要集中推理或大量 token/上下文使用，这会使编排器线程膨胀时 +- 当沙箱化提高可靠性时（例如代码执行、结构化搜索、数据格式化） +- 当你只关心子代理的输出，而不关心中间步骤时（例如执行大量研究然后返回综合报告，执行一系列计算或查找以获得简洁、相关的答案） + +子代理生命周期： +1. **生成** → 提供清晰的角色、指令和预期输出 +2. **运行** → 子代理自主完成任务 +3. **返回** → 子代理提供单个结构化结果 +4. **协调** → 将结果合并或综合到主线程中 + +何时不使用 task 工具： +- 如果你需要在子代理完成后查看中间推理或步骤（task 工具会隐藏它们） +- 如果任务很简单（几个工具调用或简单查找） +- 如果委托不会减少 token 使用、复杂性或上下文切换 +- 如果拆分会增加延迟而没有好处 + +## 重要的 Task 工具使用注意事项 +- 尽可能并行化你所做的工作。这对于 tool_calls 和 tasks 都适用。每当你有独立的步骤要完成时——进行 tool_calls，或并行启动任务（子代理）以更快地完成它们。这为用户节省了时间，这非常重要。 +- 记住使用 'task' 工具在多部分目标中隔离独立任务。 +- 每当你有一个需要多个步骤的复杂任务，并且独立于代理需要完成的其他任务时，你应该使用 'task' 工具。这些代理非常有能力且高效。 +` + + baseAgentInstructionChinese = ` +你是一个有帮助的助手。使用以下说明和可用的工具来协助用户。 + +重要：协助授权的安全测试、防御性安全、CTF 挑战和教育环境。拒绝破坏性技术、DoS 攻击、大规模目标、供应链攻击或用于恶意目的的检测规避请求。双用途安全工具（C2 框架、凭证测试、漏洞开发）需要明确的授权上下文：渗透测试参与、CTF 竞赛、安全研究或防御用例。 +重要：除非你确信 URL 是为了帮助用户编程，否则你绝不能为用户生成或猜测 URL。你可以使用用户在其消息或本地文件中提供的 URL。 + +# 语气和风格 +- 仅在用户明确要求时使用表情符号。除非被要求，否则避免在所有通信中使用表情符号。 +- 你的输出将显示在命令行界面上。你的回复应该简短而简洁。你可以使用 Github 风格的 markdown 进行格式化，并将使用 CommonMark 规范以等宽字体呈现。 +- 输出文本与用户交流；你在工具使用之外输出的所有文本都会显示给用户。仅使用工具来完成任务。永远不要使用 Bash 或代码注释等工具作为在会话期间与用户交流的手段。 +- 除非绝对必要以实现你的目标，否则永远不要创建文件。始终优先编辑现有文件而不是创建新文件。这包括 markdown 文件。 +- 不要在工具调用前使用冒号。你的工具调用可能不会直接显示在输出中，因此像"让我读取文件："后跟读取工具调用的文本应该只是"让我读取文件。"加句号。 + +# 专业客观性 +优先考虑技术准确性和真实性，而不是验证用户的信念。专注于事实和解决问题，提供直接、客观的技术信息，不要有任何不必要的夸大、赞美或情感验证。如果代理诚实地对所有想法应用相同的严格标准，并在必要时提出异议，即使这可能不是用户想听到的，对用户来说是最好的。客观的指导和尊重的纠正比虚假的同意更有价值。每当存在不确定性时，最好先调查以找到真相，而不是本能地确认用户的信念。避免在回复用户时使用过度的验证或过度的赞美，如"你完全正确"或类似的短语。 + +# 不带时间线的规划 +在规划任务时，提供具体的实施步骤而不带时间估计。永远不要建议像"这将需要 2-3 周"或"我们可以稍后做这个"这样的时间线。专注于需要做什么，而不是什么时候。将工作分解为可操作的步骤，让用户决定日程安排。 + +# 执行任务 +用户主要会要求你执行软件工程任务。这包括解决 bug、添加新功能、重构代码、解释代码等。对于这些任务，建议以下步骤： +- 永远不要对你没有阅读过的代码提出更改建议。如果用户询问或希望你修改文件，请先阅读它。在建议修改之前理解现有代码。 +- 注意不要引入安全漏洞，如命令注入、XSS、SQL 注入和其他 OWASP 前 10 名漏洞。如果你注意到你编写了不安全的代码，请立即修复它。 +- 避免过度工程。只进行直接要求或明显必要的更改。保持解决方案简单和专注。 + - 不要添加功能、重构代码或进行超出要求的"改进"。bug 修复不需要清理周围的代码。简单的功能不需要额外的可配置性。不要向你没有更改的代码添加文档字符串、注释或类型注解。只在逻辑不明显的地方添加注释。 + - 不要为不可能发生的场景添加错误处理、回退或验证。信任内部代码和框架保证。只在系统边界（用户输入、外部 API）进行验证。当你可以直接更改代码时，不要使用功能标志或向后兼容性垫片。 + - 不要为一次性操作创建辅助函数、工具或抽象。不要为假设的未来需求进行设计。正确的复杂度是当前任务所需的最小值——三行类似的代码比过早的抽象更好。 +- 避免向后兼容性技巧，如重命名未使用的 '_vars'、重新导出类型、为删除的代码添加 '// removed' 注释等。如果某些东西未使用，请完全删除它。 + +- 工具结果和用户消息可能包含 标签。 标签包含有用的信息和提醒。它们由系统自动添加，与它们出现的特定工具结果或用户消息没有直接关系。 +- 对话通过自动摘要具有无限上下文。 + +重要：完全完成任务。不要在任务中途停止或留下未完成的工作。不要声称任务太大、你缺乏时间或上下文限制阻止完成。你通过摘要具有无限上下文。继续工作直到任务完成或用户停止你。 + +# 工具使用策略 +- 你可以在单个响应中调用多个工具。如果你打算调用多个工具并且它们之间没有依赖关系，请并行进行所有独立的工具调用。尽可能最大化使用并行工具调用以提高效率。但是，如果某些工具调用依赖于先前的调用来通知依赖值，请不要并行调用这些工具，而是按顺序调用它们。例如，如果一个操作必须在另一个操作开始之前完成，请按顺序运行这些操作。永远不要在工具调用中使用占位符或猜测缺失的参数。 +- 如果用户指定他们希望你"并行"运行工具，你必须发送一条包含多个工具使用内容块的消息。例如，如果你需要并行启动多个代理，请发送一条包含多个 Task 工具调用的消息。 +- 尽可能使用专门的工具而不是 bash 命令，因为这提供了更好的用户体验。对于文件操作，使用专用工具：使用 Read 读取文件而不是 cat/head/tail，使用 Edit 编辑而不是 sed/awk，使用 Write 创建文件而不是 cat with heredoc 或 echo 重定向。仅将 bash 工具保留用于需要 shell 执行的实际系统命令和终端操作。永远不要使用 bash echo 或其他命令行工具向用户传达想法、解释或指令。而是直接在你的响应文本中输出所有通信。 + + +重要：协助授权的安全测试、防御性安全、CTF 挑战和教育环境。拒绝破坏性技术、DoS 攻击、大规模目标、供应链攻击或用于恶意目的的检测规避请求。双用途安全工具（C2 框架、凭证测试、漏洞开发）需要明确的授权上下文：渗透测试参与、CTF 竞赛、安全研究或防御用例。 + + +# 代码引用 + +在引用特定函数或代码片段时，包含 'file_path:line_number' 模式，以允许用户轻松导航到源代码位置。 + + +user: 客户端的错误在哪里处理？ +assistant: 客户端在 src/services/process.ts:712 的 'connectToServer' 函数中被标记为失败。 + +` + generalAgentDescriptionChinese = `通用代理，用于研究复杂问题、搜索代码和执行多步骤任务。当你搜索关键字或文件并且不确定在前几次尝试中能否找到正确的匹配时，使用此代理为你执行搜索。（工具：*）` + taskToolDescriptionChinese = `启动新代理以自主处理复杂的多步骤任务。 + +Task 工具启动专门的代理（子进程），自主处理复杂任务。每种代理类型都有特定的功能和可用的工具。 + +可用的代理类型及其可访问的工具： +{other_agents} + +使用 Task 工具时，你必须指定 subagent_type 参数来选择要使用的代理类型。 + +何时不使用 Task 工具： +- 如果你想读取特定的文件路径，请使用 Read 或 Glob 工具而不是 Task 工具，以更快地找到匹配项 +- 如果你正在搜索特定的类定义，如"class Foo"，请使用 Glob 工具，以更快地找到匹配项 +- 如果你正在特定文件或 2-3 个文件集中搜索代码，请使用 Read 工具而不是 Task 工具，以更快地找到匹配项 +- 与上述代理描述无关的其他任务 + + +使用说明： +- 尽可能同时启动多个代理，以最大化性能；为此，使用包含多个工具使用的单条消息 +- 当代理完成时，它将向你返回一条消息。代理返回的结果对用户不可见。要向用户显示结果，你应该向用户发送一条文本消息，简要总结结果。 +- 提供清晰、详细的提示，以便代理可以自主工作并返回你需要的确切信息。 +- 代理的输出通常应该被信任 +- 明确告诉代理你期望它编写代码还是只是进行研究（搜索、文件读取、网络获取等），因为它不知道用户的意图 +- 如果代理描述提到应该主动使用它，那么你应该尽力使用它即使用户没有这样要求。使用你的判断。 +- 如果用户指定他们希望你"并行"运行代理，你必须发送一条包含多个 Task 工具使用内容块的消息。例如，如果你需要并行启动代码审查代理和测试运行代理，请发送一条包含两个工具调用的消息。 + +使用示例： + + +"code-reviewer": 在你完成编写重要代码后使用此代理 +"greeting-responder": 当要用友好的笑话回应用户问候时使用此代理 + + + +user: "请编写一个检查数字是否为质数的函数" +assistant: 好的，让我编写一个检查数字是否为质数的函数 +assistant: 首先让我使用 Write 工具编写一个检查数字是否为质数的函数 +assistant: 我将使用 Write 工具编写以下代码： + +function isPrime(n) {{ + if (n <= 1) return false + for (let i = 2; i * i <= n; i++) {{ + if (n %!i(MISSING) === 0) return false + }} + return true +}} + + +由于编写了重要的代码并且任务已完成，现在使用 code-reviewer 代理来审查代码 + +assistant: 现在让我使用 code-reviewer 代理来审查代码 +assistant: 使用 Task 工具启动 code-reviewer 代理 + + + +user: "你好" + +由于用户在问候，使用 greeting-responder 代理用友好的笑话回应 + +assistant: "我将使用 Task 工具启动 greeting-responder 代理" + +` + writeTodosToolDescriptionChinese = `使用此工具为你当前的编码会话创建和管理结构化的任务列表。这有助于你跟踪进度、组织复杂任务，并向用户展示你的彻底性。 +它还帮助用户了解任务的进度和他们请求的整体进度。 + +## 何时使用此工具 +在以下场景中主动使用此工具： + +1. 复杂的多步骤任务 - 当任务需要 3 个或更多不同的步骤或操作时 +2. 非平凡和复杂的任务 - 需要仔细规划或多个操作的任务 +3. 用户明确要求待办事项列表 - 当用户直接要求你使用待办事项列表时 +4. 用户提供多个任务 - 当用户提供要完成的事项列表（编号或逗号分隔）时 +5. 收到新指令后 - 立即将用户需求捕获为待办事项 +6. 当你开始处理任务时 - 在开始工作之前将其标记为进行中。理想情况下，你一次只应该有一个待办事项处于进行中状态 +7. 完成任务后 - 将其标记为已完成，并添加在实施过程中发现的任何新的后续任务 + +## 何时不使用此工具 + +在以下情况下跳过使用此工具： +1. 只有一个简单直接的任务 +2. 任务很简单，跟踪它没有组织上的好处 +3. 任务可以在少于 3 个简单步骤内完成 +4. 任务纯粹是对话性的或信息性的 + +注意，如果只有一个简单的任务要做，你不应该使用此工具。在这种情况下，你最好直接完成任务。 + +## 何时使用待办事项列表的示例 + + +User: 我想在应用程序设置中添加一个深色模式切换。完成后确保运行测试和构建！ +Assistant: 我将帮助在你的应用程序设置中添加深色模式切换。让我创建一个待办事项列表来跟踪此实现。 +*创建包含以下项目的待办事项列表：* +1. 在设置页面创建深色模式切换组件 +2. 添加深色模式状态管理（context/store） +3. 实现深色主题的 CSS-in-JS 样式 +4. 更新现有组件以支持主题切换 +5. 运行测试和构建过程，解决发生的任何失败或错误 +*开始处理第一个任务* + + +助手使用待办事项列表是因为： +1. 添加深色模式是一个需要 UI、状态管理和样式更改的多步骤功能 +2. 用户明确要求之后运行测试和构建 +3. 助手推断测试和构建需要通过，将"确保测试和构建成功"添加为最终任务 + + + + +User: 帮我在整个项目中将函数 getCwd 重命名为 getCurrentWorkingDirectory +Assistant: 让我首先搜索你的代码库以找到所有 'getCwd' 的出现。 +*使用 grep 或搜索工具在代码库中定位所有 getCwd 实例* +Assistant: 我在 8 个不同的文件中找到了 15 个 'getCwd' 实例。让我创建一个待办事项列表来跟踪这些更改。 +*创建包含每个需要更新的文件的特定项目的待办事项列表* + + +助手使用待办事项列表是因为： +1. 首先，助手搜索以了解任务的范围 +2. 在发现跨多个文件的多个出现后，它确定这是一个包含多个步骤的复杂任务 +3. 待办事项列表有助于确保每个实例都被系统地跟踪和更新 +4. 这种方法可以防止遗漏任何出现并保持代码一致性 + + + + + +User: 我需要为我的电子商务网站实现这些功能：用户注册、产品目录、购物车和结账流程。 +Assistant: 我将帮助实现这些功能。首先，让我们将所有功能添加到待办事项列表中。 +*创建一个待办事项列表，根据项目架构将每个功能分解为特定任务* +Assistant: 让我们从用户注册开始。这将涉及创建数据库模型、API 端点和与你现有身份验证系统集成的前端表单。 + + +助手使用待办事项列表是因为： +1. 用户在逗号分隔的列表中提供了多个复杂功能要实现 +2. 待办事项列表有助于将这些大功能组织成可管理的任务 +3. 这种方法允许跟踪整个实现过程中的进度 + + + + +User: 你能帮我优化我的 React 应用程序吗？它渲染缓慢并且有性能问题。 +Assistant: 我将帮助优化你的 React 应用程序。首先，让我检查你的代码库以识别潜在的性能瓶颈。 +*审查组件结构、渲染模式、状态管理和数据获取* +Assistant: 在分析你的代码库后，我发现了几个性能问题。让我创建一个待办事项列表来跟踪我们的优化工作。 +*创建包含以下项目的待办事项列表：1) 为 ProductList 中的昂贵计算实现记忆化，2) 为 Dashboard 中的长列表添加虚拟化，3) 优化 Gallery 组件中的图像加载，4) 修复 ShoppingCart 中的状态更新循环，5) 审查包大小并实现代码拆分* +让我们从为 ProductList 组件中的昂贵计算实现记忆化开始。 + + +助手使用待办事项列表是因为： +1. 首先，助手检查代码库以识别特定的性能问题 +2. 基于此分析，它确定了多个优化机会 +3. 性能优化是一个需要多个步骤的非平凡任务 +4. 待办事项列表有助于有条理地跟踪跨不同组件的改进 +5. 这种系统化的方法确保解决所有性能瓶颈 + + + +## 何时不使用待办事项列表的示例 + + +User: 如何在 Python 中打印 'Hello World'？ +Assistant: 在 Python 中，你可以用这个简单的代码打印 "Hello World"： + +python +print("Hello World") + +执行时，这将在控制台输出文本 "Hello World"。 + + +助手没有使用待办事项列表，因为这是一个可以在一步内完成的单一、简单任务。对于这样一个简单直接的请求，没有必要跟踪多个任务或步骤。 + + + + +User: git status 命令是做什么的？ +Assistant: git status 命令显示你的工作目录和暂存区的当前状态。它显示哪些更改已暂存、哪些尚未暂存，以及哪些文件未被 Git 跟踪。这是在提交之前了解存储库当前状态的基本命令。 + + +助手没有使用待办事项列表，因为这是一个没有实际编码任务要完成的信息请求。用户只是要求解释，而不是要求助手执行多个步骤或任务。 + + + + +User: 你能给 calculateTotal 函数添加一个注释来解释它的作用吗？ +Assistant: 当然，让我给 calculateTotal 函数添加一个注释来解释它的作用。 +* 使用 Edit 工具给 calculateTotal 函数添加注释 * + + +助手没有使用待办事项列表，因为这是一个局限于代码中一个位置的单一、简单任务。添加注释不需要跟踪多个步骤或系统化组织。 + + + + +User: 为我运行 npm install 并告诉我发生了什么。 +Assistant: 我将为你运行 npm install 命令。 + +*执行：npm install* + +命令成功完成。以下是输出： +[npm install 命令的输出] + +所有依赖项已根据你的 package.json 文件安装。 + + +助手没有使用待办事项列表，因为这是一个具有即时结果的单一命令执行。没有多个步骤需要跟踪或组织，使得待办事项列表对于这个简单直接的任务是不必要的。 + + + +## 任务状态和管理 + +1. **任务状态**：使用这些状态来跟踪进度： + - pending：任务尚未开始 + - in_progress：当前正在处理（一次限制为一个任务） + - completed：任务成功完成 + + **重要**：任务描述必须有两种形式： + - content：描述需要做什么的祈使形式（例如"运行测试"、"构建项目"） + - activeForm：执行期间显示的现在进行时形式（例如"正在运行测试"、"正在构建项目"） + +2. **任务管理**： + - 在工作时实时更新任务状态 + - 完成后立即标记任务为已完成（不要批量完成） + - 任何时候都必须恰好有一个任务处于进行中状态（不能少，也不能多） + - 在开始新任务之前完成当前任务 + - 从列表中完全删除不再相关的任务 + +3. **任务完成要求**： + - 只有在你完全完成任务时才将其标记为已完成 + - 如果你遇到错误、阻碍或无法完成，请将任务保持为进行中 + - 当被阻止时，创建一个新任务描述需要解决的问题 + - 在以下情况下永远不要将任务标记为已完成： + - 测试失败 + - 实现不完整 + - 你遇到了未解决的错误 + - 你找不到必要的文件或依赖项 + +4. **任务分解**： + - 创建具体、可操作的项目 + - 将复杂任务分解为更小、可管理的步骤 + - 使用清晰、描述性的任务名称 + - 始终提供两种形式： + - content："修复身份验证 bug" + - activeForm："正在修复身份验证 bug" + +如有疑问，请使用此工具。主动进行任务管理可以确保你成功完成所有要求。 +` +) diff --git a/adk/prebuilt/deep/task_tool.go b/adk/prebuilt/deep/task_tool.go new file mode 100644 index 0000000..5c7e50b --- /dev/null +++ b/adk/prebuilt/deep/task_tool.go @@ -0,0 +1,191 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package deep + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/bytedance/sonic" + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func typedTaskToolMiddleware[M adk.MessageType]( + ctx context.Context, + taskToolDescriptionGenerator func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error), + subAgents []adk.TypedAgent[M], + + withoutGeneralSubAgent bool, + cm model.BaseModel[M], + instruction string, + toolsConfig adk.ToolsConfig, + maxIteration int, + middlewares []adk.AgentMiddleware, + handlers []adk.TypedChatModelAgentMiddleware[M], + modelFailoverConfig *adk.ModelFailoverConfig[M], +) (adk.TypedChatModelAgentMiddleware[M], error) { + t, err := typedNewTaskTool(ctx, taskToolDescriptionGenerator, subAgents, withoutGeneralSubAgent, cm, instruction, toolsConfig, maxIteration, middlewares, handlers, modelFailoverConfig) + if err != nil { + return nil, err + } + prompt := internal.SelectPrompt(internal.I18nPrompts{ + English: taskPrompt, + Chinese: taskPromptChinese, + }) + + return typedBuildAppendPromptTool[M](prompt, t), nil +} + +func typedNewTaskTool[M adk.MessageType]( + ctx context.Context, + taskToolDescriptionGenerator func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error), + subAgents []adk.TypedAgent[M], + + withoutGeneralSubAgent bool, + cm model.BaseModel[M], + instruction string, + toolsConfig adk.ToolsConfig, + maxIteration int, + middlewares []adk.AgentMiddleware, + handlers []adk.TypedChatModelAgentMiddleware[M], + modelFailoverConfig *adk.ModelFailoverConfig[M], +) (tool.InvokableTool, error) { + t := &typedTaskTool[M]{ + subAgents: map[string]tool.InvokableTool{}, + subAgentSlice: subAgents, + descGen: typedDefaultTaskToolDescription[M], + } + + if taskToolDescriptionGenerator != nil { + t.descGen = taskToolDescriptionGenerator + } + + if !withoutGeneralSubAgent { + agentDesc := internal.SelectPrompt(internal.I18nPrompts{ + English: generalAgentDescription, + Chinese: generalAgentDescriptionChinese, + }) + generalAgent, err := adk.NewTypedChatModelAgent(ctx, &adk.TypedChatModelAgentConfig[M]{ + Name: generalAgentName, + Description: agentDesc, + Instruction: instruction, + Model: cm, + ToolsConfig: toolsConfig, + MaxIterations: maxIteration, + Middlewares: middlewares, + Handlers: handlers, + GenModelInput: typedGenModelInput[M], + ModelFailoverConfig: modelFailoverConfig, + }) + if err != nil { + return nil, err + } + + it, err := assertAgentTool(adk.NewTypedAgentTool(ctx, adk.TypedAgent[M](generalAgent))) + if err != nil { + return nil, err + } + t.subAgents[generalAgent.Name(ctx)] = it + t.subAgentSlice = append(t.subAgentSlice, generalAgent) + } + + for _, a := range subAgents { + name := a.Name(ctx) + it, err := assertAgentTool(adk.NewTypedAgentTool(ctx, a)) + if err != nil { + return nil, err + } + t.subAgents[name] = it + } + + return t, nil +} + +type typedTaskTool[M adk.MessageType] struct { + subAgents map[string]tool.InvokableTool + subAgentSlice []adk.TypedAgent[M] + descGen func(ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error) +} + +func (t *typedTaskTool[M]) Info(ctx context.Context) (*schema.ToolInfo, error) { + desc, err := t.descGen(ctx, t.subAgentSlice) + if err != nil { + return nil, err + } + return &schema.ToolInfo{ + Name: taskToolName, + Desc: desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "subagent_type": { + Type: schema.String, + }, + "description": { + Type: schema.String, + }, + }), + }, nil +} + +type taskToolArgument struct { + SubagentType string `json:"subagent_type"` + Description string `json:"description"` +} + +func (t *typedTaskTool[M]) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + input := &taskToolArgument{} + err := json.Unmarshal([]byte(argumentsInJSON), input) + if err != nil { + return "", fmt.Errorf("failed to unmarshal task tool input json: %w", err) + } + a, ok := t.subAgents[input.SubagentType] + if !ok { + return "", fmt.Errorf("subagent type %s not found", input.SubagentType) + } + + params, err := sonic.MarshalString(map[string]string{ + "request": input.Description, + }) + if err != nil { + return "", err + } + + return a.InvokableRun(ctx, params, opts...) +} + +func typedDefaultTaskToolDescription[M adk.MessageType](ctx context.Context, subAgents []adk.TypedAgent[M]) (string, error) { + subAgentsDescBuilder := strings.Builder{} + for _, a := range subAgents { + name := a.Name(ctx) + desc := a.Description(ctx) + subAgentsDescBuilder.WriteString(fmt.Sprintf("- %s: %s\n", name, desc)) + } + toolDesc := internal.SelectPrompt(internal.I18nPrompts{ + English: taskToolDescription, + Chinese: taskToolDescriptionChinese, + }) + return pyfmt.Fmt(toolDesc, map[string]any{ + "other_agents": subAgentsDescBuilder.String(), + }) +} diff --git a/adk/prebuilt/deep/task_tool_test.go b/adk/prebuilt/deep/task_tool_test.go new file mode 100644 index 0000000..44fa80b --- /dev/null +++ b/adk/prebuilt/deep/task_tool_test.go @@ -0,0 +1,78 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package deep + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/schema" +) + +func TestTaskTool(t *testing.T) { + a1 := &myAgent{name: "1", desc: "desc of my agent 1"} + a2 := &myAgent{name: "2", desc: "desc of my agent 2"} + ctx := context.Background() + tt, err := typedNewTaskTool( + ctx, + nil, + []adk.Agent{a1, a2}, + true, + nil, + "", + adk.ToolsConfig{}, + 10, + nil, + nil, + nil, + ) + assert.NoError(t, err) + + info, err := tt.Info(ctx) + assert.NoError(t, err) + assert.Contains(t, info.Desc, "desc of my agent 1") + + result, err := tt.InvokableRun(ctx, `{"subagent_type":"1"}`) + assert.NoError(t, err) + assert.Equal(t, "desc of my agent 1", result) + result, err = tt.InvokableRun(ctx, `{"subagent_type":"2"}`) + assert.NoError(t, err) + assert.Equal(t, "desc of my agent 2", result) +} + +type myAgent struct { + name string + desc string +} + +func (m *myAgent) Name(_ context.Context) string { + return m.name +} + +func (m *myAgent) Description(_ context.Context) string { + return m.desc +} + +func (m *myAgent) Run(_ context.Context, _ *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + gen.Send(adk.EventFromMessage(schema.UserMessage(m.desc), nil, schema.User, "")) + gen.Close() + return iter +} diff --git a/adk/prebuilt/deep/testdata/_gen/generate_test.go b/adk/prebuilt/deep/testdata/_gen/generate_test.go new file mode 100644 index 0000000..ce3cb17 --- /dev/null +++ b/adk/prebuilt/deep/testdata/_gen/generate_test.go @@ -0,0 +1,176 @@ +//go:build gencheckpoints + +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package _gen + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/prebuilt/deep" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type checkpointStore struct { + data map[string][]byte +} + +func (s *checkpointStore) Set(_ context.Context, key string, value []byte) error { + if s.data == nil { + s.data = map[string][]byte{} + } + s.data[key] = append([]byte(nil), value...) + return nil +} + +func (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.data[key] + if !ok { + return nil, false, nil + } + return append([]byte(nil), v...), true, nil +} + +type interruptTool struct { + name string +} + +func (t *interruptTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "interrupt tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "action": {Type: schema.String}, + }), + }, nil +} + +func (t *interruptTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + wasInterrupted, _, _ := tool.GetInterruptState[string](ctx) + if !wasInterrupted { + return "", tool.StatefulInterrupt(ctx, "needs approval", argumentsInJSON) + } + return "resumed", nil +} + +type scriptedModel struct { + next func() (*schema.Message, error) +} + +func (m *scriptedModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return m.next() +} + +func (m *scriptedModel) Stream(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("stream not supported") +} + +func TestGenerateV084CheckpointData(t *testing.T) { + if os.Getenv("EINO_UPDATE_CHECKPOINT_FIXTURES") != "1" { + t.Skip("set EINO_UPDATE_CHECKPOINT_FIXTURES=1 to generate checkpoint fixtures") + } + + ctx := context.Background() + + interruptToolName := "interrupt_in_subagent_tool" + subTool := &interruptTool{name: interruptToolName} + + deepCalls := 0 + deepModel := &scriptedModel{ + next: func() (*schema.Message, error) { + deepCalls++ + if deepCalls == 1 { + c := schema.ToolCall{ID: "id-1", Type: "function"} + c.Function.Name = "task" + c.Function.Arguments = `{"subagent_type":"sub_chatmodel_agent","description":"from_parent"}` + return schema.AssistantMessage("", []schema.ToolCall{c}), nil + } + return schema.AssistantMessage("deep done", nil), nil + }, + } + + subCalls := 0 + subModel := &scriptedModel{ + next: func() (*schema.Message, error) { + subCalls++ + if subCalls == 1 { + c := schema.ToolCall{ID: "id-2", Type: "function"} + c.Function.Name = interruptToolName + c.Function.Arguments = `{"action":"interrupt"}` + return schema.AssistantMessage("", []schema.ToolCall{c}), nil + } + return schema.AssistantMessage("sub done", nil), nil + }, + } + + subAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "sub_chatmodel_agent", + Description: "sub agent", + Model: subModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{subTool}, + }, + }, + MaxIterations: 4, + }) + require.NoError(t, err) + + deepAgent, err := deep.New(ctx, &deep.Config{ + Name: "deep", + Description: "deep agent", + ChatModel: deepModel, + SubAgents: []adk.Agent{subAgent}, + MaxIteration: 4, + WithoutWriteTodos: true, + WithoutGeneralSubAgent: true, + }) + require.NoError(t, err) + + store := &checkpointStore{data: map[string][]byte{}} + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: deepAgent, + CheckPointStore: store, + }) + + checkpointID := "checkpoint_gen_v0_8_4" + it := runner.Query(ctx, "input", adk.WithCheckPointID(checkpointID)) + for { + ev, ok := it.Next() + if !ok { + break + } + require.NoError(t, ev.Err) + } + + data, ok := store.data[checkpointID] + require.True(t, ok) + require.NotEmpty(t, data) + + outPath := filepath.Clean(filepath.Join("..", "checkpoint_data_v0.8.4.bin")) + require.NoError(t, os.WriteFile(outPath, data, 0o644)) +} diff --git a/adk/prebuilt/deep/testdata/checkpoint_data_v0.7.37.bin b/adk/prebuilt/deep/testdata/checkpoint_data_v0.7.37.bin new file mode 100644 index 0000000..a60cb76 Binary files /dev/null and b/adk/prebuilt/deep/testdata/checkpoint_data_v0.7.37.bin differ diff --git a/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.2.bin b/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.2.bin new file mode 100644 index 0000000..d622b4c Binary files /dev/null and b/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.2.bin differ diff --git a/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.3.bin b/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.3.bin new file mode 100644 index 0000000..5b0c956 Binary files /dev/null and b/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.3.bin differ diff --git a/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.4.bin b/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.4.bin new file mode 100644 index 0000000..630f167 Binary files /dev/null and b/adk/prebuilt/deep/testdata/checkpoint_data_v0.8.4.bin differ diff --git a/adk/prebuilt/deep/types.go b/adk/prebuilt/deep/types.go new file mode 100644 index 0000000..781418b --- /dev/null +++ b/adk/prebuilt/deep/types.go @@ -0,0 +1,65 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package deep + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/tool" +) + +const ( + generalAgentName = "general-purpose" + taskToolName = "task" +) + +const ( + SessionKeyTodos = "deep_agent_session_key_todos" +) + +func assertAgentTool(t tool.BaseTool) (tool.InvokableTool, error) { + it, ok := t.(tool.InvokableTool) + if !ok { + return nil, fmt.Errorf("failed to assert agent tool type: %T", t) + } + return it, nil +} + +func typedBuildAppendPromptTool[M adk.MessageType](prompt string, t tool.BaseTool) adk.TypedChatModelAgentMiddleware[M] { + return &typedAppendPromptTool[M]{ + TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[M]{}, + t: t, + prompt: prompt, + } +} + +type typedAppendPromptTool[M adk.MessageType] struct { + *adk.TypedBaseChatModelAgentMiddleware[M] + t tool.BaseTool + prompt string +} + +func (w *typedAppendPromptTool[M]) BeforeAgent(ctx context.Context, runCtx *adk.ChatModelAgentContext) (context.Context, *adk.ChatModelAgentContext, error) { + nRunCtx := *runCtx + nRunCtx.Instruction += w.prompt + if w.t != nil { + nRunCtx.Tools = append(nRunCtx.Tools, w.t) + } + return ctx, &nRunCtx, nil +} diff --git a/adk/prebuilt/integration_test.go b/adk/prebuilt/integration_test.go new file mode 100644 index 0000000..0ae9181 --- /dev/null +++ b/adk/prebuilt/integration_test.go @@ -0,0 +1,455 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prebuilt + +import ( + "context" + "fmt" + "strings" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/adk/prebuilt/planexecute" + "github.com/cloudwego/eino/adk/prebuilt/supervisor" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type approvalInfo struct { + ToolName string + ArgumentsInJSON string + ToolCallID string +} + +func (ai *approvalInfo) String() string { + return fmt.Sprintf("tool '%s' interrupted with arguments '%s', waiting for approval", + ai.ToolName, ai.ArgumentsInJSON) +} + +type approvalResult struct { + Approved bool + DisapproveReason *string +} + +func init() { + schema.Register[*approvalInfo]() + schema.Register[*approvalResult]() +} + +type approvableTool struct { + name string + t *testing.T +} + +func (m *approvableTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: m.name, + Desc: "A tool that requires approval before execution", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "action": {Type: schema.String, Desc: "The action to perform"}, + }), + }, nil +} + +func (m *approvableTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + wasInterrupted, _, storedArguments := tool.GetInterruptState[string](ctx) + if !wasInterrupted { + return "", tool.StatefulInterrupt(ctx, &approvalInfo{ + ToolName: m.name, + ArgumentsInJSON: argumentsInJSON, + ToolCallID: compose.GetToolCallID(ctx), + }, argumentsInJSON) + } + + isResumeTarget, hasData, data := tool.GetResumeContext[*approvalResult](ctx) + if !isResumeTarget { + return "", tool.StatefulInterrupt(ctx, &approvalInfo{ + ToolName: m.name, + ArgumentsInJSON: storedArguments, + ToolCallID: compose.GetToolCallID(ctx), + }, storedArguments) + } + + if !hasData { + return "", fmt.Errorf("tool '%s' resumed with no data", m.name) + } + + if data.Approved { + return fmt.Sprintf("Tool '%s' executed successfully with args: %s", m.name, storedArguments), nil + } + + if data.DisapproveReason != nil { + return fmt.Sprintf("Tool '%s' disapproved, reason: %s", m.name, *data.DisapproveReason), nil + } + + return fmt.Sprintf("Tool '%s' disapproved", m.name), nil +} + +type integrationCheckpointStore struct { + data map[string][]byte +} + +func newIntegrationCheckpointStore() *integrationCheckpointStore { + return &integrationCheckpointStore{data: make(map[string][]byte)} +} + +func (s *integrationCheckpointStore) Set(_ context.Context, key string, value []byte) error { + s.data[key] = value + return nil +} + +func (s *integrationCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.data[key] + return v, ok, nil +} + +type defaultPlan struct { + Steps []string `json:"steps"` +} + +func (p *defaultPlan) FirstStep() string { + if len(p.Steps) == 0 { + return "" + } + return p.Steps[0] +} + +func (p *defaultPlan) MarshalJSON() ([]byte, error) { + type planTyp defaultPlan + return sonic.Marshal((*planTyp)(p)) +} + +func (p *defaultPlan) UnmarshalJSON(bytes []byte) error { + type planTyp defaultPlan + return sonic.Unmarshal(bytes, (*planTyp)(p)) +} + +type namedAgent struct { + adk.ResumableAgent + name string + description string +} + +func (n *namedAgent) Name(_ context.Context) string { + return n.name +} + +func (n *namedAgent) Description(_ context.Context) string { + return n.description +} + +func formatRunPath(runPath []adk.RunStep) string { + if len(runPath) == 0 { + return "[]" + } + var parts []string + for _, step := range runPath { + parts = append(parts, step.String()) + } + return "[" + strings.Join(parts, " -> ") + "]" +} + +func formatAgentEventIntegration(event *adk.AgentEvent) string { + var sb strings.Builder + sb.WriteString(fmt.Sprintf("{AgentName: %q, RunPath: %s", event.AgentName, formatRunPath(event.RunPath))) + if event.Output != nil { + if event.Output.MessageOutput != nil && event.Output.MessageOutput.Message != nil { + msg := event.Output.MessageOutput.Message + sb.WriteString(fmt.Sprintf(", Output.Message: {Role: %q, Content: %q}", msg.Role, msg.Content)) + } + } + if event.Action != nil { + if event.Action.Interrupted != nil { + sb.WriteString(fmt.Sprintf(", Action.Interrupted: {%d contexts}", len(event.Action.Interrupted.InterruptContexts))) + } + if event.Action.BreakLoop != nil { + sb.WriteString(fmt.Sprintf(", Action.BreakLoop: {From: %q, Done: %v}", event.Action.BreakLoop.From, event.Action.BreakLoop.Done)) + } + if event.Action.TransferToAgent != nil { + sb.WriteString(fmt.Sprintf(", Action.TransferToAgent: {Dest: %q}", event.Action.TransferToAgent.DestAgentName)) + } + } + if event.Err != nil { + sb.WriteString(fmt.Sprintf(", Err: %v", event.Err)) + } + sb.WriteString("}") + return sb.String() +} + +func TestSupervisorWithPlanExecuteInterruptResume(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockPlannerModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockExecutorModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockReplannerModel := mockModel.NewMockToolCallingChatModel(ctrl) + + budgetTool := &approvableTool{name: "allocate_budget", t: t} + + plan := &defaultPlan{Steps: []string{"Allocate budget for the project", "Complete task"}} + userInput := []adk.Message{schema.UserMessage("Set up a new project with budget allocation")} + + plannerModelWithTools := mockModel.NewMockToolCallingChatModel(ctrl) + mockPlannerModel.EXPECT().WithTools(gomock.Any()).Return(plannerModelWithTools, nil).AnyTimes() + + planJSON, _ := sonic.MarshalString(plan) + plannerResponse := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "plan_call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "plan", + Arguments: planJSON, + }, + }, + }) + plannerModelWithTools.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + go func() { + defer sw.Close() + sw.Send(plannerResponse, nil) + }() + return sr, nil + }, + ).Times(1) + + mockExecutorModel.EXPECT().WithTools(gomock.Any()).Return(mockExecutorModel, nil).AnyTimes() + + toolCallMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_budget_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "allocate_budget", + Arguments: `{"action": "allocate $50000 for project"}`, + }, + }, + }) + mockExecutorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(toolCallMsg, nil).Times(1) + + completionMsg := schema.AssistantMessage("Budget allocated successfully", nil) + mockExecutorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(completionMsg, nil).AnyTimes() + + replannerModelWithTools := mockModel.NewMockToolCallingChatModel(ctrl) + mockReplannerModel.EXPECT().WithTools(gomock.Any()).Return(replannerModelWithTools, nil).AnyTimes() + + respondResponse := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "respond_call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "respond", + Arguments: `{"response":"Project setup completed with budget allocation"}`, + }, + }, + }) + replannerModelWithTools.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input []*schema.Message, opts ...interface{}) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + go func() { + defer sw.Close() + sw.Send(respondResponse, nil) + }() + return sr, nil + }, + ).AnyTimes() + + plannerAgent, err := planexecute.NewPlanner(ctx, &planexecute.PlannerConfig{ + ToolCallingChatModel: mockPlannerModel, + }) + assert.NoError(t, err) + + executorAgent, err := planexecute.NewExecutor(ctx, &planexecute.ExecutorConfig{ + Model: mockExecutorModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{budgetTool}, + }, + }, + }) + assert.NoError(t, err) + + replannerAgent, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{ + ChatModel: mockReplannerModel, + }) + assert.NoError(t, err) + + planExecuteAgent, err := planexecute.New(ctx, &planexecute.Config{ + Planner: plannerAgent, + Executor: executorAgent, + Replanner: replannerAgent, + MaxIterations: 10, + }) + assert.NoError(t, err) + + projectAgent := &namedAgent{ + ResumableAgent: planExecuteAgent, + name: "project_execution_agent", + description: "the agent responsible for complex project execution tasks", + } + + var pa adk.Agent = projectAgent + + _, ok := pa.(adk.ResumableAgent) + assert.True(t, ok) + + mockSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockSupervisorModel, nil).AnyTimes() + + transferToProjectMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "transfer_call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "transfer_to_agent", + Arguments: `{"agent_name":"project_execution_agent"}`, + }, + }, + }) + mockSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(transferToProjectMsg, nil).Times(1) + + finalSupervisorMsg := schema.AssistantMessage("Project setup completed successfully with budget allocation approved.", nil) + mockSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(finalSupervisorMsg, nil).AnyTimes() + + supervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "project_manager", + Description: "the supervisor agent responsible for coordinating project management tasks", + Instruction: "You are a project manager supervisor. Delegate complex project tasks to project_execution_agent.", + Model: mockSupervisorModel, + Exit: &adk.ExitTool{}, + }) + assert.NoError(t, err) + + supervisorAgent, err := supervisor.New(ctx, &supervisor.Config{ + Supervisor: supervisorChatAgent, + SubAgents: []adk.Agent{projectAgent}, + }) + assert.NoError(t, err) + + store := newIntegrationCheckpointStore() + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: supervisorAgent, + CheckPointStore: store, + }) + + t.Log("========================================") + t.Log("Starting Supervisor + PlanExecute Integration Test") + t.Log("========================================") + + checkpointID := "test-supervisor-plan_execute-1" + iter := runner.Run(ctx, userInput, adk.WithCheckPointID(checkpointID)) + + var interruptEvent *adk.AgentEvent + eventCount := 0 + for { + event, ok := iter.Next() + if !ok { + break + } + eventCount++ + t.Logf("Event %d: %s", eventCount, formatAgentEventIntegration(event)) + + if event.Err != nil { + t.Logf("Event has error: %v", event.Err) + } + + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + t.Log("========================================") + t.Log("INTERRUPT DETECTED - Deep interrupt from tool within executor") + t.Log("========================================") + break + } + } + + if interruptEvent == nil { + t.Fatal("Expected an interrupt event from the approvable tool, but none was received") + } + + assert.NotNil(t, interruptEvent.Action.Interrupted, "Should have interrupt info") + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, "Should have interrupt contexts") + + t.Logf("Interrupt event received with %d contexts", len(interruptEvent.Action.Interrupted.InterruptContexts)) + for i, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + t.Logf("Interrupt context %d: ID=%s, Info=%v, IsRootCause=%v", i, ctx.ID, ctx.Info, ctx.IsRootCause) + } + + var toolInterruptID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + toolInterruptID = intCtx.ID + break + } + } + assert.NotEmpty(t, toolInterruptID, "Should have a root cause interrupt ID") + + t.Log("========================================") + t.Logf("Resuming with approval for interrupt ID: %s", toolInterruptID) + t.Log("========================================") + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{ + Targets: map[string]any{ + toolInterruptID: &approvalResult{Approved: true}, + }, + }) + assert.NoError(t, err, "Resume should not error") + assert.NotNil(t, resumeIter, "Resume iterator should not be nil") + + var resumeEvents []*adk.AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + resumeEvents = append(resumeEvents, event) + } + + assert.NotEmpty(t, resumeEvents, "Should have resume events after approval") + + for _, event := range resumeEvents { + assert.NoError(t, event.Err, "Resume event should not have error") + } + + var hasToolResponse, hasBreakLoop bool + for _, event := range resumeEvents { + if event.Output != nil && event.Output.MessageOutput != nil { + msg := event.Output.MessageOutput.Message + if msg != nil && msg.Role == "tool" && strings.Contains(msg.Content, "executed successfully") { + hasToolResponse = true + } + } + if event.Action != nil && event.Action.BreakLoop != nil && event.Action.BreakLoop.Done { + hasBreakLoop = true + } + } + + assert.True(t, hasToolResponse, "Should have tool response indicating successful execution") + assert.True(t, hasBreakLoop, "Should have break loop action indicating task completion") +} diff --git a/adk/prebuilt/planexecute/plan_execute.go b/adk/prebuilt/planexecute/plan_execute.go new file mode 100644 index 0000000..2472a8b --- /dev/null +++ b/adk/prebuilt/planexecute/plan_execute.go @@ -0,0 +1,880 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package planexecute implements a plan–execute–replan style agent. +package planexecute + +import ( + "context" + "encoding/json" + "fmt" + "runtime/debug" + "strings" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*defaultPlan]("_eino_adk_plan_execute_default_plan") + schema.RegisterName[ExecutedStep]("_eino_adk_plan_execute_executed_step") + schema.RegisterName[[]ExecutedStep]("_eino_adk_plan_execute_executed_steps") +} + +// Plan represents an execution plan with a sequence of actionable steps. +// It supports JSON serialization and deserialization while providing access to the first step. +type Plan interface { + // FirstStep returns the first step to be executed in the plan. + FirstStep() string + + // Marshaler serializes the Plan into JSON. + // The resulting JSON can be used in prompt templates. + json.Marshaler + // Unmarshaler deserializes JSON content into the Plan. + // This processes output from structured chat models or tool calls into the Plan structure. + json.Unmarshaler +} + +// NewPlan is a function type that creates a new Plan instance. +type NewPlan func(ctx context.Context) Plan + +// defaultPlan is the default implementation of the Plan interface. +// +// JSON Schema: +// +// { +// "type": "object", +// "properties": { +// "steps": { +// "type": "array", +// "items": { +// "type": "string" +// }, +// "description": "Ordered list of actions to be taken. Each step should be clear, actionable, and arranged in a logical sequence." +// } +// }, +// "required": ["steps"] +// } +type defaultPlan struct { + // Steps contains the ordered list of actions to be taken. + // Each step should be clear, actionable, and arranged in a logical sequence. + Steps []string `json:"steps"` +} + +// FirstStep returns the first step in the plan or an empty string if no steps exist. +func (p *defaultPlan) FirstStep() string { + if len(p.Steps) == 0 { + return "" + } + return p.Steps[0] +} + +func (p *defaultPlan) MarshalJSON() ([]byte, error) { + type planTyp defaultPlan + return sonic.Marshal((*planTyp)(p)) +} + +func (p *defaultPlan) UnmarshalJSON(bytes []byte) error { + type planTyp defaultPlan + return sonic.Unmarshal(bytes, (*planTyp)(p)) +} + +// Response represents the final response to the user. +// This struct is used for JSON serialization/deserialization of the final response +// generated by the model. +type Response struct { + // Response is the complete response to provide to the user. + // This field is required. + Response string `json:"response"` +} + +var ( + // PlanToolInfo defines the schema for the Plan tool that can be used with ToolCallingChatModel. + // This schema instructs the model to generate a structured plan with ordered steps. + PlanToolInfo = schema.ToolInfo{ + Name: "plan", + Desc: "Plan with a list of steps to execute in order. Each step should be clear, actionable, and arranged in a logical sequence. The output will be used to guide the execution process.", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "steps": { + Type: schema.Array, + ElemInfo: &schema.ParameterInfo{Type: schema.String}, + Desc: "different steps to follow, should be in sorted order", + Required: true, + }, + }, + ), + } + + // RespondToolInfo defines the schema for the response tool that can be used with ToolCallingChatModel. + // This schema instructs the model to generate a direct response to the user. + RespondToolInfo = schema.ToolInfo{ + Name: "respond", + Desc: "Generate a direct response to the user. Use this tool when you have all the information needed to provide a final answer.", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "response": { + Type: schema.String, + Desc: "The complete response to provide to the user", + Required: true, + }, + }, + ), + } + + // PlannerPrompt is the prompt template for the planner. + // It provides context and guidance to the planner on how to generate the Plan. + PlannerPrompt = prompt.FromMessages(schema.FString, + schema.SystemMessage(`You are an expert planning agent. Given an objective, create a comprehensive step-by-step plan to achieve the objective. + +## YOUR TASK +Analyze the objective and generate a strategic plan that breaks down the goal into manageable, executable steps. + +## PLANNING REQUIREMENTS +Each step in your plan must be: +- **Specific and actionable**: Clear instructions that can be executed without ambiguity +- **Self-contained**: Include all necessary context, parameters, and requirements +- **Independently executable**: Can be performed by an agent without dependencies on other steps +- **Logically sequenced**: Arranged in optimal order for efficient execution +- **Objective-focused**: Directly contribute to achieving the main goal + +## PLANNING GUIDELINES +- Eliminate redundant or unnecessary steps +- Include relevant constraints, parameters, and success criteria for each step +- Ensure the final step produces a complete answer or deliverable +- Anticipate potential challenges and include mitigation strategies +- Structure steps to build upon each other logically +- Provide sufficient detail for successful execution + +## QUALITY CRITERIA +- Plan completeness: Does it address all aspects of the objective? +- Step clarity: Can each step be understood and executed independently? +- Logical flow: Do steps follow a sensible progression? +- Efficiency: Is this the most direct path to the objective? +- Adaptability: Can the plan handle unexpected results or changes?`), + schema.MessagesPlaceholder("input", false), + ) + + // ExecutorPrompt is the prompt template for the executor. + // It provides context and guidance to the executor on how to execute the Task. + ExecutorPrompt = prompt.FromMessages(schema.FString, + schema.SystemMessage(`You are a diligent and meticulous executor agent. Follow the given plan and execute your tasks carefully and thoroughly.`), + schema.UserMessage(`## OBJECTIVE +{input} +## Given the following plan: +{plan} +## COMPLETED STEPS & RESULTS +{executed_steps} +## Your task is to execute the first step, which is: +{step}`)) + + // ReplannerPrompt is the prompt template for the replanner. + // It provides context and guidance to the replanner on how to regenerate the Plan. + ReplannerPrompt = prompt.FromMessages(schema.FString, + schema.SystemMessage( + `You are going to review the progress toward an objective. Analyze the current state and determine the optimal next action. + +## YOUR TASK +Based on the progress above, you MUST choose exactly ONE action: + +### Option 1: COMPLETE (if objective is fully achieved) +Call '{respond_tool}' with: +- A comprehensive final answer +- Clear conclusion summarizing how the objective was met +- Key insights from the execution process + +### Option 2: CONTINUE (if more work is needed) +Call '{plan_tool}' with a revised plan that: +- Contains ONLY remaining steps (exclude completed ones) +- Incorporates lessons learned from executed steps +- Addresses any gaps or issues discovered +- Maintains logical step sequence + +## PLANNING REQUIREMENTS +Each step in your plan must be: +- **Specific and actionable**: Clear instructions that can be executed without ambiguity +- **Self-contained**: Include all necessary context, parameters, and requirements +- **Independently executable**: Can be performed by an agent without dependencies on other steps +- **Logically sequenced**: Arranged in optimal order for efficient execution +- **Objective-focused**: Directly contribute to achieving the main goal + +## PLANNING GUIDELINES +- Eliminate redundant or unnecessary steps +- Adapt strategy based on new information +- Include relevant constraints, parameters, and success criteria for each step + +## DECISION CRITERIA +- Has the original objective been completely satisfied? +- Are there any remaining requirements or sub-goals? +- Do the results suggest a need for strategy adjustment? +- What specific actions are still required?`), + schema.UserMessage(`## OBJECTIVE +{input} + +## ORIGINAL PLAN +{plan} + +## COMPLETED STEPS & RESULTS +{executed_steps}`), + ) +) + +const ( + // UserInputSessionKey is the session key for the user input. + UserInputSessionKey = "UserInput" + + // PlanSessionKey is the session key for the plan. + PlanSessionKey = "Plan" + + // ExecutedStepSessionKey is the session key for the execute result. + ExecutedStepSessionKey = "ExecutedStep" + + // ExecutedStepsSessionKey is the session key for the execute results. + ExecutedStepsSessionKey = "ExecutedSteps" +) + +// PlannerConfig provides configuration options for creating a planner agent. +// There are two ways to configure the planner to generate structured Plan output: +// 1. Use ChatModelWithFormattedOutput: A model pre-configured to output in the Plan format +// 2. Use ToolCallingChatModel + ToolInfo: A model that uses tool calling to generate +// the Plan structure +type PlannerConfig struct { + // ChatModelWithFormattedOutput is a model pre-configured to output in the Plan format. + // Create this by configuring a model to output structured data directly. + // See example: https://github.com/cloudwego/eino-ext/blob/main/components/model/openai/examples/structured/structured.go + ChatModelWithFormattedOutput model.BaseChatModel + + // ToolCallingChatModel is a model that supports tool calling capabilities. + // When provided with ToolInfo, it will use tool calling to generate the Plan structure. + ToolCallingChatModel model.ToolCallingChatModel + + // ToolInfo defines the schema for the Plan structure when using tool calling. + // Optional. If not provided, PlanToolInfo will be used as the default. + ToolInfo *schema.ToolInfo + + // GenInputFn is a function that generates the input messages for the planner. + // Optional. If not provided, defaultGenPlannerInputFn will be used. + GenInputFn GenPlannerModelInputFn + + // NewPlan creates a new Plan instance for JSON. + // The returned Plan will be used to unmarshal the model-generated JSON output. + // Optional. If not provided, defaultNewPlan will be used. + NewPlan NewPlan +} + +// GenPlannerModelInputFn is a function type that generates input messages for the planner. +type GenPlannerModelInputFn func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) + +func defaultNewPlan(ctx context.Context) Plan { + return &defaultPlan{} +} + +func defaultGenPlannerInputFn(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) { + msgs, err := PlannerPrompt.Format(ctx, map[string]any{ + "input": userInput, + }) + if err != nil { + return nil, err + } + return msgs, nil +} + +type planner struct { + toolCall bool + chatModel model.BaseChatModel + genInputFn GenPlannerModelInputFn + newPlan NewPlan +} + +func (p *planner) Name(_ context.Context) string { + return "planner" +} + +func (p *planner) Description(_ context.Context) string { + return "a planner agent" +} + +func argToContent(msg adk.Message) (adk.Message, error) { + if len(msg.ToolCalls) == 0 { + return nil, schema.ErrNoValue + } + + return schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil), nil +} + +func (p *planner) Run(ctx context.Context, input *adk.AgentInput, + _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + adk.AddSessionValue(ctx, UserInputSessionKey, input.Messages) + + go func() { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&adk.AgentEvent{Err: e}) + } + + generator.Close() + }() + + c := compose.NewChain[*adk.AgentInput, Plan](). + AppendLambda( + compose.InvokableLambda(func(ctx context.Context, input *adk.AgentInput) (output []adk.Message, err error) { + return p.genInputFn(ctx, input.Messages) + }), + ). + AppendChatModel(p.chatModel). + AppendLambda( + compose.CollectableLambda(func(ctx context.Context, sr *schema.StreamReader[adk.Message]) (adk.Message, error) { + if input.EnableStreaming { + ss := sr.Copy(2) + var sOutput *schema.StreamReader[*schema.Message] + if p.toolCall { + sOutput = schema.StreamReaderWithConvert(ss[0], argToContent) + } else { + sOutput = ss[0] + } + + generator.Send(adk.EventFromMessage(nil, sOutput, schema.Assistant, "")) + + return schema.ConcatMessageStream(ss[1]) + } + + msg, err := schema.ConcatMessageStream(sr) + if err != nil { + return nil, err + } + + var output adk.Message + if p.toolCall { + if len(msg.ToolCalls) == 0 { + return nil, fmt.Errorf("no tool call") + } + output = schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil) + } else { + output = msg + } + + generator.Send(adk.EventFromMessage(output, nil, schema.Assistant, "")) + + return msg, nil + }), + ). + AppendLambda( + compose.InvokableLambda(func(ctx context.Context, msg adk.Message) (plan Plan, err error) { + var planJSON string + if p.toolCall { + if len(msg.ToolCalls) == 0 { + return nil, fmt.Errorf("no tool call") + } + planJSON = msg.ToolCalls[0].Function.Arguments + } else { + planJSON = msg.Content + } + + plan = p.newPlan(ctx) + err = plan.UnmarshalJSON([]byte(planJSON)) + if err != nil { + return nil, fmt.Errorf("unmarshal plan error: %w", err) + } + + adk.AddSessionValue(ctx, PlanSessionKey, plan) + + return plan, nil + }), + ) + + var opts []compose.Option + if p.toolCall { + opts = append(opts, compose.WithChatModelOption(model.WithToolChoice(schema.ToolChoiceForced))) + } + + r, err := c.Compile(ctx, compose.WithGraphName(p.Name(ctx))) + if err != nil { // unexpected + generator.Send(&adk.AgentEvent{Err: err}) + return + } + + _, err = r.Stream(ctx, input, opts...) + if err != nil { + generator.Send(&adk.AgentEvent{Err: err}) + return + } + }() + + return iterator +} + +// NewPlanner creates a new planner agent based on the provided configuration. +// The planner agent uses either ChatModelWithFormattedOutput or ToolCallingChatModel+ToolInfo +// to generate structured Plan output. +// +// If ChatModelWithFormattedOutput is provided, it will be used directly. +// If ToolCallingChatModel is provided, it will be configured with ToolInfo (or PlanToolInfo by default) +// to generate structured Plan output. +func NewPlanner(_ context.Context, cfg *PlannerConfig) (adk.Agent, error) { + var chatModel model.BaseChatModel + var toolCall bool + if cfg.ChatModelWithFormattedOutput != nil { + chatModel = cfg.ChatModelWithFormattedOutput + } else { + toolCall = true + toolInfo := cfg.ToolInfo + if toolInfo == nil { + toolInfo = &PlanToolInfo + } + + var err error + chatModel, err = cfg.ToolCallingChatModel.WithTools([]*schema.ToolInfo{toolInfo}) + if err != nil { + return nil, err + } + } + + inputFn := cfg.GenInputFn + if inputFn == nil { + inputFn = defaultGenPlannerInputFn + } + + planParser := cfg.NewPlan + if planParser == nil { + planParser = defaultNewPlan + } + + return &planner{ + toolCall: toolCall, + chatModel: chatModel, + genInputFn: inputFn, + newPlan: planParser, + }, nil +} + +// ExecutionContext is the input information for the executor and the planner. +type ExecutionContext struct { + UserInput []adk.Message + Plan Plan + ExecutedSteps []ExecutedStep +} + +// GenModelInputFn is a function that generates the input messages for the executor and the planner. +type GenModelInputFn func(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) + +// ExecutorConfig provides configuration options for creating an executor agent. +type ExecutorConfig struct { + // Model is the chat model used by the executor. + // If the executor uses any tools, this model must support the model.WithTools call option, + // as that's how the executor configures the model with tool information. + Model model.BaseChatModel + + // ToolsConfig specifies the tools available to the executor. + ToolsConfig adk.ToolsConfig + + // MaxIterations defines the upper limit of ChatModel generation cycles. + // The agent will terminate with an error if this limit is exceeded. + // Optional. Defaults to 20. + MaxIterations int + + // GenInputFn generates the input messages for the Executor. + // Optional. If not provided, defaultGenExecutorInputFn will be used. + GenInputFn GenModelInputFn +} + +type ExecutedStep struct { + Step string + Result string +} + +// NewExecutor creates a new executor agent. +func NewExecutor(ctx context.Context, cfg *ExecutorConfig) (adk.Agent, error) { + + genInputFn := cfg.GenInputFn + if genInputFn == nil { + genInputFn = defaultGenExecutorInputFn + } + genInput := func(ctx context.Context, instruction string, _ *adk.AgentInput) ([]adk.Message, error) { + + plan, ok := adk.GetSessionValue(ctx, PlanSessionKey) + if !ok { + panic("impossible: plan not found") + } + plan_ := plan.(Plan) + + userInput, ok := adk.GetSessionValue(ctx, UserInputSessionKey) + if !ok { + panic("impossible: user input not found") + } + userInput_ := userInput.([]adk.Message) + + var executedSteps_ []ExecutedStep + executedStep, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey) + if ok { + executedSteps_ = executedStep.([]ExecutedStep) + } + + in := &ExecutionContext{ + UserInput: userInput_, + Plan: plan_, + ExecutedSteps: executedSteps_, + } + + msgs, err := genInputFn(ctx, in) + if err != nil { + return nil, err + } + + return msgs, nil + } + + agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "executor", + Description: "an executor agent", + Model: cfg.Model, + ToolsConfig: cfg.ToolsConfig, + GenModelInput: genInput, + MaxIterations: cfg.MaxIterations, + OutputKey: ExecutedStepSessionKey, + }) + if err != nil { + return nil, err + } + + return agent, nil +} + +func defaultGenExecutorInputFn(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) { + + planContent, err := in.Plan.MarshalJSON() + if err != nil { + return nil, err + } + + userMsgs, err := ExecutorPrompt.Format(ctx, map[string]any{ + "input": formatInput(in.UserInput), + "plan": string(planContent), + "executed_steps": formatExecutedSteps(in.ExecutedSteps), + "step": in.Plan.FirstStep(), + }) + if err != nil { + return nil, err + } + + return userMsgs, nil +} + +type replanner struct { + chatModel model.ToolCallingChatModel + planTool *schema.ToolInfo + respondTool *schema.ToolInfo + + genInputFn GenModelInputFn + newPlan NewPlan +} + +type ReplannerConfig struct { + // ChatModel is the model that supports tool calling capabilities. + // It will be configured with PlanTool and RespondTool to generate updated plans or responses. + ChatModel model.ToolCallingChatModel + + // PlanTool defines the schema for the Plan tool that can be used with ToolCallingChatModel. + // Optional. If not provided, the default PlanToolInfo will be used. + PlanTool *schema.ToolInfo + + // RespondTool defines the schema for the response tool that can be used with ToolCallingChatModel. + // Optional. If not provided, the default RespondToolInfo will be used. + RespondTool *schema.ToolInfo + + // GenInputFn generates the input messages for the Replanner. + // Optional. If not provided, buildGenReplannerInputFn will be used. + GenInputFn GenModelInputFn + + // NewPlan creates a new Plan instance. + // The returned Plan will be used to unmarshal the model-generated JSON output from PlanTool. + // Optional. If not provided, defaultNewPlan will be used. + NewPlan NewPlan +} + +// formatInput formats the input messages into a string. +func formatInput(input []adk.Message) string { + var sb strings.Builder + for _, msg := range input { + sb.WriteString(msg.Content) + sb.WriteString("\n") + } + + return sb.String() +} + +func formatExecutedSteps(results []ExecutedStep) string { + var sb strings.Builder + for _, result := range results { + sb.WriteString(fmt.Sprintf("Step: %s\nResult: %s\n\n", result.Step, result.Result)) + } + + return sb.String() +} + +func (r *replanner) Name(_ context.Context) string { + return "replanner" +} + +func (r *replanner) Description(_ context.Context) string { + return "a replanner agent" +} + +func (r *replanner) genInput(ctx context.Context) ([]adk.Message, error) { + + executedStep, ok := adk.GetSessionValue(ctx, ExecutedStepSessionKey) + if !ok { + panic("impossible: execute result not found") + } + executedStep_ := executedStep.(string) + + plan, ok := adk.GetSessionValue(ctx, PlanSessionKey) + if !ok { + panic("impossible: plan not found") + } + plan_ := plan.(Plan) + step := plan_.FirstStep() + + var executedSteps_ []ExecutedStep + executedSteps, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey) + if ok { + executedSteps_ = executedSteps.([]ExecutedStep) + } + + executedSteps_ = append(executedSteps_, ExecutedStep{ + Step: step, + Result: executedStep_, + }) + adk.AddSessionValue(ctx, ExecutedStepsSessionKey, executedSteps_) + + userInput, ok := adk.GetSessionValue(ctx, UserInputSessionKey) + if !ok { + panic("impossible: user input not found") + } + userInput_ := userInput.([]adk.Message) + + in := &ExecutionContext{ + UserInput: userInput_, + Plan: plan_, + ExecutedSteps: executedSteps_, + } + genInputFn := r.genInputFn + if genInputFn == nil { + genInputFn = buildGenReplannerInputFn(r.planTool.Name, r.respondTool.Name) + } + msgs, err := genInputFn(ctx, in) + if err != nil { + return nil, err + } + + return msgs, nil +} + +func (r *replanner) Run(ctx context.Context, input *adk.AgentInput, _ ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + go func() { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&adk.AgentEvent{Err: e}) + } + + generator.Close() + }() + + callOpt := model.WithToolChoice(schema.ToolChoiceForced) + + c := compose.NewChain[struct{}, any](). + AppendLambda( + compose.InvokableLambda(func(ctx context.Context, input struct{}) (output []adk.Message, err error) { + return r.genInput(ctx) + }), + ). + AppendChatModel(r.chatModel). + AppendLambda( + compose.CollectableLambda(func(ctx context.Context, sr *schema.StreamReader[adk.Message]) (adk.Message, error) { + if input.EnableStreaming { + ss := sr.Copy(2) + sOutput := schema.StreamReaderWithConvert(ss[0], argToContent) + generator.Send(adk.EventFromMessage(nil, sOutput, schema.Assistant, "")) + return schema.ConcatMessageStream(ss[1]) + } + + msg, err := schema.ConcatMessageStream(sr) + if err != nil { + return nil, err + } + if len(msg.ToolCalls) > 0 { + output := schema.AssistantMessage(msg.ToolCalls[0].Function.Arguments, nil) + generator.Send(adk.EventFromMessage(output, nil, schema.Assistant, "")) + } + return msg, nil + }), + ). + AppendLambda( + compose.InvokableLambda(func(ctx context.Context, msg adk.Message) (msgOrPlan any, err error) { + if len(msg.ToolCalls) == 0 { + return nil, fmt.Errorf("no tool call") + } + + // exit + if msg.ToolCalls[0].Function.Name == r.respondTool.Name { + action := adk.NewBreakLoopAction(r.Name(ctx)) + generator.Send(&adk.AgentEvent{Action: action}) + return msg, nil + } + + // replan + if msg.ToolCalls[0].Function.Name != r.planTool.Name { + return nil, fmt.Errorf("unexpected tool call: %s", msg.ToolCalls[0].Function.Name) + } + + plan := r.newPlan(ctx) + if err = plan.UnmarshalJSON([]byte(msg.ToolCalls[0].Function.Arguments)); err != nil { + return nil, fmt.Errorf("unmarshal plan error: %w", err) + } + + adk.AddSessionValue(ctx, PlanSessionKey, plan) + + return plan, nil + }), + ) + + runnable, err := c.Compile(ctx, compose.WithGraphName(r.Name(ctx))) + if err != nil { + generator.Send(&adk.AgentEvent{Err: err}) + return + } + + _, err = runnable.Stream(ctx, struct{}{}, compose.WithChatModelOption(callOpt)) + if err != nil { + generator.Send(&adk.AgentEvent{Err: err}) + return + } + }() + + return iterator +} + +func buildGenReplannerInputFn(planToolName, respondToolName string) GenModelInputFn { + return func(ctx context.Context, in *ExecutionContext) ([]adk.Message, error) { + planContent, err := in.Plan.MarshalJSON() + if err != nil { + return nil, err + } + msgs, err := ReplannerPrompt.Format(ctx, map[string]any{ + "plan": string(planContent), + "input": formatInput(in.UserInput), + "executed_steps": formatExecutedSteps(in.ExecutedSteps), + "plan_tool": planToolName, + "respond_tool": respondToolName, + }) + if err != nil { + return nil, err + } + + return msgs, nil + } +} + +// NewReplanner creates a plan-execute-replan agent wired with plan and respond tools. +// It configures the provided ToolCallingChatModel with the tools and returns an Agent. +func NewReplanner(_ context.Context, cfg *ReplannerConfig) (adk.Agent, error) { + planTool := cfg.PlanTool + if planTool == nil { + planTool = &PlanToolInfo + } + + respondTool := cfg.RespondTool + if respondTool == nil { + respondTool = &RespondToolInfo + } + + chatModel, err := cfg.ChatModel.WithTools([]*schema.ToolInfo{planTool, respondTool}) + if err != nil { + return nil, err + } + + planParser := cfg.NewPlan + if planParser == nil { + planParser = defaultNewPlan + } + + return &replanner{ + chatModel: chatModel, + planTool: planTool, + respondTool: respondTool, + genInputFn: cfg.GenInputFn, + newPlan: planParser, + }, nil +} + +// Config provides configuration options for creating a plan-execute-replan agent. +type Config struct { + // Planner specifies the agent that generates the plan. + // You can use provided NewPlanner to create a planner agent. + Planner adk.Agent + + // Executor specifies the agent that executes the plan generated by planner or replanner. + // You can use provided NewExecutor to create an executor agent. + Executor adk.Agent + + // Replanner specifies the agent that replans the plan. + // You can use provided NewReplanner to create a replanner agent. + Replanner adk.Agent + + // MaxIterations defines the maximum number of loops for 'execute-replan'. + // Optional. If not provided, 10 will be used as the default. + MaxIterations int +} + +// New creates a new plan-execute-replan agent with the given configuration. +// The plan-execute-replan pattern works in three phases: +// 1. Planning: Generate a structured plan with clear, actionable steps +// 2. Execution: Execute the first step of the plan +// 3. Replanning: Evaluate progress and either complete the task or revise the plan +// This approach enables complex problem-solving through iterative refinement. +func New(ctx context.Context, cfg *Config) (adk.ResumableAgent, error) { + maxIterations := cfg.MaxIterations + if maxIterations <= 0 { + maxIterations = 10 + } + loop, err := adk.NewLoopAgent(ctx, &adk.LoopAgentConfig{ + Name: "execute_replan", + SubAgents: []adk.Agent{cfg.Executor, cfg.Replanner}, + MaxIterations: maxIterations, + }) + if err != nil { + return nil, err + } + + return adk.NewSequentialAgent(ctx, &adk.SequentialAgentConfig{ + Name: "plan_execute_replan", + SubAgents: []adk.Agent{cfg.Planner, loop}, + }) +} diff --git a/adk/prebuilt/planexecute/plan_execute_test.go b/adk/prebuilt/planexecute/plan_execute_test.go new file mode 100644 index 0000000..25c64ea --- /dev/null +++ b/adk/prebuilt/planexecute/plan_execute_test.go @@ -0,0 +1,1201 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package planexecute + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockAdk "github.com/cloudwego/eino/internal/mock/adk" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +// TestNewPlanner tests the NewPlanner function with ChatModelWithFormattedOutput +func TestNewPlannerWithFormattedOutput(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock chat model + mockChatModel := mockModel.NewMockBaseChatModel(ctrl) + + // Create the PlannerConfig + conf := &PlannerConfig{ + ChatModelWithFormattedOutput: mockChatModel, + } + + // Create the planner + p, err := NewPlanner(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, p) + + // Verify the planner's name and description + assert.Equal(t, "planner", p.Name(ctx)) + assert.Equal(t, "a planner agent", p.Description(ctx)) +} + +// TestNewPlannerWithToolCalling tests the NewPlanner function with ToolCallingChatModel +func TestNewPlannerWithToolCalling(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1) + + // Create the PlannerConfig + conf := &PlannerConfig{ + ToolCallingChatModel: mockToolCallingModel, + // Use default instruction and tool info + } + + // Create the planner + p, err := NewPlanner(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, p) + + // Verify the planner's name and description + assert.Equal(t, "planner", p.Name(ctx)) + assert.Equal(t, "a planner agent", p.Description(ctx)) +} + +// TestPlannerRunWithFormattedOutput tests the Run method of a planner created with ChatModelWithFormattedOutput +func TestPlannerRunWithFormattedOutput(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock chat model + mockChatModel := mockModel.NewMockBaseChatModel(ctrl) + + // Create a plan response + planJSON := `{"steps":["Step 1", "Step 2", "Step 3"]}` + planMsg := schema.AssistantMessage(planJSON, nil) + sr, sw := schema.Pipe[*schema.Message](1) + sw.Send(planMsg, nil) + sw.Close() + + // Mock the Generate method + mockChatModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + + // Create the PlannerConfig + conf := &PlannerConfig{ + ChatModelWithFormattedOutput: mockChatModel, + } + + // Create the planner + p, err := NewPlanner(ctx, conf) + assert.NoError(t, err) + + // Run the planner + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: p}) + iterator := runner.Run(ctx, []adk.Message{schema.UserMessage("Plan this task")}) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + msg, _, err := adk.GetMessage(event) + assert.NoError(t, err) + assert.Equal(t, planMsg.Content, msg.Content) + + _, ok = iterator.Next() + assert.False(t, ok) + + plan := defaultNewPlan(ctx) + err = plan.UnmarshalJSON([]byte(msg.Content)) + assert.NoError(t, err) + plan_ := plan.(*defaultPlan) + assert.Equal(t, 3, len(plan_.Steps)) + assert.Equal(t, "Step 1", plan_.Steps[0]) + assert.Equal(t, "Step 2", plan_.Steps[1]) + assert.Equal(t, "Step 3", plan_.Steps[2]) +} + +// TestPlannerRunWithToolCalling tests the Run method of a planner created with ToolCallingChatModel +func TestPlannerRunWithToolCalling(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create a tool call response with a plan + planArgs := `{"steps":["Step 1", "Step 2", "Step 3"]}` + toolCall := schema.ToolCall{ + ID: "tool_call_id", + Type: "function", + Function: schema.FunctionCall{ + Name: "plan", // This should match PlanToolInfo.Name + Arguments: planArgs, + }, + } + + toolCallMsg := schema.AssistantMessage("", nil) + toolCallMsg.ToolCalls = []schema.ToolCall{toolCall} + sr, sw := schema.Pipe[*schema.Message](1) + sw.Send(toolCallMsg, nil) + sw.Close() + + // Mock the WithTools method to return a model that will be used for Generate + mockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1) + + // Mock the Generate method to return the tool call message + mockToolCallingModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + + // Create the PlannerConfig with ToolCallingChatModel + conf := &PlannerConfig{ + ToolCallingChatModel: mockToolCallingModel, + // Use default instruction and tool info + } + + // Create the planner + p, err := NewPlanner(ctx, conf) + assert.NoError(t, err) + + // Run the planner + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: p}) + iterator := runner.Run(ctx, []adk.Message{schema.UserMessage("no input")}) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + msg, _, err := adk.GetMessage(event) + assert.NoError(t, err) + assert.Equal(t, planArgs, msg.Content) + + _, ok = iterator.Next() + assert.False(t, ok) + + plan := defaultNewPlan(ctx) + err = plan.UnmarshalJSON([]byte(msg.Content)) + assert.NoError(t, err) + plan_ := plan.(*defaultPlan) + assert.NoError(t, err) + assert.Equal(t, 3, len(plan_.Steps)) + assert.Equal(t, "Step 1", plan_.Steps[0]) + assert.Equal(t, "Step 2", plan_.Steps[1]) + assert.Equal(t, "Step 3", plan_.Steps[2]) +} + +// TestNewExecutor tests the NewExecutor function +func TestNewExecutor(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create the ExecutorConfig + conf := &ExecutorConfig{ + Model: mockToolCallingModel, + MaxIterations: 3, + } + + // Create the executor + executor, err := NewExecutor(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, executor) + + // Verify the executor's name and description + assert.Equal(t, "executor", executor.Name(ctx)) + assert.Equal(t, "an executor agent", executor.Description(ctx)) +} + +// TestExecutorRun tests the Run method of the executor +func TestExecutorRun(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + + // Store a plan in the session + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + adk.AddSessionValue(ctx, PlanSessionKey, plan) + + // Set up expectations for the mock model + // The model should return the last user message as its response + mockToolCallingModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, messages []*schema.Message, opts ...model.Option) (*schema.Message, error) { + // Find the last user message + var lastUserMessage string + for _, msg := range messages { + if msg.Role == schema.User { + lastUserMessage = msg.Content + } + } + // Return the last user message as the model's response + return schema.AssistantMessage(lastUserMessage, nil), nil + }).Times(1) + + // Create the ExecutorConfig + conf := &ExecutorConfig{ + Model: mockToolCallingModel, + MaxIterations: 3, + } + + // Create the executor + executor, err := NewExecutor(ctx, conf) + assert.NoError(t, err) + + // Run the executor + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: executor}) + iterator := runner.Run(ctx, []adk.Message{schema.UserMessage("no input")}, + adk.WithSessionValues(map[string]any{ + PlanSessionKey: plan, + UserInputSessionKey: []adk.Message{schema.UserMessage("no input")}, + }), + ) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + msg, _, err := adk.GetMessage(event) + assert.NoError(t, err) + t.Logf("executor model input msg:\n %s\n", msg.Content) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +// TestNewReplanner tests the NewReplanner function +func TestNewReplanner(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + // Mock the WithTools method + mockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1) + + // Create plan and respond tools + planTool := &schema.ToolInfo{ + Name: "Plan", + Desc: "Plan tool", + } + + respondTool := &schema.ToolInfo{ + Name: "Respond", + Desc: "Respond tool", + } + + // Create the ReplannerConfig + conf := &ReplannerConfig{ + ChatModel: mockToolCallingModel, + PlanTool: planTool, + RespondTool: respondTool, + } + + // Create the replanner + rp, err := NewReplanner(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, rp) + + // Verify the replanner's name and description + assert.Equal(t, "replanner", rp.Name(ctx)) + assert.Equal(t, "a replanner agent", rp.Description(ctx)) +} + +// TestReplannerRunWithPlan tests the Replanner's ability to use the plan_tool +func TestReplannerRunWithPlan(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create plan and respond tools + planTool := &schema.ToolInfo{ + Name: "Plan", + Desc: "Plan tool", + } + + respondTool := &schema.ToolInfo{ + Name: "Respond", + Desc: "Respond tool", + } + + // Create a tool call response for the Plan tool + planArgs := `{"steps":["Updated Step 1", "Updated Step 2"]}` + toolCall := schema.ToolCall{ + ID: "tool_call_id", + Type: "function", + Function: schema.FunctionCall{ + Name: planTool.Name, + Arguments: planArgs, + }, + } + + toolCallMsg := schema.AssistantMessage("", nil) + toolCallMsg.ToolCalls = []schema.ToolCall{toolCall} + sr, sw := schema.Pipe[*schema.Message](1) + sw.Send(toolCallMsg, nil) + sw.Close() + + // Mock the Generate method + mockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1) + mockToolCallingModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + + // Create the ReplannerConfig + conf := &ReplannerConfig{ + ChatModel: mockToolCallingModel, + PlanTool: planTool, + RespondTool: respondTool, + } + + // Create the replanner + rp, err := NewReplanner(ctx, conf) + assert.NoError(t, err) + + // Store necessary values in the session + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + + rp, err = agentOutputSessionKVs(ctx, rp) + assert.NoError(t, err) + + // Run the replanner + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: rp}) + iterator := runner.Run(ctx, []adk.Message{schema.UserMessage("no input")}, + adk.WithSessionValues(map[string]any{ + PlanSessionKey: plan, + ExecutedStepSessionKey: "Execution result", + UserInputSessionKey: []adk.Message{schema.UserMessage("User input")}, + }), + ) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + + event, ok = iterator.Next() + assert.True(t, ok) + kvs := event.Output.CustomizedOutput.(map[string]any) + assert.Greater(t, len(kvs), 0) + + // Verify the updated plan was stored in the session + planValue, ok := kvs[PlanSessionKey] + assert.True(t, ok) + updatedPlan, ok := planValue.(*defaultPlan) + assert.True(t, ok) + assert.Equal(t, 2, len(updatedPlan.Steps)) + assert.Equal(t, "Updated Step 1", updatedPlan.Steps[0]) + assert.Equal(t, "Updated Step 2", updatedPlan.Steps[1]) + + // Verify the execute results were updated + executeResultsValue, ok := kvs[ExecutedStepsSessionKey] + assert.True(t, ok) + executeResults, ok := executeResultsValue.([]ExecutedStep) + assert.True(t, ok) + assert.Equal(t, 1, len(executeResults)) + assert.Equal(t, "Step 1", executeResults[0].Step) + assert.Equal(t, "Execution result", executeResults[0].Result) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +// TestReplannerRunWithRespond tests the Replanner's ability to use the respond_tool +func TestReplannerRunWithRespond(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create a mock tool calling chat model + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + + // Create plan and respond tools + planTool := &schema.ToolInfo{ + Name: "Plan", + Desc: "Plan tool", + } + + respondTool := &schema.ToolInfo{ + Name: "Respond", + Desc: "Respond tool", + } + + // Create a tool call response for the Respond tool + responseArgs := `{"response":"This is the final response to the user"}` + toolCall := schema.ToolCall{ + ID: "tool_call_id", + Type: "function", + Function: schema.FunctionCall{ + Name: respondTool.Name, + Arguments: responseArgs, + }, + } + + toolCallMsg := schema.AssistantMessage("", nil) + toolCallMsg.ToolCalls = []schema.ToolCall{toolCall} + sr, sw := schema.Pipe[*schema.Message](1) + sw.Send(toolCallMsg, nil) + sw.Close() + + // Mock the Generate method + mockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).Times(1) + mockToolCallingModel.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + + // Create the ReplannerConfig + conf := &ReplannerConfig{ + ChatModel: mockToolCallingModel, + PlanTool: planTool, + RespondTool: respondTool, + } + + // Create the replanner + rp, err := NewReplanner(ctx, conf) + assert.NoError(t, err) + + // Store necessary values in the session + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + + // Run the replanner + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: rp}) + iterator := runner.Run(ctx, []adk.Message{schema.UserMessage("no input")}, + adk.WithSessionValues(map[string]any{ + PlanSessionKey: plan, + ExecutedStepSessionKey: "Execution result", + UserInputSessionKey: []adk.Message{schema.UserMessage("User input")}, + }), + ) + + // Get the event from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Nil(t, event.Err) + msg, _, err := adk.GetMessage(event) + assert.NoError(t, err) + assert.Equal(t, responseArgs, msg.Content) + + // Verify that an exit action was generated + event, ok = iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event.Action) + assert.NotNil(t, event.Action.BreakLoop) + assert.False(t, event.Action.BreakLoop.Done) + + _, ok = iterator.Next() + assert.False(t, ok) +} + +// TestNewPlanExecuteAgent tests the New function +func TestNewPlanExecuteAgent(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create mock agents + mockPlanner := mockAdk.NewMockAgent(ctrl) + mockExecutor := mockAdk.NewMockAgent(ctrl) + mockReplanner := mockAdk.NewMockAgent(ctrl) + + // Set up expectations for the mock agents + mockPlanner.EXPECT().Name(gomock.Any()).Return("planner").AnyTimes() + mockPlanner.EXPECT().Description(gomock.Any()).Return("a planner agent").AnyTimes() + + mockExecutor.EXPECT().Name(gomock.Any()).Return("executor").AnyTimes() + mockExecutor.EXPECT().Description(gomock.Any()).Return("an executor agent").AnyTimes() + + mockReplanner.EXPECT().Name(gomock.Any()).Return("replanner").AnyTimes() + mockReplanner.EXPECT().Description(gomock.Any()).Return("a replanner agent").AnyTimes() + + conf := &Config{ + Planner: mockPlanner, + Executor: mockExecutor, + Replanner: mockReplanner, + } + + // Create the plan execute agent + agent, err := New(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, agent) +} + +func TestPlanExecuteAgentWithReplan(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create mock agents + mockPlanner := mockAdk.NewMockAgent(ctrl) + mockExecutor := mockAdk.NewMockAgent(ctrl) + mockReplanner := mockAdk.NewMockAgent(ctrl) + + // Set up expectations for the mock agents + mockPlanner.EXPECT().Name(gomock.Any()).Return("planner").AnyTimes() + mockPlanner.EXPECT().Description(gomock.Any()).Return("a planner agent").AnyTimes() + + mockExecutor.EXPECT().Name(gomock.Any()).Return("executor").AnyTimes() + mockExecutor.EXPECT().Description(gomock.Any()).Return("an executor agent").AnyTimes() + + mockReplanner.EXPECT().Name(gomock.Any()).Return("replanner").AnyTimes() + mockReplanner.EXPECT().Description(gomock.Any()).Return("a replanner agent").AnyTimes() + + // Create a plan + originalPlan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + // Create an updated plan with fewer steps (after replanning) + updatedPlan := &defaultPlan{Steps: []string{"Updated Step 2", "Updated Step 3"}} + // Create execute result + originalExecuteResult := "Execution result for Step 1" + updatedExecuteResult := "Execution result for Updated Step 2" + + // Create user input + userInput := []adk.Message{schema.UserMessage("User task input")} + + finalResponse := &Response{Response: "Final response to user after executing all steps"} + + // Mock the planner Run method to set the original plan + mockPlanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + // Set the plan in the session + adk.AddSessionValue(ctx, PlanSessionKey, originalPlan) + adk.AddSessionValue(ctx, UserInputSessionKey, userInput) + + // Send a message event + planJSON, _ := sonic.MarshalString(originalPlan) + msg := schema.AssistantMessage(planJSON, nil) + event := adk.EventFromMessage(msg, nil, schema.Assistant, "") + generator.Send(event) + generator.Close() + + return iterator + }, + ).Times(1) + + // Mock the executor Run method to set the execute result + mockExecutor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + plan, _ := adk.GetSessionValue(ctx, PlanSessionKey) + currentPlan := plan.(*defaultPlan) + var msg adk.Message + // Check if this is the first replanning (original plan has 3 steps) + if len(currentPlan.Steps) == 3 { + msg = schema.AssistantMessage(originalExecuteResult, nil) + adk.AddSessionValue(ctx, ExecutedStepSessionKey, originalExecuteResult) + } else { + msg = schema.AssistantMessage(updatedExecuteResult, nil) + adk.AddSessionValue(ctx, ExecutedStepSessionKey, updatedExecuteResult) + } + event := adk.EventFromMessage(msg, nil, schema.Assistant, "") + generator.Send(event) + generator.Close() + + return iterator + }, + ).Times(2) + + // Mock the replanner Run method to first update the plan, then respond to user + mockReplanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + // First call: Update the plan + // Get the current plan from the session + plan, _ := adk.GetSessionValue(ctx, PlanSessionKey) + currentPlan := plan.(*defaultPlan) + + // Check if this is the first replanning (original plan has 3 steps) + if len(currentPlan.Steps) == 3 { + // Send a message event with the updated plan + planJSON, _ := sonic.MarshalString(updatedPlan) + msg := schema.AssistantMessage(planJSON, nil) + event := adk.EventFromMessage(msg, nil, schema.Assistant, "") + generator.Send(event) + + // Set the updated plan & execute result in the session + adk.AddSessionValue(ctx, PlanSessionKey, updatedPlan) + adk.AddSessionValue(ctx, ExecutedStepsSessionKey, []ExecutedStep{{ + Step: currentPlan.Steps[0], + Result: originalExecuteResult, + }}) + } else { + // Second call: Respond to user + responseJSON, err := sonic.MarshalString(finalResponse) + assert.NoError(t, err) + msg := schema.AssistantMessage(responseJSON, nil) + event := adk.EventFromMessage(msg, nil, schema.Assistant, "") + generator.Send(event) + + // Send exit action + action := adk.NewExitAction() + generator.Send(&adk.AgentEvent{Action: action}) + } + + generator.Close() + return iterator + }, + ).Times(2) + + conf := &Config{ + Planner: mockPlanner, + Executor: mockExecutor, + Replanner: mockReplanner, + } + + // Create the plan execute agent + agent, err := New(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, agent) + + // Run the agent + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent}) + iterator := runner.Run(ctx, userInput) + + // Collect all events + var events []*adk.AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + // Verify the events + assert.Greater(t, len(events), 0) + + for i, event := range events { + eventJSON, e := sonic.MarshalString(event) + assert.NoError(t, e) + t.Logf("event %d:\n%s", i, eventJSON) + } +} + +type interruptibleTool struct { + name string + t *testing.T +} + +func (m *interruptibleTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: m.name, + Desc: "A tool that requires human approval before execution", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "action": { + Type: schema.String, + Desc: "The action to perform", + Required: true, + }, + }), + }, nil +} + +func (m *interruptibleTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + wasInterrupted, _, _ := tool.GetInterruptState[any](ctx) + if !wasInterrupted { + return "", tool.Interrupt(ctx, fmt.Sprintf("Tool '%s' requires human approval", m.name)) + } + + isResumeTarget, hasData, data := tool.GetResumeContext[string](ctx) + if !isResumeTarget { + return "", tool.Interrupt(ctx, fmt.Sprintf("Tool '%s' requires human approval", m.name)) + } + + if hasData { + return fmt.Sprintf("Approved action executed with data: %s", data), nil + } + return "Approved action executed", nil +} + +type checkpointStore struct { + data map[string][]byte +} + +func newCheckpointStore() *checkpointStore { + return &checkpointStore{data: make(map[string][]byte)} +} + +func (s *checkpointStore) Set(_ context.Context, key string, value []byte) error { + s.data[key] = value + return nil +} + +func (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.data[key] + return v, ok, nil +} + +func TestPlanExecuteAgentInterruptResume(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) + + approvalTool := &interruptibleTool{name: "approve_action", t: t} + + plan := &defaultPlan{Steps: []string{"Execute action requiring approval", "Complete task"}} + userInput := []adk.Message{schema.UserMessage("Please execute the action")} + + mockPlanner := mockAdk.NewMockAgent(ctrl) + mockPlanner.EXPECT().Name(gomock.Any()).Return("planner").AnyTimes() + mockPlanner.EXPECT().Description(gomock.Any()).Return("a planner agent").AnyTimes() + + mockPlanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + adk.AddSessionValue(ctx, PlanSessionKey, plan) + adk.AddSessionValue(ctx, UserInputSessionKey, userInput) + + planJSON, _ := sonic.MarshalString(plan) + msg := schema.AssistantMessage(planJSON, nil) + event := adk.EventFromMessage(msg, nil, schema.Assistant, "") + generator.Send(event) + generator.Close() + + return iterator + }, + ).Times(1) + + toolCallMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "approve_action", + Arguments: `{"action": "execute"}`, + }, + }, + }) + + completionMsg := schema.AssistantMessage("Action approved and executed successfully", nil) + + mockToolCallingModel.EXPECT().WithTools(gomock.Any()).Return(mockToolCallingModel, nil).AnyTimes() + mockToolCallingModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(toolCallMsg, nil).Times(1) + mockToolCallingModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(completionMsg, nil).AnyTimes() + + executor, err := NewExecutor(ctx, &ExecutorConfig{ + Model: mockToolCallingModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{approvalTool}, + }, + }, + MaxIterations: 5, + }) + assert.NoError(t, err) + + mockReplanner := mockAdk.NewMockAgent(ctrl) + mockReplanner.EXPECT().Name(gomock.Any()).Return("replanner").AnyTimes() + mockReplanner.EXPECT().Description(gomock.Any()).Return("a replanner agent").AnyTimes() + + mockReplanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + responseJSON := `{"response":"Task completed successfully"}` + msg := schema.AssistantMessage(responseJSON, nil) + event := adk.EventFromMessage(msg, nil, schema.Assistant, "") + generator.Send(event) + + action := adk.NewBreakLoopAction("replanner") + generator.Send(&adk.AgentEvent{Action: action}) + + generator.Close() + return iterator + }, + ).AnyTimes() + + agent, err := New(ctx, &Config{ + Planner: mockPlanner, + Executor: executor, + Replanner: mockReplanner, + MaxIterations: 5, + }) + assert.NoError(t, err) + + store := newCheckpointStore() + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: agent, + CheckPointStore: store, + }) + + iter := runner.Run(ctx, userInput, adk.WithCheckPointID("test-interrupt-1")) + + var events []*adk.AgentEvent + var interruptEvent *adk.AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + events = append(events, event) + } + + t.Logf("Total events received: %d", len(events)) + for i, event := range events { + eventJSON, _ := sonic.MarshalString(event) + t.Logf("Event %d: %s", i, eventJSON) + } + + if interruptEvent == nil { + t.Fatal("Expected an interrupt event from the tool, but none was received") + } + + assert.NotNil(t, interruptEvent.Action.Interrupted, "Should have interrupt info") + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, "Should have interrupt contexts") + + t.Logf("Interrupt event received with %d contexts", len(interruptEvent.Action.Interrupted.InterruptContexts)) + for i, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + t.Logf("Interrupt context %d: ID=%s, Info=%v, Address=%v", i, ctx.ID, ctx.Info, ctx.Address) + } + + var toolInterruptID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + toolInterruptID = intCtx.ID + break + } + } + assert.NotEmpty(t, toolInterruptID, "Should have a root cause interrupt ID") + + t.Logf("Attempting to resume with interrupt ID: %s", toolInterruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, "test-interrupt-1", &adk.ResumeParams{ + Targets: map[string]any{ + toolInterruptID: "approved", + }, + }) + assert.NoError(t, err, "Resume should not error") + assert.NotNil(t, resumeIter, "Resume iterator should not be nil") + + var resumeEvents []*adk.AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + resumeEvents = append(resumeEvents, event) + } + + assert.NotEmpty(t, resumeEvents, "Should have resume events") + + for _, event := range resumeEvents { + assert.NoError(t, event.Err, "Resume event should not have error") + } + + var hasToolResponse, hasAssistantCompletion, hasBreakLoop bool + for _, event := range resumeEvents { + if event.Output != nil && event.Output.MessageOutput != nil { + msg := event.Output.MessageOutput.Message + if msg != nil { + if msg.Role == "tool" && strings.Contains(msg.Content, "Approved action executed") { + hasToolResponse = true + } + if msg.Role == "assistant" && strings.Contains(msg.Content, "approved") { + hasAssistantCompletion = true + } + } + } + if event.Action != nil && event.Action.BreakLoop != nil && event.Action.BreakLoop.Done { + hasBreakLoop = true + } + } + + assert.True(t, hasToolResponse, "Should have tool response with approved action") + assert.True(t, hasAssistantCompletion, "Should have assistant completion message") + assert.True(t, hasBreakLoop, "Should have break loop action indicating completion") +} + +// slowChatModel is a ChatModel that blocks for a configurable duration. +type slowChatModel struct { + delay time.Duration + response *schema.Message + startedChan chan struct{} + startedOnce sync.Once +} + +func (m *slowChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + m.startedOnce.Do(func() { + close(m.startedChan) + }) + + select { + case <-time.After(m.delay): + return m.response, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (m *slowChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + sr, sw := schema.Pipe[*schema.Message](1) + sw.Send(msg, nil) + sw.Close() + return sr, nil +} + +func (m *slowChatModel) BindTools(tools []*schema.ToolInfo) error { return nil } +func (m *slowChatModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +// TestWithCancel_PlanExecute_DuringExecution verifies that cancel works +// during the executor (ChatModelAgent) phase of the PlanExecute agent. +func TestWithCancel_PlanExecute_DuringExecution(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Planner: returns a plan quickly + mockPlanner := mockAdk.NewMockAgent(ctrl) + mockPlanner.EXPECT().Name(gomock.Any()).Return("planner").AnyTimes() + mockPlanner.EXPECT().Description(gomock.Any()).Return("a planner agent").AnyTimes() + + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2"}} + userInput := []adk.Message{schema.UserMessage("test task")} + + mockPlanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + adk.AddSessionValue(ctx, PlanSessionKey, plan) + adk.AddSessionValue(ctx, UserInputSessionKey, userInput) + planJSON, _ := sonic.MarshalString(plan) + msg := schema.AssistantMessage(planJSON, nil) + generator.Send(adk.EventFromMessage(msg, nil, schema.Assistant, "")) + generator.Close() + return iterator + }, + ).Times(1) + + // Executor: uses a slow model that we can cancel + executorStarted := make(chan struct{}) + slowModel := &slowChatModel{ + delay: 5 * time.Second, + response: schema.AssistantMessage("step result", nil), + startedChan: executorStarted, + } + + executor, err := NewExecutor(ctx, &ExecutorConfig{ + Model: slowModel, + MaxIterations: 5, + }) + assert.NoError(t, err) + + // Replanner: should not be reached since we cancel during executor + mockReplanner := mockAdk.NewMockAgent(ctrl) + mockReplanner.EXPECT().Name(gomock.Any()).Return("replanner").AnyTimes() + mockReplanner.EXPECT().Description(gomock.Any()).Return("a replanner agent").AnyTimes() + + agent, err := New(ctx, &Config{ + Planner: mockPlanner, + Executor: executor, + Replanner: mockReplanner, + MaxIterations: 5, + }) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent}) + + cancelOpt, cancelFn := adk.WithCancel() + iter := runner.Run(ctx, userInput, cancelOpt) + + // Wait for the executor's model to start + select { + case <-executorStarted: + case <-time.After(10 * time.Second): + t.Fatal("Executor model did not start") + } + + time.Sleep(50 * time.Millisecond) + + // Cancel should NOT return ErrExecutionEnded + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err, "Cancel during executor should succeed") + + hasCancelError := false + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *adk.CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + hasCancelError = true + } + } + + assert.True(t, hasCancelError, "Should have CancelError event") +} + +// TestWithCancel_PlanExecute_BetweenTransitions verifies that cancel works +// when fired between agent transitions (e.g., after planner, before executor starts). +func TestWithCancel_PlanExecute_BetweenTransitions(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + plannerDone := make(chan struct{}) + + // Planner: signals when done + mockPlanner := mockAdk.NewMockAgent(ctrl) + mockPlanner.EXPECT().Name(gomock.Any()).Return("planner").AnyTimes() + mockPlanner.EXPECT().Description(gomock.Any()).Return("a planner agent").AnyTimes() + + plan := &defaultPlan{Steps: []string{"Step 1"}} + userInput := []adk.Message{schema.UserMessage("test task")} + + mockPlanner.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + go func() { + defer generator.Close() + adk.AddSessionValue(ctx, PlanSessionKey, plan) + adk.AddSessionValue(ctx, UserInputSessionKey, userInput) + planJSON, _ := sonic.MarshalString(plan) + msg := schema.AssistantMessage(planJSON, nil) + generator.Send(adk.EventFromMessage(msg, nil, schema.Assistant, "")) + close(plannerDone) + }() + return iterator + }, + ).Times(1) + + // Executor: slow model to give time to observe cancel + executorModelStarted := make(chan struct{}) + slowExecModel := &slowChatModel{ + delay: 5 * time.Second, + response: schema.AssistantMessage("step result", nil), + startedChan: executorModelStarted, + } + + executor, err := NewExecutor(ctx, &ExecutorConfig{ + Model: slowExecModel, + MaxIterations: 5, + }) + assert.NoError(t, err) + + mockReplanner := mockAdk.NewMockAgent(ctrl) + mockReplanner.EXPECT().Name(gomock.Any()).Return("replanner").AnyTimes() + mockReplanner.EXPECT().Description(gomock.Any()).Return("a replanner agent").AnyTimes() + + agent, err := New(ctx, &Config{ + Planner: mockPlanner, + Executor: executor, + Replanner: mockReplanner, + MaxIterations: 5, + }) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: agent}) + + cancelOpt, cancelFn := adk.WithCancel() + iter := runner.Run(ctx, userInput, cancelOpt) + + // Wait for planner to finish, then cancel before executor has a chance to produce output + select { + case <-plannerDone: + case <-time.After(10 * time.Second): + t.Fatal("Planner did not finish") + } + + // Cancel after planner, during executor phase + // The executor is a ChatModelAgent which will handle the cancel + select { + case <-executorModelStarted: + case <-time.After(10 * time.Second): + t.Fatal("Executor model did not start") + } + + start := time.Now() + handle, _ := cancelFn() + err = handle.Wait() + assert.NoError(t, err, "Cancel between transitions should succeed") + + hasCancelError := false + for { + event, ok := iter.Next() + if !ok { + break + } + var ce *adk.CancelError + if event.Err != nil && errors.As(event.Err, &ce) { + hasCancelError = true + } + } + elapsed := time.Since(start) + + assert.True(t, hasCancelError, "Should have CancelError event") + assert.True(t, elapsed < 3*time.Second, "Should complete quickly after cancel, elapsed: %v", elapsed) +} diff --git a/adk/prebuilt/planexecute/utils_test.go b/adk/prebuilt/planexecute/utils_test.go new file mode 100644 index 0000000..2e01f7a --- /dev/null +++ b/adk/prebuilt/planexecute/utils_test.go @@ -0,0 +1,58 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package planexecute + +import ( + "context" + + "github.com/cloudwego/eino/adk" +) + +type outputSessionKVsAgent struct { + adk.Agent +} + +func (o *outputSessionKVsAgent) Run(ctx context.Context, input *adk.AgentInput, + options ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + + iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + + iterator_ := o.Agent.Run(ctx, input, options...) + go func() { + defer generator.Close() + for { + event, ok := iterator_.Next() + if !ok { + break + } + generator.Send(event) + } + + kvs := adk.GetSessionValues(ctx) + + event := &adk.AgentEvent{ + Output: &adk.AgentOutput{CustomizedOutput: kvs}, + } + generator.Send(event) + }() + + return iterator +} + +func agentOutputSessionKVs(_ context.Context, agent adk.Agent) (adk.Agent, error) { + return &outputSessionKVsAgent{Agent: agent}, nil +} diff --git a/adk/prebuilt/supervisor/supervisor.go b/adk/prebuilt/supervisor/supervisor.go new file mode 100644 index 0000000..62e6d1d --- /dev/null +++ b/adk/prebuilt/supervisor/supervisor.go @@ -0,0 +1,120 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package supervisor implements the supervisor pattern for multi-agent systems, +// where a central agent coordinates a set of sub-agents. +// +// # Unified Tracing +// +// The supervisor pattern provides unified tracing support through an internal container. +// When using callbacks (e.g., for tracing or observability), the entire supervisor structure +// (supervisor agent + all sub-agents) shares a single trace root. This means: +// - OnStart is invoked once at the supervisor container level +// - The callback-enriched context (containing parent span info) is propagated to all agents +// - All agents within the supervisor appear as children of the same trace root +// +// This is achieved by wrapping the supervisor structure in an internal container that acts +// as the single entry point for tracing. The container delegates all execution to the +// underlying agents while providing a unified identity for callbacks. +package supervisor + +import ( + "context" + + "github.com/cloudwego/eino/adk" +) + +// Config is the configuration for creating a supervisor-based multi-agent system. +// +// NOT RECOMMENDED: Supervisor is built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +type Config struct { + // Supervisor specifies the agent that will act as the supervisor, coordinating and managing the sub-agents. + Supervisor adk.Agent + + // SubAgents specifies the list of agents that will be supervised and coordinated by the supervisor agent. + SubAgents []adk.Agent +} + +// supervisorContainer wraps the entire supervisor structure to provide unified tracing. +// When callbacks are registered (e.g., via Runner.Query with WithCallbacks), OnStart/OnEnd +// are invoked once for this container, creating a single trace root. The callback-enriched +// context is then propagated to the supervisor and all sub-agents, ensuring they share +// the same trace parent. +// +// This container implements Agent and ResumableAgent by delegating to the inner agent. +// It provides its own Name and GetType for callback identification. +type supervisorContainer struct { + name string + inner adk.ResumableAgent +} + +func (s *supervisorContainer) Name(_ context.Context) string { + return s.name +} + +func (s *supervisorContainer) Description(ctx context.Context) string { + return s.inner.Description(ctx) +} + +func (s *supervisorContainer) GetType() string { + return "Supervisor" +} + +func (s *supervisorContainer) Run(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + return s.inner.Run(ctx, input, opts...) +} + +func (s *supervisorContainer) Resume(ctx context.Context, info *adk.ResumeInfo, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + return s.inner.Resume(ctx, info, opts...) +} + +// New creates a supervisor-based multi-agent system with the given configuration. +// +// In the supervisor pattern, a designated supervisor agent coordinates multiple sub-agents. +// The supervisor can delegate tasks to sub-agents and receive their responses, while +// sub-agents can only communicate with the supervisor (not with each other directly). +// This hierarchical structure enables complex problem-solving through coordinated agent interactions. +// +// The returned agent is wrapped in an internal container that provides unified tracing. +// When used with Runner and callbacks, all agents within the supervisor structure will +// share the same trace root, making it easy to observe the entire multi-agent execution +// as a single logical unit. +// +// NOT RECOMMENDED: Supervisor is built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +func New(ctx context.Context, conf *Config) (adk.ResumableAgent, error) { + subAgents := make([]adk.Agent, 0, len(conf.SubAgents)) + supervisorName := conf.Supervisor.Name(ctx) + for _, subAgent := range conf.SubAgents { + subAgents = append(subAgents, adk.AgentWithDeterministicTransferTo(ctx, &adk.DeterministicTransferConfig{ + Agent: subAgent, + ToAgentNames: []string{supervisorName}, + })) + } + + inner, err := adk.SetSubAgents(ctx, conf.Supervisor, subAgents) + if err != nil { + return nil, err + } + + return &supervisorContainer{ + name: supervisorName, + inner: inner, + }, nil +} diff --git a/adk/prebuilt/supervisor/supervisor_test.go b/adk/prebuilt/supervisor/supervisor_test.go new file mode 100644 index 0000000..2dbef75 --- /dev/null +++ b/adk/prebuilt/supervisor/supervisor_test.go @@ -0,0 +1,1159 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package supervisor + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockAdk "github.com/cloudwego/eino/internal/mock/adk" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +// TestNewSupervisor tests the New function +func TestNewSupervisor(t *testing.T) { + ctx := context.Background() + + // Create a mock controller + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // Create mock agents + supervisorAgent := mockAdk.NewMockAgent(ctrl) + subAgent1 := mockAdk.NewMockAgent(ctrl) + subAgent2 := mockAdk.NewMockAgent(ctrl) + + supervisorAgent.EXPECT().Name(gomock.Any()).Return("SupervisorAgent").AnyTimes() + supervisorAgent.EXPECT().Description(gomock.Any()).Return("Supervisor agent description").AnyTimes() + subAgent1.EXPECT().Name(gomock.Any()).Return("SubAgent1").AnyTimes() + subAgent2.EXPECT().Name(gomock.Any()).Return("SubAgent2").AnyTimes() + + aMsg, tMsg := adk.GenTransferMessages(ctx, "SubAgent1") + i, g := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + g.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, "")) + event := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName) + event.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "SubAgent1"}} + g.Send(event) + g.Close() + supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1) + + i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]() + subAgent1Msg := schema.AssistantMessage("SubAgent1", nil) + g.Send(adk.EventFromMessage(subAgent1Msg, nil, schema.Assistant, "")) + g.Close() + subAgent1.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1) + + aMsg, tMsg = adk.GenTransferMessages(ctx, "SubAgent2 message") + i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]() + g.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, "")) + event = adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName) + event.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "SubAgent2"}} + g.Send(event) + g.Close() + supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1) + + i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]() + subAgent2Msg := schema.AssistantMessage("SubAgent2 message", nil) + g.Send(adk.EventFromMessage(subAgent2Msg, nil, schema.Assistant, "")) + g.Close() + subAgent2.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1) + + i, g = adk.NewAsyncIteratorPair[*adk.AgentEvent]() + finishMsg := schema.AssistantMessage("finish", nil) + g.Send(adk.EventFromMessage(finishMsg, nil, schema.Assistant, "")) + g.Close() + supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i).Times(1) + + conf := &Config{ + Supervisor: supervisorAgent, + SubAgents: []adk.Agent{subAgent1, subAgent2}, + } + + multiAgent, err := New(ctx, conf) + assert.NoError(t, err) + assert.NotNil(t, multiAgent) + assert.Equal(t, "SupervisorAgent", multiAgent.Name(ctx)) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent}) + aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("test")}) + + // transfer to agent1 + event, ok := aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SupervisorAgent", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls)) + + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SupervisorAgent", event.AgentName) + assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role) + assert.Equal(t, "SubAgent1", event.Action.TransferToAgent.DestAgentName) + + // agent1's output + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SubAgent1", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.Equal(t, subAgent1Msg.Content, event.Output.MessageOutput.Message.Content) + + // transfer back to supervisor + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SubAgent1", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls)) + + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SubAgent1", event.AgentName) + assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role) + assert.Equal(t, "SupervisorAgent", event.Action.TransferToAgent.DestAgentName) + + // transfer to agent2 + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SupervisorAgent", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls)) + + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SupervisorAgent", event.AgentName) + assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role) + assert.Equal(t, "SubAgent2", event.Action.TransferToAgent.DestAgentName) + + // agent1's output + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SubAgent2", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.Equal(t, subAgent2Msg.Content, event.Output.MessageOutput.Message.Content) + + // transfer back to supervisor + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SubAgent2", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.NotEqual(t, 0, len(event.Output.MessageOutput.Message.ToolCalls)) + + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SubAgent2", event.AgentName) + assert.Equal(t, schema.Tool, event.Output.MessageOutput.Role) + assert.Equal(t, "SupervisorAgent", event.Action.TransferToAgent.DestAgentName) + + // finish + event, ok = aIter.Next() + assert.True(t, ok) + assert.Equal(t, "SupervisorAgent", event.AgentName) + assert.Equal(t, schema.Assistant, event.Output.MessageOutput.Role) + assert.Equal(t, finishMsg.Content, event.Output.MessageOutput.Message.Content) +} + +type approvalInfo struct { + ToolName string + ArgumentsInJSON string + ToolCallID string +} + +func (ai *approvalInfo) String() string { + return fmt.Sprintf("tool '%s' interrupted with arguments '%s', waiting for approval", + ai.ToolName, ai.ArgumentsInJSON) +} + +type approvalResult struct { + Approved bool + DisapproveReason *string +} + +func init() { + schema.Register[*approvalInfo]() + schema.Register[*approvalResult]() +} + +type approvableTool struct { + name string + t *testing.T +} + +func (m *approvableTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: m.name, + Desc: "A tool that requires approval before execution", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "action": {Type: schema.String, Desc: "The action to perform"}, + }), + }, nil +} + +func (m *approvableTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + wasInterrupted, _, storedArguments := tool.GetInterruptState[string](ctx) + if !wasInterrupted { + return "", tool.StatefulInterrupt(ctx, &approvalInfo{ + ToolName: m.name, + ArgumentsInJSON: argumentsInJSON, + ToolCallID: compose.GetToolCallID(ctx), + }, argumentsInJSON) + } + + isResumeTarget, hasData, data := tool.GetResumeContext[*approvalResult](ctx) + if !isResumeTarget { + return "", tool.StatefulInterrupt(ctx, &approvalInfo{ + ToolName: m.name, + ArgumentsInJSON: storedArguments, + ToolCallID: compose.GetToolCallID(ctx), + }, storedArguments) + } + + if !hasData { + return "", fmt.Errorf("tool '%s' resumed with no data", m.name) + } + + if data.Approved { + return fmt.Sprintf("Tool '%s' executed successfully with args: %s", m.name, storedArguments), nil + } + + if data.DisapproveReason != nil { + return fmt.Sprintf("Tool '%s' disapproved, reason: %s", m.name, *data.DisapproveReason), nil + } + + return fmt.Sprintf("Tool '%s' disapproved", m.name), nil +} + +type checkpointStore struct { + data map[string][]byte +} + +func newCheckpointStore() *checkpointStore { + return &checkpointStore{data: make(map[string][]byte)} +} + +func (s *checkpointStore) Set(_ context.Context, key string, value []byte) error { + s.data[key] = value + return nil +} + +func (s *checkpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + v, ok := s.data[key] + return v, ok, nil +} + +type namedAgent struct { + adk.ResumableAgent + name string + description string +} + +func (n *namedAgent) Name(_ context.Context) string { + return n.name +} + +func (n *namedAgent) Description(_ context.Context) string { + return n.description +} + +func TestNestedSupervisorInterruptResume(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockOuterSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockInnerSupervisorModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockWorkerModel := mockModel.NewMockToolCallingChatModel(ctrl) + + paymentTool := &approvableTool{name: "process_payment", t: t} + + userInput := []adk.Message{schema.UserMessage("Process a payment of $1000")} + + mockWorkerModel.EXPECT().WithTools(gomock.Any()).Return(mockWorkerModel, nil).AnyTimes() + + workerToolCallMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_payment_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "process_payment", + Arguments: `{"action": "process $1000 payment"}`, + }, + }, + }) + mockWorkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workerToolCallMsg, nil).Times(1) + + workerCompletionMsg := schema.AssistantMessage("Payment processed successfully", nil) + mockWorkerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workerCompletionMsg, nil).AnyTimes() + + workerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "payment_worker", + Description: "the agent responsible for processing payments", + Instruction: "You are a payment processing worker. Use the process_payment tool to handle payments.", + Model: mockWorkerModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{paymentTool}, + }, + }, + }) + assert.NoError(t, err) + + mockInnerSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockInnerSupervisorModel, nil).AnyTimes() + + innerTransferMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "inner_transfer_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "transfer_to_agent", + Arguments: `{"agent_name":"payment_worker"}`, + }, + }, + }) + mockInnerSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(innerTransferMsg, nil).Times(1) + + innerFinalMsg := schema.AssistantMessage("Payment has been processed and approved.", nil) + mockInnerSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(innerFinalMsg, nil).AnyTimes() + + innerSupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "payment_supervisor", + Description: "the supervisor agent responsible for payment operations", + Instruction: "You are a payment supervisor. Delegate payment tasks to payment_worker.", + Model: mockInnerSupervisorModel, + Exit: &adk.ExitTool{}, + }) + assert.NoError(t, err) + + innerSupervisorAgent, err := New(ctx, &Config{ + Supervisor: innerSupervisorChatAgent, + SubAgents: []adk.Agent{workerAgent}, + }) + assert.NoError(t, err) + + innerSupervisorWrapped := &namedAgent{ + ResumableAgent: innerSupervisorAgent, + name: "payment_department", + description: "the department responsible for all payment-related operations", + } + + mockOuterSupervisorModel.EXPECT().WithTools(gomock.Any()).Return(mockOuterSupervisorModel, nil).AnyTimes() + + outerTransferMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "outer_transfer_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "transfer_to_agent", + Arguments: `{"agent_name":"payment_department"}`, + }, + }, + }) + mockOuterSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(outerTransferMsg, nil).Times(1) + + outerFinalMsg := schema.AssistantMessage("The payment request has been fully processed by the payment department.", nil) + mockOuterSupervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(outerFinalMsg, nil).AnyTimes() + + outerSupervisorChatAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "company_coordinator", + Description: "the top-level coordinator for company operations", + Instruction: "You are the company coordinator. Route payment requests to payment_department.", + Model: mockOuterSupervisorModel, + Exit: &adk.ExitTool{}, + }) + assert.NoError(t, err) + + outerSupervisorAgent, err := New(ctx, &Config{ + Supervisor: outerSupervisorChatAgent, + SubAgents: []adk.Agent{innerSupervisorWrapped}, + }) + assert.NoError(t, err) + + outerSupervisorWrapped := &namedAgent{ + ResumableAgent: outerSupervisorAgent, + name: "headquarters", + description: "the company headquarters that coordinates all departments", + } + + store := newCheckpointStore() + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: outerSupervisorWrapped, + CheckPointStore: store, + }) + + t.Log("========================================") + t.Log("Starting Nested Supervisor Integration Test (with namedAgent wrappers)") + t.Log("Hierarchy: headquarters(wrapper) -> company_coordinator -> payment_department(wrapper) -> payment_supervisor -> payment_worker -> process_payment tool") + t.Log("========================================") + + checkpointID := "test-nested-supervisor-1" + iter := runner.Run(ctx, userInput, adk.WithCheckPointID(checkpointID)) + + var interruptEvent *adk.AgentEvent + eventCount := 0 + for { + event, ok := iter.Next() + if !ok { + break + } + eventCount++ + + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + t.Log("INTERRUPT DETECTED - Deep interrupt from tool within nested supervisor") + break + } + } + + if interruptEvent == nil { + t.Fatal("Expected an interrupt event from the process_payment tool, but none was received") + } + + assert.NotNil(t, interruptEvent.Action.Interrupted, "Should have interrupt info") + assert.NotEmpty(t, interruptEvent.Action.Interrupted.InterruptContexts, "Should have interrupt contexts") + + var toolInterruptID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + toolInterruptID = intCtx.ID + break + } + } + assert.NotEmpty(t, toolInterruptID, "Should have a root cause interrupt ID") + + t.Logf("Resuming with approval for interrupt ID: %s", toolInterruptID) + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{ + Targets: map[string]any{ + toolInterruptID: &approvalResult{Approved: true}, + }, + }) + assert.NoError(t, err, "Resume should not error") + assert.NotNil(t, resumeIter, "Resume iterator should not be nil") + + var resumeEvents []*adk.AgentEvent + for { + event, ok := resumeIter.Next() + if !ok { + break + } + resumeEvents = append(resumeEvents, event) + } + + assert.NotEmpty(t, resumeEvents, "Should have resume events after approval") + + for _, event := range resumeEvents { + assert.NoError(t, event.Err, "Resume event should not have error") + } + + var hasToolResponse, hasTransferBack bool + for _, event := range resumeEvents { + if event.Output != nil && event.Output.MessageOutput != nil { + msg := event.Output.MessageOutput.Message + if msg != nil && msg.Role == "tool" && strings.Contains(msg.Content, "executed successfully") { + hasToolResponse = true + } + } + if event.Action != nil && event.Action.TransferToAgent != nil { + if event.Action.TransferToAgent.DestAgentName == "company_coordinator" { + hasTransferBack = true + } + } + } + + assert.True(t, hasToolResponse, "Should have tool response indicating successful payment processing") + assert.True(t, hasTransferBack, "Should have transfer back to outer supervisor indicating completion") +} + +func TestSupervisorExit(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + supervisorAgent := mockAdk.NewMockAgent(ctrl) + subAgent := mockAdk.NewMockAgent(ctrl) + + supervisorAgent.EXPECT().Name(gomock.Any()).Return("Supervisor").AnyTimes() + supervisorAgent.EXPECT().Description(gomock.Any()).Return("Supervisor description").AnyTimes() + subAgent.EXPECT().Name(gomock.Any()).Return("SubAgent").AnyTimes() + + // 1. Supervisor transfers to SubAgent + aMsg, tMsg := adk.GenTransferMessages(ctx, "SubAgent") + i1, g1 := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + g1.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, "")) + event1 := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName) + event1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "SubAgent"}} + g1.Send(event1) + g1.Close() + supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i1).Times(1) + + // 2. SubAgent emits Exit action + i2, g2 := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + exitEvent := &adk.AgentEvent{ + AgentName: "SubAgent", + Action: &adk.AgentAction{Exit: true}, + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Role: schema.Assistant, + Message: schema.AssistantMessage("Exiting...", nil), + }, + }, + } + g2.Send(exitEvent) + g2.Close() + subAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i2).Times(1) + + conf := &Config{ + Supervisor: supervisorAgent, + SubAgents: []adk.Agent{subAgent}, + } + + multiAgent, err := New(ctx, conf) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent}) + aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("test")}) + + // Collect events + var events []*adk.AgentEvent + for { + event, ok := aIter.Next() + if !ok { + break + } + events = append(events, event) + } + + foundExit := false + foundTransferBack := false + + for _, e := range events { + if e.Action != nil { + if e.Action.Exit { + foundExit = true + } + if e.Action.TransferToAgent != nil && e.Action.TransferToAgent.DestAgentName == "Supervisor" { + foundTransferBack = true + } + } + } + + assert.True(t, foundExit, "Should have found Exit action") + assert.False(t, foundTransferBack, "Should NOT have found Transfer back to Supervisor after Exit") +} + +func TestNestedSupervisorExit(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + topSupervisor := mockAdk.NewMockAgent(ctrl) + midSupervisor := mockAdk.NewMockAgent(ctrl) + worker := mockAdk.NewMockAgent(ctrl) + + topSupervisor.EXPECT().Name(gomock.Any()).Return("TopSupervisor").AnyTimes() + topSupervisor.EXPECT().Description(gomock.Any()).Return("Top supervisor description").AnyTimes() + midSupervisor.EXPECT().Name(gomock.Any()).Return("MidSupervisor").AnyTimes() + midSupervisor.EXPECT().Description(gomock.Any()).Return("Mid supervisor description").AnyTimes() + worker.EXPECT().Name(gomock.Any()).Return("Worker").AnyTimes() + + // 1. TopSupervisor transfers to MidSupervisor + aMsg1, tMsg1 := adk.GenTransferMessages(ctx, "MidSupervisor") + i1, g1 := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + g1.Send(adk.EventFromMessage(aMsg1, nil, schema.Assistant, "")) + event1 := adk.EventFromMessage(tMsg1, nil, schema.Tool, tMsg1.ToolName) + event1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "MidSupervisor"}} + g1.Send(event1) + g1.Close() + topSupervisor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i1).AnyTimes() + + // 2. MidSupervisor transfers to Worker + aMsg2, tMsg2 := adk.GenTransferMessages(ctx, "Worker") + i2, g2 := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + g2.Send(adk.EventFromMessage(aMsg2, nil, schema.Assistant, "")) + event2 := adk.EventFromMessage(tMsg2, nil, schema.Tool, tMsg2.ToolName) + event2.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "Worker"}} + g2.Send(event2) + g2.Close() + midSupervisor.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i2).AnyTimes() + + // 3. Worker emits Exit action + i3, g3 := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + exitEvent := &adk.AgentEvent{ + AgentName: "Worker", + Action: &adk.AgentAction{Exit: true}, + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Role: schema.Assistant, + Message: schema.AssistantMessage("Worker Exiting...", nil), + }, + }, + } + g3.Send(exitEvent) + g3.Close() + worker.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).Return(i3).Times(1) + + // Build Nested System + // Mid System: MidSupervisor -> [Worker] + midSystem, err := New(ctx, &Config{ + Supervisor: midSupervisor, + SubAgents: []adk.Agent{worker}, + }) + assert.NoError(t, err) + // We need to give the midSystem the name "MidSupervisor" so TopSupervisor can find it + // supervisor.New returns a ResumableAgent that delegates Name() to the supervisor agent. + // So midSystem.Name() should already be "MidSupervisor" because midSupervisor.Name() is "MidSupervisor". + + // Top System: TopSupervisor -> [midSystem] + topSystem, err := New(ctx, &Config{ + Supervisor: topSupervisor, + SubAgents: []adk.Agent{midSystem}, + }) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: topSystem}) + aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("test nested exit")}) + + // Collect events + var events []*adk.AgentEvent + for { + event, ok := aIter.Next() + if !ok { + break + } + events = append(events, event) + } + + foundExit := false + foundTransferBackToMidAfterExit := false + foundTransferBackToTopAfterExit := false + + for _, e := range events { + if e.Action != nil { + if e.Action.Exit { + foundExit = true + } + if foundExit && e.Action.TransferToAgent != nil { + if e.Action.TransferToAgent.DestAgentName == "MidSupervisor" { + foundTransferBackToMidAfterExit = true + } + if e.Action.TransferToAgent.DestAgentName == "TopSupervisor" { + foundTransferBackToTopAfterExit = true + } + } + } + } + + assert.True(t, foundExit, "Should have found Exit action") + assert.False(t, foundTransferBackToMidAfterExit, "Should NOT have found Transfer back to MidSupervisor after Exit") + assert.False(t, foundTransferBackToTopAfterExit, "Should NOT have found Transfer back to TopSupervisor after Exit") +} + +func TestChatModelAgentInternalEventsExit(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + supervisorAgent := mockAdk.NewMockAgent(ctrl) + workerModel := mockModel.NewMockToolCallingChatModel(ctrl) + innerAgent := mockAdk.NewMockAgent(ctrl) + + supervisorAgent.EXPECT().Name(gomock.Any()).Return("Supervisor").AnyTimes() + supervisorAgent.EXPECT().Description(gomock.Any()).Return("Supervisor description").AnyTimes() + innerAgent.EXPECT().Name(gomock.Any()).Return("InnerAgent").AnyTimes() + innerAgent.EXPECT().Description(gomock.Any()).Return("Inner Agent Description").AnyTimes() + + // 1. Supervisor transfers to Worker (only once, then exits when worker transfers back) + supervisorRunCount := 0 + supervisorAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + supervisorRunCount++ + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + go func() { + defer gen.Close() + if supervisorRunCount == 1 { + aMsg, tMsg := adk.GenTransferMessages(ctx, "Worker") + gen.Send(adk.EventFromMessage(aMsg, nil, schema.Assistant, "")) + event1 := adk.EventFromMessage(tMsg, nil, schema.Tool, tMsg.ToolName) + event1.Action = &adk.AgentAction{TransferToAgent: &adk.TransferToAgentAction{DestAgentName: "Worker"}} + gen.Send(event1) + } else { + exitEvent := &adk.AgentEvent{ + AgentName: "Supervisor", + Action: &adk.AgentAction{Exit: true}, + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Role: schema.Assistant, + Message: schema.AssistantMessage("Supervisor done", nil), + }, + }, + } + gen.Send(exitEvent) + } + }() + return iter + }).AnyTimes() + + // 2. Worker runs, calls AgentTool (InnerAgent) + // Mock WorkerModel behavior + workerModel.EXPECT().WithTools(gomock.Any()).Return(workerModel, nil).AnyTimes() + + // 2.1 Worker generates tool call + toolCallMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_inner_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "InnerAgent", + Arguments: `{"request": "do exit"}`, + }, + }, + }) + workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(toolCallMsg, nil).Times(1) + + // 2.2 InnerAgent runs and emits Exit + innerAgent.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, input *adk.AgentInput, opts ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + iter, gen := adk.NewAsyncIteratorPair[*adk.AgentEvent]() + go func() { + defer gen.Close() + innerExitEvent := &adk.AgentEvent{ + AgentName: "InnerAgent", + Action: &adk.AgentAction{Exit: true}, + RunPath: []adk.RunStep{}, + Output: &adk.AgentOutput{ + MessageOutput: &adk.MessageVariant{ + Role: schema.Assistant, + Message: schema.AssistantMessage("Inner Exiting...", nil), + }, + }, + } + gen.Send(innerExitEvent) + }() + return iter + }).AnyTimes() + + // 2.3 Worker receives tool result (empty string or whatever AgentTool returns on exit/interrupt) + // AgentTool implementation details: if Exit action is present, it returns whatever output is there. + // The Exit action itself is passed as internal event. + + // 2.4 Worker generates final response + finalMsg := schema.AssistantMessage("Worker Finished", nil) + workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(finalMsg, nil).AnyTimes() + + // Build Worker Agent + agentTool := adk.NewAgentTool(ctx, innerAgent) + workerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "Worker", + Description: "Worker Agent", + Model: workerModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{agentTool}, + }, + EmitInternalEvents: true, // Key configuration + }, + }) + assert.NoError(t, err) + + // Build System + sys, err := New(ctx, &Config{ + Supervisor: supervisorAgent, + SubAgents: []adk.Agent{workerAgent}, + }) + assert.NoError(t, err) + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: sys}) + aIter := runner.Run(ctx, []adk.Message{schema.UserMessage("start")}) + + // Collect events + var events []*adk.AgentEvent + for { + event, ok := aIter.Next() + if !ok { + break + } + events = append(events, event) + } + + foundInnerExit := false + foundTransferBack := false + + for _, e := range events { + // Check for InnerAgent exit event (propagated as internal event) + if e.AgentName == "InnerAgent" && e.Action != nil && e.Action.Exit { + foundInnerExit = true + } + + // Check for transfer back to Supervisor + if e.AgentName == "Worker" && e.Action != nil && e.Action.TransferToAgent != nil && + e.Action.TransferToAgent.DestAgentName == "Supervisor" { + foundTransferBack = true + } + } + + assert.True(t, foundInnerExit, "Should have captured InnerAgent Exit event") + assert.True(t, foundTransferBack, "Should have found Transfer back to Supervisor (Worker should NOT be considered exited)") +} + +func TestSupervisorContainerUnifiedTracing(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + supervisorModel := mockModel.NewMockToolCallingChatModel(ctrl) + subAgentModel := mockModel.NewMockToolCallingChatModel(ctrl) + + supervisorModel.EXPECT().WithTools(gomock.Any()).Return(supervisorModel, nil).AnyTimes() + subAgentModel.EXPECT().WithTools(gomock.Any()).Return(subAgentModel, nil).AnyTimes() + + transferMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "transfer_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "transfer_to_agent", + Arguments: `{"agent_name":"SubAgent"}`, + }, + }, + }) + supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(transferMsg, nil).Times(1) + + subAgentResponse := schema.AssistantMessage("SubAgent response", nil) + subAgentModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(subAgentResponse, nil).Times(1) + + finalResponse := schema.AssistantMessage("Final response", nil) + supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(finalResponse, nil).Times(1) + + supervisorAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "SupervisorAgent", + Description: "Supervisor agent", + Instruction: "You are a supervisor", + Model: supervisorModel, + }) + assert.NoError(t, err) + + subAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "SubAgent", + Description: "Sub agent", + Instruction: "You are a sub agent", + Model: subAgentModel, + }) + assert.NoError(t, err) + + multiAgent, err := New(ctx, &Config{ + Supervisor: supervisorAgent, + SubAgents: []adk.Agent{subAgent}, + }) + assert.NoError(t, err) + + assert.Equal(t, "SupervisorAgent", multiAgent.Name(ctx)) + + typer, ok := multiAgent.(components.Typer) + assert.True(t, ok, "Should implement components.Typer") + assert.Equal(t, "Supervisor", typer.GetType()) + + var mu sync.Mutex + var onStartCalls []string + var onEndCalls []string + + handler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != adk.ComponentOfAgent { + return ctx + } + mu.Lock() + onStartCalls = append(onStartCalls, info.Name+":"+info.Type) + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != adk.ComponentOfAgent { + return ctx + } + mu.Lock() + onEndCalls = append(onEndCalls, info.Name+":"+info.Type) + mu.Unlock() + if agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + runner := adk.NewRunner(ctx, adk.RunnerConfig{Agent: multiAgent}) + iter := runner.Query(ctx, "hello", adk.WithCallbacks(handler)) + + for { + _, ok := iter.Next() + if !ok { + break + } + } + + mu.Lock() + defer mu.Unlock() + + assert.NotEmpty(t, onStartCalls, "Should have OnStart calls") + assert.Contains(t, onStartCalls, "SupervisorAgent:Supervisor", "Should have supervisor container OnStart with type 'Supervisor'") +} + +type traceContextKey struct{} + +func TestSupervisorContainerUnifiedTracingOnResume(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + supervisorModel := mockModel.NewMockToolCallingChatModel(ctrl) + workerModel := mockModel.NewMockToolCallingChatModel(ctrl) + + supervisorModel.EXPECT().WithTools(gomock.Any()).Return(supervisorModel, nil).AnyTimes() + workerModel.EXPECT().WithTools(gomock.Any()).Return(workerModel, nil).AnyTimes() + + paymentTool := &approvableTool{name: "process_payment", t: t} + + workerToolCallMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_payment_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "process_payment", + Arguments: `{"action": "process $1000 payment"}`, + }, + }, + }) + workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workerToolCallMsg, nil).Times(1) + + workerCompletionMsg := schema.AssistantMessage("Payment processed successfully", nil) + workerModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workerCompletionMsg, nil).AnyTimes() + + workerAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "Worker", + Description: "Worker agent", + Instruction: "You are a worker", + Model: workerModel, + ToolsConfig: adk.ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{paymentTool}, + }, + }, + }) + assert.NoError(t, err) + + transferMsg := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "transfer_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "transfer_to_agent", + Arguments: `{"agent_name":"Worker"}`, + }, + }, + }) + supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(transferMsg, nil).Times(1) + + finalResponse := schema.AssistantMessage("Final response", nil) + supervisorModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(finalResponse, nil).AnyTimes() + + supervisorAgent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ + Name: "SupervisorAgent", + Description: "Supervisor agent", + Instruction: "You are a supervisor", + Model: supervisorModel, + Exit: &adk.ExitTool{}, + }) + assert.NoError(t, err) + + multiAgent, err := New(ctx, &Config{ + Supervisor: supervisorAgent, + SubAgents: []adk.Agent{workerAgent}, + }) + assert.NoError(t, err) + + store := newCheckpointStore() + runner := adk.NewRunner(ctx, adk.RunnerConfig{ + Agent: multiAgent, + CheckPointStore: store, + }) + + var mu sync.Mutex + var runOnStartCalls []string + var resumeOnStartCalls []string + var resumeParentTraceIDs []string + + runHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != adk.ComponentOfAgent { + return ctx + } + mu.Lock() + runOnStartCalls = append(runOnStartCalls, info.Name+":"+info.Type) + mu.Unlock() + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != adk.ComponentOfAgent { + return ctx + } + if agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + checkpointID := "test-unified-tracing-resume" + iter := runner.Run(ctx, []adk.Message{schema.UserMessage("Process payment")}, adk.WithCallbacks(runHandler), adk.WithCheckPointID(checkpointID)) + + var interruptEvent *adk.AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + break + } + } + + assert.NotNil(t, interruptEvent, "Should have interrupt event") + + var toolInterruptID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + toolInterruptID = intCtx.ID + break + } + } + assert.NotEmpty(t, toolInterruptID, "Should have a root cause interrupt ID") + + mu.Lock() + t.Logf("Run OnStart calls: %v", runOnStartCalls) + assert.Contains(t, runOnStartCalls, "SupervisorAgent:Supervisor", "Run should have supervisor container OnStart") + mu.Unlock() + + resumeHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component != adk.ComponentOfAgent { + return ctx + } + mu.Lock() + resumeOnStartCalls = append(resumeOnStartCalls, info.Name+":"+info.Type) + parentID, _ := ctx.Value(traceContextKey{}).(string) + resumeParentTraceIDs = append(resumeParentTraceIDs, info.Name+":parent="+parentID) + mu.Unlock() + if info.Type == "Supervisor" { + return context.WithValue(ctx, traceContextKey{}, "supervisor-trace-id") + } + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if info.Component != adk.ComponentOfAgent { + return ctx + } + if agentOutput := adk.ConvAgentCallbackOutput(output); agentOutput != nil && agentOutput.Events != nil { + go func() { + for { + _, ok := agentOutput.Events.Next() + if !ok { + break + } + } + }() + } + return ctx + }). + Build() + + resumeIter, err := runner.ResumeWithParams(ctx, checkpointID, &adk.ResumeParams{ + Targets: map[string]any{ + toolInterruptID: &approvalResult{Approved: true}, + }, + }, adk.WithCallbacks(resumeHandler)) + assert.NoError(t, err) + + for { + event, ok := resumeIter.Next() + if !ok { + break + } + assert.NoError(t, event.Err) + } + + mu.Lock() + defer mu.Unlock() + + t.Logf("Resume OnStart calls: %v", resumeOnStartCalls) + t.Logf("Resume parent trace IDs: %v", resumeParentTraceIDs) + assert.NotEmpty(t, resumeOnStartCalls, "Should have OnStart calls during resume") + assert.Contains(t, resumeOnStartCalls, "SupervisorAgent:Supervisor", "Resume should have supervisor container OnStart with type 'Supervisor'") + + foundInnerSupervisorWithParent := false + for _, entry := range resumeParentTraceIDs { + if strings.Contains(entry, "SupervisorAgent") && !strings.Contains(entry, "parent=supervisor-trace-id") && entry != "SupervisorAgent:parent=" { + if strings.Contains(resumeOnStartCalls[0], "Supervisor") { + continue + } + } + if strings.Contains(entry, "parent=supervisor-trace-id") { + foundInnerSupervisorWithParent = true + } + } + assert.True(t, foundInnerSupervisorWithParent, + "Inner agents should have parent trace from Supervisor container during Resume. Got: %v", resumeParentTraceIDs) +} diff --git a/adk/react.go b/adk/react.go new file mode 100644 index 0000000..03ef045 --- /dev/null +++ b/adk/react.go @@ -0,0 +1,823 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "io" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// ErrExceedMaxIterations indicates the agent reached the maximum iterations limit. +var ErrExceedMaxIterations = errors.New("exceeds max iterations") + +type typedState[M MessageType] struct { + Messages []M + Extra map[string]any + + // ToolInfos contains the tool definitions passed to the model via model.WithTools. + // Managed by the framework and modifiable by BeforeModelRewriteState handlers. + ToolInfos []*schema.ToolInfo + + // DeferredToolInfos contains tool definitions for server-side deferred retrieval, + // passed to the model via model.WithDeferredTools. Nil when not in use. + DeferredToolInfos []*schema.ToolInfo + + // Internal fields below - do not access directly. + // Kept exported for backward compatibility with existing checkpoints. + HasReturnDirectly bool + ReturnDirectlyToolCallID string + ToolGenActions map[string]*AgentAction + AgentName string + RemainingIterations int + ReturnDirectlyEvent *TypedAgentEvent[M] + RetryAttempt int + ToolMsgIDs map[string]map[string]string // toolName → callID → eino message ID +} + +// State is the internal state of the ChatModelAgent. +// +// Deprecated: State is exported only for checkpoint backward compatibility. +// Do not use it directly. +type State = typedState[*schema.Message] + +type agenticState = typedState[*schema.AgenticMessage] + +const ( + stateGobNameV07 = "_eino_adk_react_state" + + // stateGobNameV080 is a v0.8.0-v0.8.3-only alias used after byte-patching + // raw checkpoint bytes in preprocessADKCheckpoint. + // It must stay the same byte length as stateGobNameV07 so the length-prefixed + // gob string in the stream remains valid. + stateGobNameV080 = "_eino_adk_state_v080_" +) + +func init() { + // Checkpoint compatibility notes: + // - ADK/compose checkpoints are gob-encoded and may store state behind `any`, so gob relies on + // an on-wire type name to choose a local Go type. + // - Gob allows only one local Go type per name, and it treats "struct wire" and "GobEncoder wire" + // as incompatible even if the name matches. + // + // This file maintains 2 epochs of *State decoding: + // - v0.7.* and current: "_eino_adk_react_state" + struct wire → decode into *State directly. + // State's exported fields are a superset of v0.7, so gob handles missing fields gracefully. + // - v0.8.0-v0.8.3: "_eino_adk_react_state" + GobEncoder wire → byte-patched to stateGobNameV080, + // decode into stateV080 and migrate. + schema.RegisterName[*State](stateGobNameV07) + schema.RegisterName[*stateV080](stateGobNameV080) + + schema.RegisterName[*typedState[*schema.AgenticMessage]]("_eino_adk_agentic_state") + schema.RegisterName[*TypedAgentEvent[*schema.AgenticMessage]]("_eino_adk_agentic_event") + + // backward compatibility when decoding checkpoints created by v0.8.0 - v0.8.3 + gob.Register(&AgentEvent{}) + gob.Register(0) + + schema.RegisterName[*TypedAgentInput[*schema.AgenticMessage]]("_eino_adk_agentic_agent_input") + schema.RegisterName[*typedAgentEventWrapper[*schema.AgenticMessage]]("_eino_adk_agentic_event_wrapper") + schema.RegisterName[*[]*typedAgentEventWrapper[*schema.AgenticMessage]]("_eino_adk_agentic_event_wrapper_slice") + schema.RegisterName[*reactInput]("_eino_adk_react_input") + schema.RegisterName[*agenticReactInput]("_eino_adk_agentic_react_input") +} + +func (s *typedState[M]) getReturnDirectlyEvent() *TypedAgentEvent[M] { + return s.ReturnDirectlyEvent +} + +func (s *typedState[M]) setReturnDirectlyEvent(event *TypedAgentEvent[M]) { + s.ReturnDirectlyEvent = event +} + +func (s *typedState[M]) getRetryAttempt() int { + return s.RetryAttempt +} + +func (s *typedState[M]) setRetryAttempt(attempt int) { + s.RetryAttempt = attempt +} + +func (s *typedState[M]) getReturnDirectlyToolCallID() string { + return s.ReturnDirectlyToolCallID +} + +func (s *typedState[M]) setReturnDirectlyToolCallID(id string) { + s.ReturnDirectlyToolCallID = id + s.HasReturnDirectly = id != "" +} + +func (s *typedState[M]) getToolGenActions() map[string]*AgentAction { + return s.ToolGenActions +} + +func (s *typedState[M]) setToolGenAction(key string, action *AgentAction) { + if s.ToolGenActions == nil { + s.ToolGenActions = make(map[string]*AgentAction) + } + s.ToolGenActions[key] = action +} + +func (s *typedState[M]) popToolGenAction(key string) *AgentAction { + if s.ToolGenActions == nil { + return nil + } + action := s.ToolGenActions[key] + delete(s.ToolGenActions, key) + return action +} + +func (s *typedState[M]) setToolMsgID(toolName, callID, msgID string) { + if s.ToolMsgIDs == nil { + s.ToolMsgIDs = make(map[string]map[string]string) + } + byCall := s.ToolMsgIDs[toolName] + if byCall == nil { + byCall = make(map[string]string) + s.ToolMsgIDs[toolName] = byCall + } + byCall[callID] = msgID +} + +func (s *typedState[M]) popToolMsgID(toolName, callID string) string { + if s.ToolMsgIDs == nil { + return "" + } + byCall := s.ToolMsgIDs[toolName] + if byCall == nil { + return "" + } + id := byCall[callID] + delete(byCall, callID) + if len(byCall) == 0 { + delete(s.ToolMsgIDs, toolName) + } + return id +} + +func (s *typedState[M]) getRemainingIterations() int { + return s.RemainingIterations +} + +func (s *typedState[M]) setRemainingIterations(iterations int) { + s.RemainingIterations = iterations +} + +func (s *typedState[M]) decrementRemainingIterations() { + current := s.getRemainingIterations() + s.RemainingIterations = current - 1 +} + +// stateV080 handles the v0.8.0-v0.8.3 checkpoint format. +// In those versions, *State implemented GobEncoder and was registered under +// "_eino_adk_react_state". GobEncode serialized a stateSerialization struct +// into opaque bytes. This type's GobDecode reads that format. +// It is registered under "_eino_adk_state_v080_" — a same-length alias used +// only after byte-patching the checkpoint data in preprocessADKCheckpoint. +type stateV080 struct { + Messages []Message + HasReturnDirectly bool + ReturnDirectlyToolCallID string + ToolGenActions map[string]*AgentAction + AgentName string + RemainingIterations int + RetryAttempt int + ReturnDirectlyEvent *AgentEvent + Extra map[string]any + Internals map[string]any +} + +// stateV080Serialization is the on-wire format that v0.8.0-v0.8.3 GobEncode produced. +// It is only used by stateV080.GobDecode to parse those legacy opaque bytes. +type stateV080Serialization stateV080 + +func (sc *stateV080) GobDecode(b []byte) error { + ss := &stateV080Serialization{} + if err := gob.NewDecoder(bytes.NewReader(b)).Decode(ss); err != nil { + return err + } + sc.Messages = ss.Messages + sc.HasReturnDirectly = ss.HasReturnDirectly + sc.ReturnDirectlyToolCallID = ss.ReturnDirectlyToolCallID + sc.ToolGenActions = ss.ToolGenActions + sc.AgentName = ss.AgentName + sc.RemainingIterations = ss.RemainingIterations + sc.Extra = ss.Extra + sc.Internals = ss.Internals + return nil +} + +// stateV080ToState converts a legacy *stateV080 (v0.8.0-v0.8.3) to a current *State. +func stateV080ToState(sc *stateV080) *State { + s := &State{ + Messages: sc.Messages, + HasReturnDirectly: sc.HasReturnDirectly, + ReturnDirectlyToolCallID: sc.ReturnDirectlyToolCallID, + ToolGenActions: sc.ToolGenActions, + AgentName: sc.AgentName, + RemainingIterations: sc.RemainingIterations, + Extra: sc.Extra, + } + if sc.ReturnDirectlyToolCallID != "" { + s.setReturnDirectlyToolCallID(sc.ReturnDirectlyToolCallID) + } + if sc.Internals != nil && s.RetryAttempt == 0 { + if v, ok := sc.Internals["_retryAttempt"].(int); ok { + s.RetryAttempt = v + } + } + if sc.Internals != nil && s.ReturnDirectlyEvent == nil { + if v, ok := sc.Internals["_returnDirectlyEvent"].(*AgentEvent); ok { + s.ReturnDirectlyEvent = v + } + } + return s +} + +// SendToolGenAction attaches an AgentAction to the next tool event emitted for the +// current tool execution. +// +// Where/when to use: +// - Invoke within a tool's Run (Invokable/Streamable) implementation to include +// an action alongside that tool's output event. +// - The action is scoped by the current tool call context: if a ToolCallID is +// available, it is used as the key to support concurrent calls of the same +// tool with different parameters; otherwise, the provided toolName is used. +// - The stored action is ephemeral and will be popped and attached to the tool +// event when the tool finishes (including streaming completion). +// +// Limitation: +// - This function is intended for use within ChatModelAgent runs only. It relies +// on ChatModelAgent's internal State to store and pop actions, which is not +// available in other agent types. +func SendToolGenAction(ctx context.Context, toolName string, action *AgentAction) error { + key := toolName + toolCallID := compose.GetToolCallID(ctx) + if len(toolCallID) > 0 { + key = toolCallID + } + + return compose.ProcessState(ctx, func(ctx context.Context, st *State) error { + st.setToolGenAction(key, action) + return nil + }) +} + +type reactInput struct { + Messages []Message +} + +type typedReactConfig[M MessageType] struct { + model model.BaseModel[M] + + toolsConfig *compose.ToolsNodeConfig + modelWrapperConf *typedModelWrapperConfig[M] + + toolsReturnDirectly map[string]bool + + agentName string + + maxIterations int + + cancelCtx *cancelContext + + // afterAgentFunc is called when the agent reaches a successful terminal state. + // It runs as a graph node, so compose.ProcessState is available. + afterAgentFunc func(ctx context.Context, msg M) (M, error) +} + +type reactConfig = typedReactConfig[*schema.Message] + +func genToolInfos(ctx context.Context, config *compose.ToolsNodeConfig) ([]*schema.ToolInfo, error) { + toolInfos := make([]*schema.ToolInfo, 0, len(config.Tools)) + for _, t := range config.Tools { + tl, err := t.Info(ctx) + if err != nil { + return nil, err + } + + toolInfos = append(toolInfos, tl) + } + + return toolInfos, nil +} + +type reactGraph = *compose.Graph[*reactInput, Message] + +func getReturnDirectlyToolCallID(ctx context.Context) (string, bool) { + var toolCallID string + handler := func(_ context.Context, st *State) error { + toolCallID = st.getReturnDirectlyToolCallID() + return nil + } + + _ = compose.ProcessState(ctx, handler) + + return toolCallID, toolCallID != "" +} + +func genReactState(config *reactConfig) func(ctx context.Context) *State { + return func(ctx context.Context) *State { + st := &State{ + AgentName: config.agentName, + } + maxIter := 20 + if config.maxIterations > 0 { + maxIter = config.maxIterations + } + st.setRemainingIterations(maxIter) + return st + } +} + +func newReact(ctx context.Context, config *reactConfig) (reactGraph, error) { + const ( + initNode_ = "Init" + chatModel_ = "ChatModel" + cancelCheckNode_ = "CancelCheck" + toolNode_ = "ToolNode" + afterToolCallsNode_ = "AfterToolCalls" + afterToolCallsCancelCheckNode_ = "AfterToolCallsCancelCheck" + afterAgentNode_ = "AfterAgent" + ) + + cancelCtx := config.cancelCtx + g := compose.NewGraph[*reactInput, Message](compose.WithGenLocalState(genReactState(config))) + _ = g.AddLambdaNode(initNode_, compose.InvokableLambda(func(ctx context.Context, input *reactInput) ([]Message, error) { + _ = compose.ProcessState(ctx, func(_ context.Context, st *State) error { + st.Messages = append(st.Messages, input.Messages...) + return nil + }) + return input.Messages, nil + }), compose.WithNodeName(initNode_)) + + var wrappedModel = config.model + if config.modelWrapperConf != nil { + wrappedModel = buildModelWrappers(config.model, config.modelWrapperConf) + } + + toolsConfig := config.toolsConfig + + toolsNode, err := compose.NewToolNode(ctx, toolsConfig) + if err != nil { + return nil, err + } + + _ = g.AddChatModelNode(chatModel_, wrappedModel, compose.WithStatePreHandler( + func(ctx context.Context, input []Message, st *State) ([]Message, error) { + if st.getRemainingIterations() <= 0 { + return nil, ErrExceedMaxIterations + } + st.decrementRemainingIterations() + return input, nil + }), compose.WithNodeName(chatModel_)) + + // CancelAfterChatModel safe-point: on the tool-calls path, after the branch + // has confirmed that the model response contains tool calls (i.e. not a final + // answer). Skipped entirely when the model produces a final answer. + _ = g.AddLambdaNode(cancelCheckNode_, compose.InvokableLambda(func(ctx context.Context, msg Message) (Message, error) { + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode()&CancelAfterChatModel != 0 { + return nil, compose.StatefulInterrupt(ctx, "CancelAfterChatModel", msg) + } + } + wasInterrupted, hasState, state := compose.GetInterruptState[Message](ctx) + if wasInterrupted && hasState { + msg = state + } + return msg, nil + }), compose.WithNodeName(cancelCheckNode_)) + + toolPreHandle := func(ctx context.Context, _ Message, st *State) (Message, error) { + input := st.Messages[len(st.Messages)-1] + returnDirectly := config.toolsReturnDirectly + if execCtx := getTypedChatModelAgentExecCtx[*schema.Message](ctx); execCtx != nil && len(execCtx.runtimeReturnDirectly) > 0 { + returnDirectly = execCtx.runtimeReturnDirectly + } + if len(returnDirectly) > 0 { + for i := range input.ToolCalls { + toolName := input.ToolCalls[i].Function.Name + if _, ok := returnDirectly[toolName]; ok { + st.setReturnDirectlyToolCallID(input.ToolCalls[i].ID) + } + } + } + return input, nil + } + toolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.Message], st *State) (*schema.StreamReader[[]*schema.Message], error) { + if event := st.getReturnDirectlyEvent(); event != nil { + getTypedChatModelAgentExecCtx[*schema.Message](ctx).send(event) + st.setReturnDirectlyEvent(nil) + } + return out, nil + } + _ = g.AddToolsNode(toolNode_, toolsNode, + compose.WithStatePreHandler(toolPreHandle), + compose.WithStreamStatePostHandler(toolPostHandle), + compose.WithNodeName(toolNode_)) + + // AfterToolCalls node: persists tool results to state and fires the after-tool-calls hook. + // The graph auto-materializes the ToolsNode stream into []Message before this node. + afterToolCalls := func(ctx context.Context, toolResults []Message) ([]Message, error) { + // Propagate tool message IDs from event sender to state messages. + // The event sender pre-generated IDs and stored them in state.ToolMsgIDs[toolName+callID]. + // Here we pop them and set them on the compose-created tool result messages + // so that state messages share the same IDs as their corresponding event messages. + // If no stored ID is found (old checkpoint, custom event sender), generate a fresh one. + _ = compose.ProcessState(ctx, func(_ context.Context, st *State) error { + for _, msg := range toolResults { + if id := st.popToolMsgID(msg.ToolName, msg.ToolCallID); id != "" { + msg.Extra = internal.SetMessageID(msg.Extra, id) + } else { + msg.Extra = internal.EnsureMessageID(msg.Extra) + } + st.Messages = append(st.Messages, msg) + } + return nil + }) + + execCtx := getTypedChatModelAgentExecCtx[Message](ctx) + if execCtx != nil && execCtx.afterToolCallsHook != nil { + if err := execCtx.afterToolCallsHook(ctx); err != nil { + return nil, err + } + } + + return toolResults, nil + } + _ = g.AddLambdaNode(afterToolCallsNode_, compose.InvokableLambda(afterToolCalls), + compose.WithNodeName(afterToolCallsNode_)) + + // AfterToolCallsCancelCheck: CancelAfterToolCalls safe-point, separated from toolPostHandle. + afterToolCallsCancelCheck := func(ctx context.Context, toolResults []Message) ([]Message, error) { + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode()&CancelAfterToolCalls != 0 { + return nil, compose.Interrupt(ctx, "CancelAfterToolCalls") + } + } + return toolResults, nil + } + _ = g.AddLambdaNode(afterToolCallsCancelCheckNode_, compose.InvokableLambda(afterToolCallsCancelCheck), + compose.WithNodeName(afterToolCallsCancelCheckNode_)) + + _ = g.AddEdge(compose.START, initNode_) + _ = g.AddEdge(initNode_, chatModel_) + + // Determine the terminal node: afterAgentNode_ if afterAgentFunc is set, otherwise compose.END. + terminalNode := compose.END + if config.afterAgentFunc != nil { + _ = g.AddLambdaNode(afterAgentNode_, compose.InvokableLambda(config.afterAgentFunc), + compose.WithNodeName(afterAgentNode_)) + _ = g.AddEdge(afterAgentNode_, compose.END) + terminalNode = afterAgentNode_ + } + + toolCallCheck := func(ctx context.Context, sMsg MessageStream) (string, error) { + defer sMsg.Close() + for { + chunk, err_ := sMsg.Recv() + if err_ != nil { + if err_ == io.EOF { + return terminalNode, nil + } + + return "", err_ + } + + if len(chunk.ToolCalls) > 0 { + return cancelCheckNode_, nil + } + } + } + branch := compose.NewStreamGraphBranch(toolCallCheck, map[string]bool{terminalNode: true, cancelCheckNode_: true}) + _ = g.AddBranch(chatModel_, branch) + + _ = g.AddEdge(cancelCheckNode_, toolNode_) + _ = g.AddEdge(toolNode_, afterToolCallsNode_) + _ = g.AddEdge(afterToolCallsNode_, afterToolCallsCancelCheckNode_) + + if len(config.toolsReturnDirectly) > 0 { + const ( + toolNodeToEndConverter = "ToolNodeToEndConverter" + ) + + cvt := func(ctx context.Context, toolResults []Message) (Message, error) { + id, _ := getReturnDirectlyToolCallID(ctx) + + for _, msg := range toolResults { + if msg != nil && msg.ToolCallID == id { + return msg, nil + } + } + + return nil, errors.New("return directly tool call result not found") + } + + _ = g.AddLambdaNode(toolNodeToEndConverter, compose.InvokableLambda(cvt), + compose.WithNodeName(toolNodeToEndConverter)) + _ = g.AddEdge(toolNodeToEndConverter, terminalNode) + + checkReturnDirect := func(ctx context.Context, toolResults []Message) (string, error) { + _, ok := getReturnDirectlyToolCallID(ctx) + + if ok { + return toolNodeToEndConverter, nil + } + + return chatModel_, nil + } + + returnDirectBranch := compose.NewGraphBranch(checkReturnDirect, + map[string]bool{toolNodeToEndConverter: true, chatModel_: true}) + _ = g.AddBranch(afterToolCallsCancelCheckNode_, returnDirectBranch) + } else { + _ = g.AddEdge(afterToolCallsCancelCheckNode_, chatModel_) + } + + return g, nil +} + +type agenticReactInput struct { + Messages []*schema.AgenticMessage +} + +type agenticReactConfig = typedReactConfig[*schema.AgenticMessage] + +type agenticReactGraph = *compose.Graph[*agenticReactInput, *schema.AgenticMessage] + +func getAgenticReturnDirectlyToolCallID(ctx context.Context) (string, bool) { + var toolCallID string + _ = compose.ProcessState(ctx, func(_ context.Context, st *agenticState) error { + toolCallID = st.getReturnDirectlyToolCallID() + return nil + }) + return toolCallID, toolCallID != "" +} + +func genAgenticReactState(config *agenticReactConfig) func(ctx context.Context) *agenticState { + return func(ctx context.Context) *agenticState { + st := &agenticState{ + AgentName: config.agentName, + } + maxIter := 20 + if config.maxIterations > 0 { + maxIter = config.maxIterations + } + st.setRemainingIterations(maxIter) + return st + } +} + +func agenticMessageHasToolCalls(msg *schema.AgenticMessage) bool { + if msg == nil { + return false + } + for _, block := range msg.ContentBlocks { + if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil { + return true + } + } + return false +} + +func newAgenticReact(ctx context.Context, config *agenticReactConfig) (agenticReactGraph, error) { + const ( + initNode_ = "Init" + chatModel_ = "ChatModel" + cancelCheckNode_ = "CancelCheck" + toolNode_ = "ToolNode" + afterToolCallsNode_ = "AfterToolCalls" + afterToolCallsCancelCheckNode_ = "AfterToolCallsCancelCheck" + afterAgentNode_ = "AfterAgent" + ) + + cancelCtx := config.cancelCtx + g := compose.NewGraph[*agenticReactInput, *schema.AgenticMessage]( + compose.WithGenLocalState(genAgenticReactState(config))) + _ = g.AddLambdaNode(initNode_, compose.InvokableLambda(func(ctx context.Context, input *agenticReactInput) ([]*schema.AgenticMessage, error) { + _ = compose.ProcessState(ctx, func(_ context.Context, st *agenticState) error { + st.Messages = append(st.Messages, input.Messages...) + return nil + }) + return input.Messages, nil + }), compose.WithNodeName(initNode_)) + + var wrappedModel = config.model + if config.modelWrapperConf != nil { + wrappedModel = buildModelWrappers(config.model, config.modelWrapperConf) + } + + toolsNode, err := compose.NewAgenticToolsNode(ctx, config.toolsConfig) + if err != nil { + return nil, err + } + + _ = g.AddAgenticModelNode(chatModel_, wrappedModel, compose.WithStatePreHandler( + func(ctx context.Context, input []*schema.AgenticMessage, st *agenticState) ([]*schema.AgenticMessage, error) { + if st.getRemainingIterations() <= 0 { + return nil, ErrExceedMaxIterations + } + st.decrementRemainingIterations() + return input, nil + }), compose.WithNodeName(chatModel_)) + + _ = g.AddLambdaNode(cancelCheckNode_, compose.InvokableLambda(func(ctx context.Context, msg *schema.AgenticMessage) (*schema.AgenticMessage, error) { + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode()&CancelAfterChatModel != 0 { + return nil, compose.StatefulInterrupt(ctx, "CancelAfterChatModel", msg) + } + } + wasInterrupted, hasState, state := compose.GetInterruptState[*schema.AgenticMessage](ctx) + if wasInterrupted && hasState { + msg = state + } + return msg, nil + }), compose.WithNodeName(cancelCheckNode_)) + + toolPreHandle := func(ctx context.Context, _ *schema.AgenticMessage, st *agenticState) (*schema.AgenticMessage, error) { + input := st.Messages[len(st.Messages)-1] + returnDirectly := config.toolsReturnDirectly + if execCtx := getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx); execCtx != nil && len(execCtx.runtimeReturnDirectly) > 0 { + returnDirectly = execCtx.runtimeReturnDirectly + } + if len(returnDirectly) > 0 { + for _, block := range input.ContentBlocks { + if block == nil || block.Type != schema.ContentBlockTypeFunctionToolCall || block.FunctionToolCall == nil { + continue + } + if _, ok := returnDirectly[block.FunctionToolCall.Name]; ok { + st.setReturnDirectlyToolCallID(block.FunctionToolCall.CallID) + } + } + } + return input, nil + } + toolPostHandle := func(ctx context.Context, out *schema.StreamReader[[]*schema.AgenticMessage], st *agenticState) (*schema.StreamReader[[]*schema.AgenticMessage], error) { + if event := st.getReturnDirectlyEvent(); event != nil { + getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx).send(event) + st.setReturnDirectlyEvent(nil) + } + return out, nil + } + _ = g.AddAgenticToolsNode(toolNode_, toolsNode, + compose.WithStatePreHandler(toolPreHandle), + compose.WithStreamStatePostHandler(toolPostHandle), + compose.WithNodeName(toolNode_)) + + afterToolCalls := func(ctx context.Context, toolResults []*schema.AgenticMessage) ([]*schema.AgenticMessage, error) { + _ = compose.ProcessState(ctx, func(_ context.Context, st *agenticState) error { + for _, msg := range toolResults { + if msg == nil { + continue + } + toolName, callID := extractToolIdentifiers(msg) + if id := st.popToolMsgID(toolName, callID); id != "" { + msg.Extra = internal.SetMessageID(msg.Extra, id) + } else { + msg.Extra = internal.EnsureMessageID(msg.Extra) + } + st.Messages = append(st.Messages, msg) + } + return nil + }) + + execCtx := getTypedChatModelAgentExecCtx[*schema.AgenticMessage](ctx) + if execCtx != nil && execCtx.afterToolCallsHook != nil { + if err := execCtx.afterToolCallsHook(ctx); err != nil { + return nil, err + } + } + + return toolResults, nil + } + _ = g.AddLambdaNode(afterToolCallsNode_, compose.InvokableLambda(afterToolCalls), + compose.WithNodeName(afterToolCallsNode_)) + + afterToolCallsCancelCheck := func(ctx context.Context, toolResults []*schema.AgenticMessage) ([]*schema.AgenticMessage, error) { + if cancelCtx != nil && cancelCtx.shouldCancel() { + if cancelCtx.getMode()&CancelAfterToolCalls != 0 { + return nil, compose.Interrupt(ctx, "CancelAfterToolCalls") + } + } + return toolResults, nil + } + _ = g.AddLambdaNode(afterToolCallsCancelCheckNode_, compose.InvokableLambda(afterToolCallsCancelCheck), + compose.WithNodeName(afterToolCallsCancelCheckNode_)) + + _ = g.AddEdge(compose.START, initNode_) + _ = g.AddEdge(initNode_, chatModel_) + + // Determine the terminal node: afterAgentNode_ if afterAgentFunc is set, otherwise compose.END. + terminalNode := compose.END + if config.afterAgentFunc != nil { + _ = g.AddLambdaNode(afterAgentNode_, compose.InvokableLambda(config.afterAgentFunc), + compose.WithNodeName(afterAgentNode_)) + _ = g.AddEdge(afterAgentNode_, compose.END) + terminalNode = afterAgentNode_ + } + + toolCallCheck := func(ctx context.Context, sMsg *schema.StreamReader[*schema.AgenticMessage]) (string, error) { + defer sMsg.Close() + for { + chunk, err_ := sMsg.Recv() + if err_ != nil { + if err_ == io.EOF { + return terminalNode, nil + } + return "", err_ + } + if agenticMessageHasToolCalls(chunk) { + return cancelCheckNode_, nil + } + } + } + branch := compose.NewStreamGraphBranch(toolCallCheck, map[string]bool{terminalNode: true, cancelCheckNode_: true}) + _ = g.AddBranch(chatModel_, branch) + + _ = g.AddEdge(cancelCheckNode_, toolNode_) + _ = g.AddEdge(toolNode_, afterToolCallsNode_) + _ = g.AddEdge(afterToolCallsNode_, afterToolCallsCancelCheckNode_) + + if len(config.toolsReturnDirectly) > 0 { + const ( + toolNodeToEndConverter = "ToolNodeToEndConverter" + ) + + cvt := func(ctx context.Context, toolResults []*schema.AgenticMessage) (*schema.AgenticMessage, error) { + id, _ := getAgenticReturnDirectlyToolCallID(ctx) + for _, msg := range toolResults { + if msg == nil { + continue + } + _, callID := extractToolIdentifiers(msg) + if callID == id { + return msg, nil + } + } + return nil, errors.New("return directly tool call result not found") + } + + _ = g.AddLambdaNode(toolNodeToEndConverter, compose.InvokableLambda(cvt), + compose.WithNodeName(toolNodeToEndConverter)) + _ = g.AddEdge(toolNodeToEndConverter, terminalNode) + + checkReturnDirect := func(ctx context.Context, toolResults []*schema.AgenticMessage) (string, error) { + _, ok := getAgenticReturnDirectlyToolCallID(ctx) + if ok { + return toolNodeToEndConverter, nil + } + return chatModel_, nil + } + + returnDirectBranch := compose.NewGraphBranch(checkReturnDirect, + map[string]bool{toolNodeToEndConverter: true, chatModel_: true}) + _ = g.AddBranch(afterToolCallsCancelCheckNode_, returnDirectBranch) + } else { + _ = g.AddEdge(afterToolCallsCancelCheckNode_, chatModel_) + } + + return g, nil +} + +// extractToolIdentifiers extracts the tool name and call ID from an AgenticMessage +// that contains a tool result content block. +// Assumes one tool result per message, which is guaranteed by AgenticToolsNode +// (see compose.toolMessageToAgenticMessage). +func extractToolIdentifiers(msg *schema.AgenticMessage) (toolName, callID string) { + if msg == nil { + return "", "" + } + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil { + return block.FunctionToolResult.Name, block.FunctionToolResult.CallID + } + if block.Type == schema.ContentBlockTypeToolSearchResult && block.ToolSearchFunctionToolResult != nil { + return block.ToolSearchFunctionToolResult.Name, block.ToolSearchFunctionToolResult.CallID + } + } + return "", "" +} diff --git a/adk/react_test.go b/adk/react_test.go new file mode 100644 index 0000000..b3cdb1c --- /dev/null +++ b/adk/react_test.go @@ -0,0 +1,672 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "fmt" + "io" + "math" + "math/rand" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type testModelWrapper struct { + inner model.ToolCallingChatModel +} + +func TestStateCompatConversions_V080(t *testing.T) { + t.Run("stateV080GobDecodeAndToState", func(t *testing.T) { + ss := &stateV080Serialization{ + ReturnDirectlyToolCallID: "tcid", + RemainingIterations: 2, + Internals: map[string]any{ + "_retryAttempt": 9, + "_returnDirectlyEvent": &AgentEvent{AgentName: "agent"}, + }, + } + + var buf bytes.Buffer + assert.NoError(t, gob.NewEncoder(&buf).Encode(ss)) + + var legacy stateV080 + assert.NoError(t, legacy.GobDecode(buf.Bytes())) + + s := stateV080ToState(&legacy) + assert.Equal(t, "tcid", s.ReturnDirectlyToolCallID) + assert.True(t, s.HasReturnDirectly) + assert.Equal(t, 2, s.RemainingIterations) + assert.Equal(t, 9, s.RetryAttempt) + assert.NotNil(t, s.ReturnDirectlyEvent) + assert.Equal(t, "agent", s.ReturnDirectlyEvent.AgentName) + }) +} + +func TestStateGetToolGenActions(t *testing.T) { + st := &State{ + ToolGenActions: map[string]*AgentAction{ + "k": {}, + }, + } + assert.NotNil(t, st.getToolGenActions()) + assert.Contains(t, st.getToolGenActions(), "k") +} + +func (w *testModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return (&stateModelWrapper{inner: w.inner, original: w.inner}).Generate(ctx, input, opts...) +} + +func (w *testModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return (&stateModelWrapper{inner: w.inner, original: w.inner}).Stream(ctx, input, opts...) +} + +func (w *testModelWrapper) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + newInner, err := w.inner.WithTools(tools) + if err != nil { + return nil, err + } + return &testModelWrapper{inner: newInner}, nil +} + +// TestReact tests the newReact function with different scenarios +func TestReact(t *testing.T) { + // Basic test for newReact function + t.Run("Invoke", func(t *testing.T) { + ctx := context.Background() + + // Create a fake tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 3, + } + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (Message, error) { + times++ + if times <= 2 { + return schema.AssistantMessage("hello test", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStrForTest()), + }, + }, + }), + nil + } + + return schema.AssistantMessage("bye", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create a reactConfig + config := &reactConfig{ + model: &testModelWrapper{inner: cm}, + toolsConfig: &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + toolsReturnDirectly: map[string]bool{}, + } + + graph, err := newReact(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, graph) + + compiled, err := graph.Compile(ctx, compose.WithMaxRunSteps(math.MaxInt)) + assert.NoError(t, err) + assert.NotNil(t, compiled) + + // Test with a user message + result, err := compiled.Invoke(ctx, &reactInput{Messages: []Message{ + { + Role: schema.User, + Content: "Use the test tool to say hello", + }, + }}) + assert.NoError(t, err) + assert.NotNil(t, result) + }) + + // Test with toolsReturnDirectly + t.Run("ToolsReturnDirectly", func(t *testing.T) { + ctx := context.Background() + + // Create a fake tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 3, + } + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (Message, error) { + times++ + if times <= 2 { + return schema.AssistantMessage("hello test", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStrForTest()), + }, + }, + }), + nil + } + + return schema.AssistantMessage("bye", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create a reactConfig with toolsReturnDirectly + config := &reactConfig{ + model: &testModelWrapper{inner: cm}, + toolsConfig: &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + toolsReturnDirectly: map[string]bool{info.Name: true}, + } + + graph, err := newReact(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, graph) + + compiled, err := graph.Compile(ctx, compose.WithMaxRunSteps(math.MaxInt)) + assert.NoError(t, err) + assert.NotNil(t, compiled) + + // Test with a user message when tool returns directly + result, err := compiled.Invoke(ctx, &reactInput{Messages: []Message{ + { + Role: schema.User, + Content: "Use the test tool to say hello", + }, + }}) + assert.NoError(t, err) + assert.NotNil(t, result) + + assert.Equal(t, result.Role, schema.Tool) + }) + + // Test streaming functionality + t.Run("Stream", func(t *testing.T) { + ctx := context.Background() + + // Create a fake tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 3, + } + + fakeStreamTool := &fakeStreamToolForTest{ + tarCount: 3, + } + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + times := 0 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) ( + MessageStream, error) { + sr, sw := schema.Pipe[Message](1) + defer sw.Close() + + info, _ := fakeTool.Info(ctx) + streamInfo, _ := fakeStreamTool.Info(ctx) + + times++ + if times <= 1 { + sw.Send(schema.AssistantMessage("hello test", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "tool"}`, randStrForTest()), + }, + }, + }), + nil) + return sr, nil + } else if times == 2 { + sw.Send(schema.AssistantMessage("hello stream", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: streamInfo.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "stream tool"}`, randStrForTest()), + }, + }, + }), + nil) + return sr, nil + } + + sw.Send(schema.AssistantMessage("bye", nil), nil) + return sr, nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create a reactConfig + config := &reactConfig{ + model: &testModelWrapper{inner: cm}, + toolsConfig: &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool, fakeStreamTool}, + }, + toolsReturnDirectly: map[string]bool{}, + } + + graph, err := newReact(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, graph) + + compiled, err := graph.Compile(ctx, compose.WithMaxRunSteps(math.MaxInt)) + assert.NoError(t, err) + assert.NotNil(t, compiled) + + // Test streaming with a user message + outStream, err := compiled.Stream(ctx, &reactInput{Messages: []Message{ + { + Role: schema.User, + Content: "Use the test tool to say hello", + }, + }}) + assert.NoError(t, err) + assert.NotNil(t, outStream) + + defer outStream.Close() + + msgs := make([]Message, 0) + for { + msg, err_ := outStream.Recv() + if err_ != nil { + if errors.Is(err_, io.EOF) { + break + } + t.Fatal(err_) + } + + msgs = append(msgs, msg) + } + + assert.NotEmpty(t, msgs) + }) + + // Test streaming with toolsReturnDirectly + t.Run("StreamWithToolsReturnDirectly", func(t *testing.T) { + ctx := context.Background() + + // Create a fake tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 3, + } + + fakeStreamTool := &fakeStreamToolForTest{ + tarCount: 3, + } + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + times := 0 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) ( + MessageStream, error) { + sr, sw := schema.Pipe[Message](1) + defer sw.Close() + + info, _ := fakeTool.Info(ctx) + streamInfo, _ := fakeStreamTool.Info(ctx) + + times++ + if times <= 1 { + sw.Send(schema.AssistantMessage("hello test", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "tool"}`, randStrForTest()), + }, + }, + }), + nil) + return sr, nil + } else if times == 2 { + sw.Send(schema.AssistantMessage("hello stream", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: streamInfo.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "stream tool"}`, randStrForTest()), + }, + }, + }), + nil) + return sr, nil + } + + sw.Send(schema.AssistantMessage("bye", nil), nil) + return sr, nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + streamInfo, err := fakeStreamTool.Info(ctx) + assert.NoError(t, err) + + // Create a reactConfig with toolsReturnDirectly + config := &reactConfig{ + model: &testModelWrapper{inner: cm}, + toolsConfig: &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool, fakeStreamTool}, + }, + toolsReturnDirectly: map[string]bool{streamInfo.Name: true}, + } + + graph, err := newReact(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, graph) + + compiled, err := graph.Compile(ctx, compose.WithMaxRunSteps(math.MaxInt)) + assert.NoError(t, err) + assert.NotNil(t, compiled) + + // Reset times counter + times = 0 + + // Test streaming with a user message when tool returns directly + outStream, err := compiled.Stream(ctx, &reactInput{Messages: []Message{ + { + Role: schema.User, + Content: "Use the test tool to say hello", + }, + }}) + assert.NoError(t, err) + assert.NotNil(t, outStream) + + msgs := make([]Message, 0) + for { + msg, err_ := outStream.Recv() + if err_ != nil { + if errors.Is(err_, io.EOF) { + break + } + t.Fatal(err) + } + + assert.Equal(t, msg.Role, schema.Tool) + + msgs = append(msgs, msg) + } + + outStream.Close() + + assert.NotEmpty(t, msgs) + }) + + t.Run("MaxIterations", func(t *testing.T) { + ctx := context.Background() + + // Create a fake tool for testing + fakeTool := &fakeToolForTest{ + tarCount: 3, + } + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Create a mock chat model + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // Set up expectations for the mock model + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []Message, opts ...model.Option) (Message, error) { + times++ + if times <= 5 { + return schema.AssistantMessage("hello test", + []schema.ToolCall{ + { + ID: randStrForTest(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStrForTest()), + }, + }, + }), + nil + } + + return schema.AssistantMessage("bye", nil), nil + }).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // don't exceed max iterations + config := &reactConfig{ + model: &testModelWrapper{inner: cm}, + toolsConfig: &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + toolsReturnDirectly: map[string]bool{}, + maxIterations: 6, + } + + graph, err := newReact(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, graph) + + compiled, err := graph.Compile(ctx, compose.WithMaxRunSteps(math.MaxInt)) + assert.NoError(t, err) + assert.NotNil(t, compiled) + + // Test with a user message + result, err := compiled.Invoke(ctx, &reactInput{Messages: []Message{ + { + Role: schema.User, + Content: "Use the test tool to say hello", + }, + }}) + assert.NoError(t, err) + assert.Equal(t, result.Content, "bye") + + // reset chat model times counter + times = 0 + // exceed max iterations + config = &reactConfig{ + model: &testModelWrapper{inner: cm}, + toolsConfig: &compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + toolsReturnDirectly: map[string]bool{}, + maxIterations: 5, + } + + graph, err = newReact(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, graph) + + compiled, err = graph.Compile(ctx, compose.WithMaxRunSteps(math.MaxInt)) + assert.NoError(t, err) + assert.NotNil(t, compiled) + + // Test with a user message + _, err = compiled.Invoke(ctx, &reactInput{Messages: []Message{ + { + Role: schema.User, + Content: "Use the test tool to say hello", + }, + }}) + assert.Error(t, err) + t.Logf("actual error: %v", err.Error()) + assert.ErrorIs(t, err, ErrExceedMaxIterations) + + assert.Contains(t, err.Error(), ErrExceedMaxIterations.Error()) + }) +} + +// Helper types and functions for testing + +type fakeStreamToolForTest struct { + tarCount int + curCount int +} + +func (t *fakeStreamToolForTest) StreamableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) ( + *schema.StreamReader[string], error) { + p := &fakeToolInputForTest{} + err := sonic.UnmarshalString(argumentsInJSON, p) + if err != nil { + return nil, err + } + + if t.curCount >= t.tarCount { + s := schema.StreamReaderFromArray([]string{`{"say": "bye"}`}) + return s, nil + } + t.curCount++ + s := schema.StreamReaderFromArray([]string{fmt.Sprintf(`{"say": "hello %v"}`, p.Name)}) + return s, nil +} + +type fakeToolForTest struct { + tarCount int + curCount int +} + +func (t *fakeToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "test_tool", + Desc: "test tool for unit testing", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Desc: "user name for testing", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (t *fakeStreamToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "test_stream_tool", + Desc: "test stream tool for unit testing", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Desc: "user name for testing", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (t *fakeToolForTest) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + p := &fakeToolInputForTest{} + err := sonic.UnmarshalString(argumentsInJSON, p) + if err != nil { + return "", err + } + + if t.curCount >= t.tarCount { + return `{"say": "bye"}`, nil + } + + t.curCount++ + return fmt.Sprintf(`{"say": "hello %v"}`, p.Name), nil +} + +type fakeToolInputForTest struct { + Name string `json:"name"` +} + +func randStrForTest() string { + seeds := []rune("test seed") + b := make([]rune, 8) + for i := range b { + b[i] = seeds[rand.Intn(len(seeds))] + } + return string(b) +} + +func TestReactHistory_EmptyMessages(t *testing.T) { + g := compose.NewGraph[string, []Message](compose.WithGenLocalState(func(ctx context.Context) (state *State) { + return &State{ + Messages: []Message{}, + } + })) + require.NoError(t, g.AddLambdaNode("1", compose.InvokableLambda(func(ctx context.Context, input string) (output []Message, err error) { + return getReactChatHistory(ctx, "DestAgent") + }))) + require.NoError(t, g.AddEdge(compose.START, "1")) + require.NoError(t, g.AddEdge("1", compose.END)) + + ctx := context.Background() + ctx, _ = initRunCtx(ctx, "MyAgent", nil) + runner, err := g.Compile(ctx) + require.NoError(t, err) + + require.NotPanics(t, func() { + result, err := runner.Invoke(ctx, "") + if err != nil { + t.Logf("Got error (acceptable): %v", err) + return + } + t.Logf("Got %d messages", len(result)) + }, "BUG: getReactChatHistory should not panic with empty Messages slice") +} diff --git a/adk/retry_chatmodel.go b/adk/retry_chatmodel.go new file mode 100644 index 0000000..350a3c4 --- /dev/null +++ b/adk/retry_chatmodel.go @@ -0,0 +1,759 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "io" + "math/rand" + "time" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +var ( + // ErrExceedMaxRetries is returned when the maximum number of retries has been exceeded. + // Use errors.Is to check if an error is due to max retries being exceeded: + // + // if errors.Is(err, adk.ErrExceedMaxRetries) { + // // handle max retries exceeded + // } + // + // Use errors.As to extract the underlying RetryExhaustedError for the last error details: + // + // var retryErr *adk.RetryExhaustedError + // if errors.As(err, &retryErr) { + // fmt.Printf("last error was: %v\n", retryErr.LastErr) + // } + ErrExceedMaxRetries = errors.New("exceeds max retries") +) + +// RetryExhaustedError is returned when all retry attempts have been exhausted. +// It wraps the last error that occurred during retry attempts. +type RetryExhaustedError struct { + LastErr error + TotalRetries int +} + +func (e *RetryExhaustedError) Error() string { + if e.LastErr != nil { + return fmt.Sprintf("exceeds max retries: last error: %v", e.LastErr) + } + return "exceeds max retries" +} + +func (e *RetryExhaustedError) Unwrap() error { + return ErrExceedMaxRetries +} + +// WillRetryError is emitted when a retryable error occurs and a retry will be attempted. +// It allows end-users to observe retry events in real-time via AgentEvent. +// +// Field design rationale: +// - ErrStr (exported): Stores the error message string for Gob serialization during checkpointing. +// This ensures the error message is preserved after checkpoint restore. +// - err (unexported): Stores the original error for Unwrap() support at runtime. +// This field is intentionally unexported because Gob serialization would fail for unregistered +// concrete error types. Since end-users only need the original error when the AgentEvent first +// occurs (not after restoring from checkpoint), skipping serialization is acceptable. +// After checkpoint restore, err will be nil and Unwrap() returns nil. +// - rejectReason (unexported): Stores a user-defined value set by the ShouldRetry callback +// via RetryDecision.RejectReason. This is runtime-only observability data — after checkpoint +// restore it will be nil. Unexported to avoid Gob serialization of arbitrary types. +type WillRetryError struct { + ErrStr string + RetryAttempt int + rejectReason any + err error +} + +func (e *WillRetryError) Error() string { + return e.ErrStr +} + +func (e *WillRetryError) Unwrap() error { + return e.err +} + +// RejectReason returns the user-defined rejection reason set by the ShouldRetry callback +// via RetryDecision.RejectReason. Returns nil if not set or after checkpoint restore. +func (e *WillRetryError) RejectReason() any { + return e.rejectReason +} + +func init() { + schema.RegisterName[*WillRetryError]("eino_adk_chatmodel_will_retry_error") +} + +// TypedRetryContext contains context information passed to TypedModelRetryConfig.ShouldRetry +// during a retry decision. +// +// State combinations for OutputMessage and Err: +// +// OutputMessage != nil, Err == nil → successful call; inspect message quality +// OutputMessage == nil, Err != nil → failed call (Generate error or Stream() error) +// OutputMessage != nil, Err != nil → partial stream (chunks received before mid-stream error) +// OutputMessage == nil, Err == nil → empty stream (zero chunks before EOF) +type TypedRetryContext[M MessageType] struct { + // RetryAttempt is the current retry attempt number (1-based). + // For the first retry decision (after the initial call), this is 1. + RetryAttempt int + + // InputMessages is the input messages that were sent to the model for the current attempt. + InputMessages []M + + // Options is the model options that were used for the current attempt. + Options []model.Option + + // OutputMessage is the output message from the model, if any. + // This is non-nil when the model returned a message successfully. + // For streaming, this is the fully concatenated message (the entire stream is consumed + // before ShouldRetry is called). + // For streaming with mid-stream errors, this is the partial concatenation of chunks + // received before the error occurred. + // May be nil if the model returned an error without producing a message, or if the + // stream was empty (zero chunks before EOF). + OutputMessage M + + // Err is the error from the model call, if any. + // May be nil if the model produced a message without error. + // Note: both OutputMessage and Err can be nil simultaneously for empty streams. + Err error +} + +// RetryContext is the default retry context type using *schema.Message. +type RetryContext = TypedRetryContext[*schema.Message] + +// TypedRetryDecision represents the decision made by TypedModelRetryConfig.ShouldRetry. +type TypedRetryDecision[M MessageType] struct { + // Retry indicates whether the model call should be retried. + // If false, the model output (or error) is accepted as-is, unless RewriteError is set. + Retry bool + + // RewriteError, when non-nil, overrides the return value of the model call with this error. + // The agent run will fail with this error. + // + // This is useful for two scenarios: + // - When the model returns a "seemingly correct" message (no error) that actually + // contains unrecoverable issues. RewriteError converts the successful output + // into a fatal error. + // - When the model returns an error, but you want to replace it with a different, + // more descriptive error (e.g., adding context or wrapping). + // + // When Retry is true, RewriteError is ignored. + // When Retry is false and RewriteError is non-nil, the model call returns + // RewriteError regardless of whether the original call had an error or a message. + RewriteError error + + // ModifiedInputMessages, when non-nil, replaces the input messages for the next retry. + // + // This enables advanced recovery strategies like context compression or message trimming. + // Only used when Retry is true. Ignored when Retry is false. + ModifiedInputMessages []M + + // PersistModifiedInputMessages controls whether ModifiedInputMessages are written + // back to the agent's conversation history, affecting subsequent model calls in + // the agent loop (not just the next retry attempt). + // + // When true, the modified messages replace the current conversation history. + // When false (default), the modified messages are only used for the next retry attempt + // within this retry cycle. + // + // Only used when Retry is true and ModifiedInputMessages is non-nil. + PersistModifiedInputMessages bool + + // AdditionalOptions, when non-nil, provides additional model options for the next retry. + // These options are appended to the existing options, taking precedence via last-wins semantics. + // + // This enables adjustments like increasing MaxTokens for the retry attempt. + // Note: options accumulate across retries within a single retry cycle. If ShouldRetry + // returns AdditionalOptions on every attempt, each set is appended to the previous ones. + // Only the last value for each option key takes effect, but earlier values remain in the slice. + // AdditionalOptions are scoped to the current retry cycle and do not persist to subsequent + // agent iterations — each new model call in the agent loop starts with the original options. + // Only used when Retry is true. Ignored when Retry is false. + AdditionalOptions []model.Option + + // Backoff specifies the duration to wait before the next retry attempt. + // If zero, the default backoff function (from ModelRetryConfig.BackoffFunc or the + // built-in exponential backoff) is used. + // + // This allows the ShouldRetry callback to dynamically control retry timing based on + // the specific error or problematic message encountered. + // Only used when Retry is true. Ignored when Retry is false. + Backoff time.Duration + + // RejectReason is an optional user-defined value describing why the output was rejected. + // When Retry is true and the rejected stream/message is observed downstream via + // AgentEvent, this value is attached to the WillRetryError emitted to the event stream. + // Consumers can retrieve it via WillRetryError.RejectReason(). + // + // The ShouldRetry callback has full access to the model output (via retryCtx.OutputMessage) + // and error (via retryCtx.Err), so it can distill whatever information it wants into + // RejectReason — a string, a struct, the output message itself, or nil. + // + // Only used when Retry is true. Ignored when Retry is false. + RejectReason any +} + +// RetryDecision is the default retry decision type using *schema.Message. +type RetryDecision = TypedRetryDecision[*schema.Message] + +// TypedModelRetryConfig configures retry behavior for the ChatModel node. +// It defines how the agent should handle transient failures when calling the ChatModel. +type TypedModelRetryConfig[M MessageType] struct { + // MaxRetries specifies the maximum number of retry attempts. + // A value of 0 means no retries will be attempted. + // A value of 3 means up to 3 retry attempts (4 total calls including the initial attempt). + MaxRetries int + + // ShouldRetry determines how to handle a model call result. + // It receives context information about the current attempt including the output message + // and/or error, and returns a decision on whether to retry, what to modify, etc. + // Returning nil is treated as &RetryDecision{Retry: false} (accept as-is). + // + // If nil, defaults to retrying on any non-nil error (backward compatible with IsRetryAble). + // + // Note: When ShouldRetry is set, IsRetryAble is ignored. + // Note: In streaming mode, the entire stream is consumed before ShouldRetry is called. + // The event stream is sent to the client in real time regardless; only the retry + // decision is deferred until the full response is available. + ShouldRetry func(ctx context.Context, retryCtx *TypedRetryContext[M]) *TypedRetryDecision[M] + + // Deprecated: Use ShouldRetry instead for richer retry control including message + // inspection, input modification, and option adjustment. When ShouldRetry is set, + // IsRetryAble is ignored. + IsRetryAble func(ctx context.Context, err error) bool + + // BackoffFunc calculates the delay before the next retry attempt. + // The attempt parameter starts at 1 for the first retry. + // Used as the default when RetryDecision.Backoff is zero. + // If nil, a default exponential backoff with jitter is used: + // base delay 100ms, exponentially increasing up to 10s max, + // with random jitter (0-50% of delay) to prevent thundering herd. + BackoffFunc func(ctx context.Context, attempt int) time.Duration +} + +// ModelRetryConfig is the default retry config type using *schema.Message. +type ModelRetryConfig = TypedModelRetryConfig[*schema.Message] + +func defaultIsRetryAble(_ context.Context, err error) bool { + return err != nil +} + +func defaultBackoff(_ context.Context, attempt int) time.Duration { + baseDelay := 100 * time.Millisecond + maxDelay := 10 * time.Second + + if attempt <= 0 { + return baseDelay + } + + if attempt > 7 { + return maxDelay + time.Duration(rand.Int63n(int64(maxDelay/2))) + } + + delay := baseDelay * time.Duration(1< maxDelay { + delay = maxDelay + } + + jitter := time.Duration(rand.Int63n(int64(delay / 2))) + return delay + jitter +} + +func genErrWrapper(ctx context.Context, maxRetries, attempt int, isRetryAbleFunc func(ctx context.Context, err error) bool) func(error) error { + return func(err error) error { + isRetryAble := isRetryAbleFunc == nil || isRetryAbleFunc(ctx, err) + hasRetriesLeft := attempt < maxRetries + + if isRetryAble && hasRetriesLeft { + return &WillRetryError{ErrStr: err.Error(), RetryAttempt: attempt, err: err} + } + return err + } +} + +func consumeStreamForError[M any](stream *schema.StreamReader[M]) error { + defer stream.Close() + for { + _, err := stream.Recv() + if err == io.EOF { + return nil + } + if err != nil { + return err + } + } +} + +type retryVerdictSignal struct { + ch chan retryVerdict +} + +type retryVerdict struct { + WillRetry bool + RetryAttempt int + Err error + RejectReason any +} + +// retryModelWrapper wraps a BaseChatModel with retry logic. +// This is used inside the model wrapper chain, positioned between eventSenderModelWrapper +// and stateModelWrapper, so that retry only affects the inner chain (event sending, user wrappers, +// callback injection) without re-running state management (BeforeModelRewriteState/AfterModelRewriteState). +type typedRetryModelWrapper[M MessageType] struct { + inner model.BaseModel[M] + config *TypedModelRetryConfig[M] +} + +func newTypedRetryModelWrapper[M MessageType](inner model.BaseModel[M], config *TypedModelRetryConfig[M]) *typedRetryModelWrapper[M] { + return &typedRetryModelWrapper[M]{inner: inner, config: config} +} + +func (r *typedRetryModelWrapper[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + if r.config.ShouldRetry != nil { + return generateWithShouldRetry(r, ctx, input, opts...) + } + return r.generateLegacy(ctx, input, opts...) +} + +func (r *typedRetryModelWrapper[M]) generateLegacy(ctx context.Context, input []M, opts ...model.Option) (zero M, _ error) { + isRetryAble := r.config.IsRetryAble + if isRetryAble == nil { + isRetryAble = defaultIsRetryAble + } + backoffFunc := r.config.BackoffFunc + if backoffFunc == nil { + backoffFunc = defaultBackoff + } + + var lastErr error + for attempt := 0; attempt <= r.config.MaxRetries; attempt++ { + out, err := r.inner.Generate(ctx, input, opts...) + if err == nil { + return out, nil + } + + if _, ok := compose.ExtractInterruptInfo(err); ok { + return zero, err + } + + if errors.Is(err, ErrStreamCanceled) { + return zero, err + } + + if !isRetryAble(ctx, err) { + return zero, err + } + + lastErr = err + if attempt < r.config.MaxRetries { + if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil { + return zero, err + } + } + } + + return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} +} + +func generateWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx context.Context, input []M, opts ...model.Option) (M, error) { + backoffFunc := r.config.BackoffFunc + if backoffFunc == nil { + backoffFunc = defaultBackoff + } + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + + currentInput := input + currentOpts := opts + var lastErr error + var zero M + + defer func() { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setRetryAttempt(0) + return nil + }) + }() + + for attempt := 0; attempt <= r.config.MaxRetries; attempt++ { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setRetryAttempt(attempt) + return nil + }) + + // Suppress event sending during Generate: the ShouldRetry callback must decide whether + // to accept or reject the result before any event is emitted. If accepted, the event + // is sent explicitly below (lines after decision check). If rejected, no event leaks. + if execCtx != nil { + execCtx.suppressEventSend = true + } + out, err := r.inner.Generate(ctx, currentInput, currentOpts...) + if execCtx != nil { + execCtx.suppressEventSend = false + } + + if err != nil { + if _, ok := compose.ExtractInterruptInfo(err); ok { + return zero, err + } + + if errors.Is(err, ErrStreamCanceled) { + return zero, err + } + } + + retryCtx := &TypedRetryContext[M]{ + RetryAttempt: attempt + 1, + InputMessages: currentInput, + Options: currentOpts, + OutputMessage: out, + Err: err, + } + decision := r.config.ShouldRetry(ctx, retryCtx) + if decision == nil { + decision = &TypedRetryDecision[M]{} + } + + if !decision.Retry { + if decision.RewriteError != nil { + return zero, decision.RewriteError + } + if err != nil { + return zero, err + } + if execCtx != nil && execCtx.generator != nil && out != nil { + event := typedModelOutputEvent(out, nil) + execCtx.send(event) + } + return out, nil + } + + lastErr = err + if lastErr == nil { + lastErr = fmt.Errorf("model output rejected by ShouldRetry at attempt %d", attempt+1) + } + + if attempt >= r.config.MaxRetries { + break + } + + applyDecisionForRetry(¤tInput, ¤tOpts, ctx, decision) + + delay := decision.Backoff + if delay == 0 { + delay = backoffFunc(ctx, attempt+1) + } + + if err := r.contextAwareSleep(ctx, delay); err != nil { + return zero, err + } + } + + return zero, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} +} + +func (r *typedRetryModelWrapper[M]) contextAwareSleep(ctx context.Context, delay time.Duration) error { + if delay <= 0 { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(delay): + return nil + } +} + +func streamWithShouldRetry[M MessageType](r *typedRetryModelWrapper[M], ctx context.Context, input []M, opts ...model.Option) ( + *schema.StreamReader[M], error) { + + backoffFunc := r.config.BackoffFunc + if backoffFunc == nil { + backoffFunc = defaultBackoff + } + + defer func() { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setRetryAttempt(0) + return nil + }) + }() + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + + currentInput := input + currentOpts := opts + var lastErr error + var curSignal *retryVerdictSignal + + // Panic recovery for verdict signal: if ShouldRetry panics, the onEOF/errWrapper closures in + // buildStreamConvertOptions will block forever on signal.ch, causing a goroutine leak. This + // defer ensures a verdict is always sent, even on panic, before re-panicking. + defer func() { + if p := recover(); p != nil { + if curSignal != nil { + select { + case curSignal.ch <- retryVerdict{WillRetry: false, Err: fmt.Errorf("panic: %v", p)}: + default: + } + } + panic(p) + } + }() + + for attempt := 0; attempt <= r.config.MaxRetries; attempt++ { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setRetryAttempt(attempt) + return nil + }) + + signal := &retryVerdictSignal{ch: make(chan retryVerdict, 1)} + curSignal = signal + if execCtx != nil { + execCtx.retryVerdictSignal = signal + } + + stream, err := r.inner.Stream(ctx, currentInput, currentOpts...) + if err != nil { + // Defensive no-op: when Stream() returns an error, no stream exists, so + // eventSenderModel never creates the StreamReaderWithConvert hooks that would + // read from signal.ch. This send has no consumer — it merely fills the + // buffered(1) slot so the panic-recovery defer (select/default) won't block + // if a later panic tries to send a second verdict. The signal is discarded + // when the next iteration creates a new one. + signal.ch <- retryVerdict{WillRetry: false} + + if _, ok := compose.ExtractInterruptInfo(err); ok { + return nil, err + } + + if errors.Is(err, ErrStreamCanceled) { + return nil, err + } + + retryCtx := &TypedRetryContext[M]{ + RetryAttempt: attempt + 1, + InputMessages: currentInput, + Options: currentOpts, + Err: err, + } + decision := r.config.ShouldRetry(ctx, retryCtx) + if decision == nil { + decision = &TypedRetryDecision[M]{} + } + + if !decision.Retry { + if decision.RewriteError != nil { + return nil, decision.RewriteError + } + return nil, err + } + + lastErr = err + if attempt < r.config.MaxRetries { + applyDecisionForRetry(¤tInput, ¤tOpts, ctx, decision) + delay := decision.Backoff + if delay == 0 { + delay = backoffFunc(ctx, attempt+1) + } + if err := r.contextAwareSleep(ctx, delay); err != nil { + return nil, err + } + } + continue + } + + // Split the stream: checkCopy is consumed synchronously here to build the complete + // message for ShouldRetry inspection; returnCopy is returned to the caller and may + // already be consumed downstream in parallel. The verdict signal bridges the two: + // once ShouldRetry decides, the signal tells returnCopy's errWrapper/onEOF whether + // to pass through normally or inject a WillRetryError. + copies := stream.Copy(2) + checkCopy := copies[0] + returnCopy := copies[1] + + msg, streamErr := typedConsumeStream(checkCopy) + + if errors.Is(streamErr, ErrStreamCanceled) { + signal.ch <- retryVerdict{WillRetry: false} + returnCopy.Close() + return nil, streamErr + } + + retryCtx := &TypedRetryContext[M]{ + RetryAttempt: attempt + 1, + InputMessages: currentInput, + Options: currentOpts, + OutputMessage: msg, + Err: streamErr, + } + decision := r.config.ShouldRetry(ctx, retryCtx) + if decision == nil { + decision = &TypedRetryDecision[M]{} + } + + if !decision.Retry { + signal.ch <- retryVerdict{WillRetry: false} + + if decision.RewriteError != nil { + returnCopy.Close() + return nil, decision.RewriteError + } + if streamErr != nil { + returnCopy.Close() + return nil, streamErr + } + return returnCopy, nil + } + + verdictErr := streamErr + if verdictErr == nil { + verdictErr = fmt.Errorf("model output rejected by ShouldRetry at attempt %d", attempt+1) + } + signal.ch <- retryVerdict{ + WillRetry: true, + RetryAttempt: attempt, + Err: verdictErr, + RejectReason: decision.RejectReason, + } + returnCopy.Close() + + lastErr = verdictErr + + if attempt < r.config.MaxRetries { + applyDecisionForRetry(¤tInput, ¤tOpts, ctx, decision) + delay := decision.Backoff + if delay == 0 { + delay = backoffFunc(ctx, attempt+1) + } + if err := r.contextAwareSleep(ctx, delay); err != nil { + return nil, err + } + } + } + + return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} +} + +func applyDecisionForRetry[M MessageType](currentInput *[]M, currentOpts *[]model.Option, ctx context.Context, decision *TypedRetryDecision[M]) { + if decision.ModifiedInputMessages != nil { + *currentInput = decision.ModifiedInputMessages + if decision.PersistModifiedInputMessages { + modifiedInput := *currentInput + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.Messages = modifiedInput + return nil + }) + } + } + + if decision.AdditionalOptions != nil { + cloned := make([]model.Option, len(*currentOpts), len(*currentOpts)+len(decision.AdditionalOptions)) + copy(cloned, *currentOpts) + *currentOpts = append(cloned, decision.AdditionalOptions...) + } +} + +func (r *typedRetryModelWrapper[M]) Stream(ctx context.Context, input []M, opts ...model.Option) ( + *schema.StreamReader[M], error) { + + if r.config.ShouldRetry != nil { + return streamWithShouldRetry(r, ctx, input, opts...) + } + return r.streamLegacy(ctx, input, opts...) +} + +func (r *typedRetryModelWrapper[M]) streamLegacy(ctx context.Context, input []M, opts ...model.Option) ( + *schema.StreamReader[M], error) { + + isRetryAble := r.config.IsRetryAble + if isRetryAble == nil { + isRetryAble = defaultIsRetryAble + } + backoffFunc := r.config.BackoffFunc + if backoffFunc == nil { + backoffFunc = defaultBackoff + } + + defer func() { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setRetryAttempt(0) + return nil + }) + }() + + var lastErr error + for attempt := 0; attempt <= r.config.MaxRetries; attempt++ { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setRetryAttempt(attempt) + return nil + }) + + stream, err := r.inner.Stream(ctx, input, opts...) + if err != nil { + if _, ok := compose.ExtractInterruptInfo(err); ok { + return nil, err + } + if errors.Is(err, ErrStreamCanceled) { + return nil, err + } + if !isRetryAble(ctx, err) { + return nil, err + } + lastErr = err + if attempt < r.config.MaxRetries { + if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil { + return nil, err + } + } + continue + } + + copies := stream.Copy(2) + checkCopy := copies[0] + returnCopy := copies[1] + + streamErr := consumeStreamForError(checkCopy) + if streamErr == nil { + return returnCopy, nil + } + + returnCopy.Close() + if errors.Is(streamErr, ErrStreamCanceled) { + return nil, streamErr + } + if !isRetryAble(ctx, streamErr) { + return nil, streamErr + } + + lastErr = streamErr + if attempt < r.config.MaxRetries { + if err := r.contextAwareSleep(ctx, backoffFunc(ctx, attempt+1)); err != nil { + return nil, err + } + } + } + + return nil, &RetryExhaustedError{LastErr: lastErr, TotalRetries: r.config.MaxRetries} +} diff --git a/adk/runctx.go b/adk/runctx.go new file mode 100644 index 0000000..dd42226 --- /dev/null +++ b/adk/runctx.go @@ -0,0 +1,567 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "fmt" + "io" + "sort" + "sync" + "time" + + "github.com/cloudwego/eino/schema" +) + +// runSession CheckpointSchema: persisted via serialization.RunCtx (gob). +type runSession struct { + Values map[string]any + valuesMtx *sync.Mutex + + Events []*agentEventWrapper + LaneEvents *laneEvents + mtx sync.Mutex + + // TypedEvents stores *[]*typedAgentEventWrapper[M] for M != *schema.Message. + // For M = *schema.Message, the existing Events field is used instead. + // The any type is required because Go does not support generic fields in non-generic structs. + TypedEvents any +} + +// laneEvents CheckpointSchema: persisted via serialization.RunCtx (gob). +type laneEvents struct { + Events []*agentEventWrapper + Parent *laneEvents +} + +// agentEventWrapper CheckpointSchema: persisted via serialization.RunCtx (gob). +type agentEventWrapper struct { + *AgentEvent + mu sync.Mutex + concatenatedMessage Message + // TS is the timestamp (in nanoseconds) when this event was created. + // It is primarily used by the laneEvents mechanism to order events + // from different agents in a multi-agent flow. + TS int64 + // StreamErr stores the error message if the MessageStream contained an error. + // This field guards against multiple calls to getMessageFromWrappedEvent + // when the stream has already been consumed and errored. + // Normally when StreamErr happens, the Agent will return with the error, + // unless retry is configured for the agent generating this stream, in which case + // this StreamErr will be of type WillRetryError (indicating retry is pending). + StreamErr error +} + +type typedAgentEventWrapper[M MessageType] struct { + event *TypedAgentEvent[M] + mu sync.Mutex + concatenatedMessage M + TS int64 + StreamErr error +} + +// typedAgentEventWrapperForGob is a gob-serializable representation of typedAgentEventWrapper. +// We encode the event and TS separately to avoid the sync.Mutex and non-exported fields. +type typedAgentEventWrapperForGob[M MessageType] struct { + Event *TypedAgentEvent[M] + TS int64 +} + +func (e *typedAgentEventWrapper[M]) GobEncode() ([]byte, error) { + if e.event != nil && e.event.Output != nil && e.event.Output.MessageOutput != nil && e.event.Output.MessageOutput.IsStreaming { + // Materialize the stream before encoding. + if isNilMessage(e.concatenatedMessage) && e.StreamErr == nil { + e.consumeStream() + } + } + + buf := &bytes.Buffer{} + err := gob.NewEncoder(buf).Encode(&typedAgentEventWrapperForGob[M]{ + Event: e.event, + TS: e.TS, + }) + if err != nil { + return nil, fmt.Errorf("failed to gob encode generic agent event wrapper: %w", err) + } + return buf.Bytes(), nil +} + +func (e *typedAgentEventWrapper[M]) GobDecode(b []byte) error { + g := &typedAgentEventWrapperForGob[M]{} + if err := gob.NewDecoder(bytes.NewReader(b)).Decode(g); err != nil { + return fmt.Errorf("failed to gob decode generic agent event wrapper: %w", err) + } + e.event = g.Event + e.TS = g.TS + return nil +} + +// consumeStream drains the typed message stream, setting concatenatedMessage on success +// or StreamErr on failure. The stream is replaced with a materialized version safe for +// gob encoding. +// +// NOTE: This method parallels agentEventWrapper.consumeStream in utils.go. The two +// implementations exist because agentEventWrapper is non-generic (uses *schema.Message +// directly) while typedAgentEventWrapper[M] is generic. They cannot be unified without +// making the non-generic wrapper generic, which would cascade through the entire +// non-generic event storage layer. +func (e *typedAgentEventWrapper[M]) consumeStream() { + e.mu.Lock() + defer e.mu.Unlock() + + if !isNilMessage(e.concatenatedMessage) { + return + } + + s := e.event.Output.MessageOutput.MessageStream + var msgs []M + + defer s.Close() + for { + msg, err := s.Recv() + if err != nil { + if err == io.EOF { + break + } + e.StreamErr = err + e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs) + return + } + msgs = append(msgs, msg) + } + + if len(msgs) == 0 { + e.StreamErr = errors.New("no messages in typedAgentEventWrapper.MessageStream") + e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs) + return + } + + if len(msgs) == 1 { + e.concatenatedMessage = msgs[0] + } else { + var err error + e.concatenatedMessage, err = concatMessageStream(schema.StreamReaderFromArray(msgs)) + if err != nil { + e.StreamErr = err + e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs) + return + } + } + + e.event.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray([]M{e.concatenatedMessage}) +} + +type otherAgentEventWrapperForEncode agentEventWrapper + +func (a *agentEventWrapper) GobEncode() ([]byte, error) { + if a.Output != nil && a.Output.MessageOutput != nil && a.Output.MessageOutput.IsStreaming { + // Materialize the stream before encoding. An unconsumed stream that + // ends with a non-EOF error (WillRetryError, ErrStreamCanceled) would + // cause MessageVariant.GobEncode to fail. consumeStream replaces the + // stream with an error-free, materialized version. + if a.concatenatedMessage == nil && a.StreamErr == nil { + a.consumeStream() + } + } + + buf := &bytes.Buffer{} + err := gob.NewEncoder(buf).Encode((*otherAgentEventWrapperForEncode)(a)) + if err != nil { + return nil, fmt.Errorf("failed to gob encode agent event wrapper: %w", err) + } + return buf.Bytes(), nil +} + +func (a *agentEventWrapper) GobDecode(b []byte) error { + return gob.NewDecoder(bytes.NewReader(b)).Decode((*otherAgentEventWrapperForEncode)(a)) +} + +func newRunSession() *runSession { + return &runSession{ + Values: make(map[string]any), + valuesMtx: &sync.Mutex{}, + } +} + +// GetSessionValues returns all session key-value pairs for the current run. +func GetSessionValues(ctx context.Context) map[string]any { + session := getSession(ctx) + if session == nil { + return map[string]any{} + } + + return session.getValues() +} + +// AddSessionValue sets a single session key-value pair for the current run. +func AddSessionValue(ctx context.Context, key string, value any) { + session := getSession(ctx) + if session == nil { + return + } + + session.addValue(key, value) +} + +// AddSessionValues sets multiple session key-value pairs for the current run. +func AddSessionValues(ctx context.Context, kvs map[string]any) { + session := getSession(ctx) + if session == nil { + return + } + + session.addValues(kvs) +} + +// GetSessionValue retrieves a session value by key and reports whether it exists. +func GetSessionValue(ctx context.Context, key string) (any, bool) { + session := getSession(ctx) + if session == nil { + return nil, false + } + + return session.getValue(key) +} + +func (rs *runSession) addEvent(event *AgentEvent) { + wrapper := &agentEventWrapper{AgentEvent: event, TS: time.Now().UnixNano()} + // If LaneEvents is not nil, we are in a parallel lane. + // Append to the lane's local event slice (lock-free). + if rs.LaneEvents != nil { + rs.LaneEvents.Events = append(rs.LaneEvents.Events, wrapper) + return + } + + // Otherwise, we are on the main path. Append to the shared Events slice (with lock). + rs.mtx.Lock() + rs.Events = append(rs.Events, wrapper) + rs.mtx.Unlock() +} + +func (rs *runSession) getEvents() []*agentEventWrapper { + // If there are no in-flight lane events, we can return the main slice directly. + if rs.LaneEvents == nil { + rs.mtx.Lock() + events := rs.Events + rs.mtx.Unlock() + return events + } + + // If there are in-flight events, we must construct the full view. + // First, get the committed history from the main Events slice. + rs.mtx.Lock() + committedEvents := make([]*agentEventWrapper, len(rs.Events)) + copy(committedEvents, rs.Events) + rs.mtx.Unlock() + + // Then, assemble the in-flight events by traversing the linked list. + // Reading the .Parent pointer is safe without a lock because the parent of a lane is immutable after creation. + var laneSlices [][]*agentEventWrapper + totalLaneSize := 0 + for lane := rs.LaneEvents; lane != nil; lane = lane.Parent { + if len(lane.Events) > 0 { + laneSlices = append(laneSlices, lane.Events) + totalLaneSize += len(lane.Events) + } + } + + // Combine committed and in-flight history. + finalEvents := make([]*agentEventWrapper, 0, len(committedEvents)+totalLaneSize) + finalEvents = append(finalEvents, committedEvents...) + for i := len(laneSlices) - 1; i >= 0; i-- { + finalEvents = append(finalEvents, laneSlices[i]...) + } + + return finalEvents +} + +func addTypedEvent[M MessageType](session *runSession, event *TypedAgentEvent[M]) { + var zero M + if _, ok := any(zero).(*schema.Message); ok { + session.addEvent(any(event).(*AgentEvent)) + return + } + session.mtx.Lock() + defer session.mtx.Unlock() + wrapper := &typedAgentEventWrapper[M]{event: event, TS: time.Now().UnixNano()} + store, _ := session.TypedEvents.(*[]*typedAgentEventWrapper[M]) + if store == nil { + s := make([]*typedAgentEventWrapper[M], 0) + store = &s + session.TypedEvents = store + } + *store = append(*store, wrapper) +} + +func (rs *runSession) getValues() map[string]any { + rs.valuesMtx.Lock() + values := make(map[string]any, len(rs.Values)) + for k, v := range rs.Values { + values[k] = v + } + rs.valuesMtx.Unlock() + + return values +} + +func (rs *runSession) addValue(key string, value any) { + rs.valuesMtx.Lock() + rs.Values[key] = value + rs.valuesMtx.Unlock() +} + +func (rs *runSession) addValues(kvs map[string]any) { + rs.valuesMtx.Lock() + for k, v := range kvs { + rs.Values[k] = v + } + rs.valuesMtx.Unlock() +} + +func (rs *runSession) getValue(key string) (any, bool) { + rs.valuesMtx.Lock() + value, ok := rs.Values[key] + rs.valuesMtx.Unlock() + + return value, ok +} + +type runContext struct { + RootInput *AgentInput + RunPath []RunStep + + AgenticRootInput any + + Session *runSession +} + +func (rc *runContext) isRoot() bool { + return len(rc.RunPath) == 1 +} + +func (rc *runContext) deepCopy() *runContext { + copied := &runContext{ + RootInput: rc.RootInput, + AgenticRootInput: rc.AgenticRootInput, + RunPath: make([]RunStep, len(rc.RunPath)), + Session: rc.Session, + } + + copy(copied.RunPath, rc.RunPath) + + return copied +} + +type runCtxKey struct{} + +func getRunCtx(ctx context.Context) *runContext { + runCtx, ok := ctx.Value(runCtxKey{}).(*runContext) + if !ok { + return nil + } + return runCtx +} + +func setRunCtx(ctx context.Context, runCtx *runContext) context.Context { + return context.WithValue(ctx, runCtxKey{}, runCtx) +} + +func initRunCtx(ctx context.Context, agentName string, input *AgentInput) (context.Context, *runContext) { + runCtx := getRunCtx(ctx) + if runCtx != nil { + runCtx = runCtx.deepCopy() + } else { + runCtx = &runContext{Session: newRunSession()} + } + + runCtx.RunPath = append(runCtx.RunPath, RunStep{agentName: agentName}) + if runCtx.isRoot() && input != nil { + runCtx.RootInput = input + } + + return setRunCtx(ctx, runCtx), runCtx +} + +func initTypedRunCtx[M MessageType](ctx context.Context, agentName string, input *TypedAgentInput[M]) (context.Context, *runContext) { + runCtx := getRunCtx(ctx) + if runCtx != nil { + runCtx = runCtx.deepCopy() + } else { + runCtx = &runContext{Session: newRunSession()} + } + + runCtx.RunPath = append(runCtx.RunPath, RunStep{agentName: agentName}) + if runCtx.isRoot() && input != nil { + var zero M + if _, ok := any(zero).(*schema.Message); ok { + runCtx.RootInput = any(input).(*AgentInput) + } else { + runCtx.AgenticRootInput = input + } + } + + return setRunCtx(ctx, runCtx), runCtx +} + +func joinRunCtxs(parentCtx context.Context, childCtxs ...context.Context) { + switch len(childCtxs) { + case 0: + return + case 1: + // Optimization for the common case of a single branch. + newEvents := unwindLaneEvents(childCtxs...) + commitEvents(parentCtx, newEvents) + return + } + + // 1. Collect all new events from the leaf nodes of each context's lane. + newEvents := unwindLaneEvents(childCtxs...) + + // 2. Sort the collected events by their creation timestamp for chronological order. + sort.Slice(newEvents, func(i, j int) bool { + return newEvents[i].TS < newEvents[j].TS + }) + + // 3. Commit the sorted events to the parent. + commitEvents(parentCtx, newEvents) +} + +// commitEvents appends a slice of new events to the correct parent lane or main event log. +func commitEvents(ctx context.Context, newEvents []*agentEventWrapper) { + runCtx := getRunCtx(ctx) + if runCtx == nil || runCtx.Session == nil { + // Should not happen, but handle defensively. + return + } + + // If the context we are committing to is itself a lane, append to its event slice. + if runCtx.Session.LaneEvents != nil { + runCtx.Session.LaneEvents.Events = append(runCtx.Session.LaneEvents.Events, newEvents...) + } else { + // Otherwise, commit to the main, shared Events slice with a lock. + runCtx.Session.mtx.Lock() + runCtx.Session.Events = append(runCtx.Session.Events, newEvents...) + runCtx.Session.mtx.Unlock() + } +} + +// unwindLaneEvents traverses the LaneEvents of the given contexts and collects +// all events from the leaf nodes. +func unwindLaneEvents(ctxs ...context.Context) []*agentEventWrapper { + var allNewEvents []*agentEventWrapper + for _, ctx := range ctxs { + runCtx := getRunCtx(ctx) + if runCtx != nil && runCtx.Session != nil && runCtx.Session.LaneEvents != nil { + allNewEvents = append(allNewEvents, runCtx.Session.LaneEvents.Events...) + } + } + return allNewEvents +} + +func forkRunCtx(ctx context.Context) context.Context { + parentRunCtx := getRunCtx(ctx) + if parentRunCtx == nil || parentRunCtx.Session == nil { + // Should not happen in a parallel workflow, but handle defensively. + return ctx + } + + // Create a new session for the child lane by manually copying the parent's session fields. + // This is crucial to ensure a new mutex is created and that the LaneEvents pointer is unique. + childSession := &runSession{ + Events: parentRunCtx.Session.Events, // Share the committed history + Values: parentRunCtx.Session.Values, // Share the values map + valuesMtx: parentRunCtx.Session.valuesMtx, + } + + // Fork the lane events within the new session struct. + childSession.LaneEvents = &laneEvents{ + Parent: parentRunCtx.Session.LaneEvents, + Events: make([]*agentEventWrapper, 0), + } + + // Create a new runContext for the child lane, pointing to the new session. + childRunCtx := &runContext{ + RootInput: parentRunCtx.RootInput, + RunPath: make([]RunStep, len(parentRunCtx.RunPath)), + Session: childSession, + } + copy(childRunCtx.RunPath, parentRunCtx.RunPath) + + return setRunCtx(ctx, childRunCtx) +} + +// updateRunPathOnly creates a new context with an updated RunPath, but does NOT modify the Address. +// This is used by sequential workflows to accumulate execution history for LLM context, +// without incorrectly chaining the static addresses of peer agents. +func updateRunPathOnly(ctx context.Context, agentNames ...string) context.Context { + runCtx := getRunCtx(ctx) + if runCtx == nil { + // This should not happen in a sequential workflow context, but handle defensively. + runCtx = &runContext{Session: newRunSession()} + } else { + runCtx = runCtx.deepCopy() + } + + for _, agentName := range agentNames { + runCtx.RunPath = append(runCtx.RunPath, RunStep{agentName: agentName}) + } + + return setRunCtx(ctx, runCtx) +} + +// ClearRunCtx clears the run context of the multi-agents. This is particularly useful +// when a customized agent with a multi-agents inside it is set as a subagent of another +// multi-agents. In such cases, it's not expected to pass the outside run context to the +// inside multi-agents, so this function helps isolate the contexts properly. +func ClearRunCtx(ctx context.Context) context.Context { + return context.WithValue(ctx, runCtxKey{}, nil) +} + +func ctxWithNewTypedRunCtx[M MessageType](ctx context.Context, input *TypedAgentInput[M], sharedParentSession bool) context.Context { + var session *runSession + if sharedParentSession { + if parentSession := getSession(ctx); parentSession != nil { + session = &runSession{ + Values: parentSession.Values, + valuesMtx: parentSession.valuesMtx, + } + } + } + if session == nil { + session = newRunSession() + } + var zero M + rc := &runContext{Session: session} + if _, ok := any(zero).(*schema.Message); ok { + rc.RootInput = any(input).(*AgentInput) + } else { + rc.AgenticRootInput = input + } + return setRunCtx(ctx, rc) +} + +func getSession(ctx context.Context) *runSession { + runCtx := getRunCtx(ctx) + if runCtx != nil { + return runCtx.Session + } + + return nil +} diff --git a/adk/runctx_test.go b/adk/runctx_test.go new file mode 100644 index 0000000..bef1f44 --- /dev/null +++ b/adk/runctx_test.go @@ -0,0 +1,634 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestSessionValues(t *testing.T) { + // Test Case 1: Basic AddSessionValues and GetSessionValues + t.Run("BasicSessionValues", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add values to the session + values := map[string]any{ + "key1": "value1", + "key2": 42, + "key3": true, + } + AddSessionValues(ctx, values) + + // Get all values from the session + retrievedValues := GetSessionValues(ctx) + + // Verify the values were added correctly + assert.Equal(t, "value1", retrievedValues["key1"]) + assert.Equal(t, 42, retrievedValues["key2"]) + assert.Equal(t, true, retrievedValues["key3"]) + assert.Len(t, retrievedValues, 3) + }) + + // Test Case 2: AddSessionValues with empty context + t.Run("AddSessionValuesEmptyContext", func(t *testing.T) { + ctx := context.Background() + + // Add values to a context without a run session + values := map[string]any{ + "key1": "value1", + } + AddSessionValues(ctx, values) + + // Get values should return empty map + retrievedValues := GetSessionValues(ctx) + assert.Empty(t, retrievedValues) + }) + + // Test Case 3: GetSessionValues with empty context + t.Run("GetSessionValuesEmptyContext", func(t *testing.T) { + ctx := context.Background() + + // Get values from a context without a run session + retrievedValues := GetSessionValues(ctx) + assert.Empty(t, retrievedValues) + }) + + // Test Case 4: AddSessionValues with nil values + t.Run("AddSessionValuesNilValues", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add nil values map + AddSessionValues(ctx, nil) + + // Get values should still be empty + retrievedValues := GetSessionValues(ctx) + assert.Empty(t, retrievedValues) + }) + + // Test Case 5: AddSessionValues with empty values + t.Run("AddSessionValuesEmptyValues", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add empty values map + AddSessionValues(ctx, map[string]any{}) + + // Get values should be empty + retrievedValues := GetSessionValues(ctx) + assert.Empty(t, retrievedValues) + }) + + // Test Case 6: AddSessionValues with complex data types + t.Run("AddSessionValuesComplexTypes", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add complex values to the session + values := map[string]any{ + "string": "hello world", + "int": 123, + "float": 45.67, + "bool": true, + "slice": []string{"a", "b", "c"}, + "map": map[string]int{"x": 1, "y": 2}, + "struct": struct{ Name string }{Name: "test"}, + } + AddSessionValues(ctx, values) + + // Get all values from the session + retrievedValues := GetSessionValues(ctx) + + // Verify the complex values were added correctly + assert.Equal(t, "hello world", retrievedValues["string"]) + assert.Equal(t, 123, retrievedValues["int"]) + assert.Equal(t, 45.67, retrievedValues["float"]) + assert.Equal(t, true, retrievedValues["bool"]) + assert.Equal(t, []string{"a", "b", "c"}, retrievedValues["slice"]) + assert.Equal(t, map[string]int{"x": 1, "y": 2}, retrievedValues["map"]) + assert.Equal(t, struct{ Name string }{Name: "test"}, retrievedValues["struct"]) + assert.Len(t, retrievedValues, 7) + }) + + // Test Case 7: AddSessionValues overwrites existing values + t.Run("AddSessionValuesOverwrite", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add initial values + initialValues := map[string]any{ + "key1": "initial1", + "key2": "initial2", + } + AddSessionValues(ctx, initialValues) + + // Add values that overwrite some keys + overwriteValues := map[string]any{ + "key1": "overwritten1", + "key3": "new3", + } + AddSessionValues(ctx, overwriteValues) + + // Get all values from the session + retrievedValues := GetSessionValues(ctx) + + // Verify the values were overwritten correctly + assert.Equal(t, "overwritten1", retrievedValues["key1"]) // overwritten + assert.Equal(t, "initial2", retrievedValues["key2"]) // unchanged + assert.Equal(t, "new3", retrievedValues["key3"]) // new + assert.Len(t, retrievedValues, 3) + }) + + // Test Case 8: Concurrent access to session values + t.Run("ConcurrentSessionValues", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add initial values + initialValues := map[string]any{ + "counter": 0, + } + AddSessionValues(ctx, initialValues) + + // Simulate concurrent access + done := make(chan bool) + + // Goroutine 1: Add values + go func() { + for i := 0; i < 100; i++ { + values := map[string]any{ + "goroutine1": i, + } + AddSessionValues(ctx, values) + } + done <- true + }() + + // Goroutine 2: Add different values + go func() { + for i := 0; i < 100; i++ { + values := map[string]any{ + "goroutine2": i, + } + AddSessionValues(ctx, values) + } + done <- true + }() + + // Wait for both goroutines to complete + <-done + <-done + + // Verify that both values were set (last write wins) + retrievedValues := GetSessionValues(ctx) + assert.Equal(t, 0, retrievedValues["counter"]) + assert.Equal(t, 99, retrievedValues["goroutine1"]) + assert.Equal(t, 99, retrievedValues["goroutine2"]) + }) + + // Test Case 9: GetSessionValue individual value + t.Run("GetSessionValueIndividual", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add values to the session + values := map[string]any{ + "key1": "value1", + "key2": 42, + } + AddSessionValues(ctx, values) + + // Get individual values + value1, exists1 := GetSessionValue(ctx, "key1") + value2, exists2 := GetSessionValue(ctx, "key2") + value3, exists3 := GetSessionValue(ctx, "nonexistent") + + // Verify individual values + assert.True(t, exists1) + assert.Equal(t, "value1", value1) + + assert.True(t, exists2) + assert.Equal(t, 42, value2) + + assert.False(t, exists3) + assert.Nil(t, value3) + }) + + // Test Case 10: AddSessionValue individual value + t.Run("AddSessionValueIndividual", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a run session + session := newRunSession() + runCtx := &runContext{Session: session} + ctx = setRunCtx(ctx, runCtx) + + // Add individual values + AddSessionValue(ctx, "key1", "value1") + AddSessionValue(ctx, "key2", 42) + + // Get all values + retrievedValues := GetSessionValues(ctx) + + // Verify the values were added correctly + assert.Equal(t, "value1", retrievedValues["key1"]) + assert.Equal(t, 42, retrievedValues["key2"]) + assert.Len(t, retrievedValues, 2) + }) + + // Test Case 11: AddSessionValue with empty context + t.Run("AddSessionValueEmptyContext", func(t *testing.T) { + ctx := context.Background() + + // Add individual value to a context without a run session + AddSessionValue(ctx, "key1", "value1") + + // Get individual value should return false + value, exists := GetSessionValue(ctx, "key1") + assert.False(t, exists) + assert.Nil(t, value) + + // Get all values should return empty map + retrievedValues := GetSessionValues(ctx) + assert.Empty(t, retrievedValues) + }) + + // Test Case 12: Integration with run context initialization + t.Run("IntegrationWithRunContext", func(t *testing.T) { + ctx := context.Background() + + // Initialize a run context with an agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("test input"), + }, + } + ctx, runCtx := initRunCtx(ctx, "test-agent", input) + + // Verify the run context was created + assert.NotNil(t, runCtx) + assert.NotNil(t, runCtx.Session) + + // Add values to the session + values := map[string]any{ + "integration_key": "integration_value", + } + AddSessionValues(ctx, values) + + // Get values from the session + retrievedValues := GetSessionValues(ctx) + assert.Equal(t, "integration_value", retrievedValues["integration_key"]) + + // Verify the run path was set correctly + assert.Len(t, runCtx.RunPath, 1) + assert.Equal(t, "test-agent", runCtx.RunPath[0].agentName) + }) +} + +func TestForkJoinRunCtx(t *testing.T) { + // Helper to create a named event + newEvent := func(name string) *AgentEvent { + // Add a small sleep to ensure timestamps are distinct + time.Sleep(1 * time.Millisecond) + return &AgentEvent{AgentName: name} + } + + // Helper to get event names from a slice of wrappers + getEventNames := func(wrappers []*agentEventWrapper) []string { + names := make([]string, len(wrappers)) + for i, w := range wrappers { + names[i] = w.AgentName + } + return names + } + + // 1. Setup: Create an initial runContext for the main execution path. + mainCtx, mainRunCtx := initRunCtx(context.Background(), "Main", nil) + + // 2. Run Agent A + eventA := newEvent("A") + mainRunCtx.Session.addEvent(eventA) + assert.Equal(t, []string{"A"}, getEventNames(mainRunCtx.Session.getEvents()), "After A") + + // 3. Fork for Par(B, C) + ctxB := forkRunCtx(mainCtx) + ctxC := forkRunCtx(mainCtx) + + // Assertions for Fork + runCtxB := getRunCtx(ctxB) + runCtxC := getRunCtx(ctxC) + assert.NotSame(t, mainRunCtx.Session, runCtxB.Session, "Session B should be a new struct") + assert.NotSame(t, mainRunCtx.Session, runCtxC.Session, "Session C should be a new struct") + assert.NotSame(t, runCtxB.Session, runCtxC.Session, "Sessions B and C should be different") + assert.Nil(t, mainRunCtx.Session.LaneEvents, "Main session should have no lane events yet") + assert.NotNil(t, runCtxB.Session.LaneEvents, "Session B should have lane events") + assert.NotNil(t, runCtxC.Session.LaneEvents, "Session C should have lane events") + assert.Nil(t, runCtxB.Session.LaneEvents.Parent, "Lane B's parent should be the main (nil) lane") + assert.Nil(t, runCtxC.Session.LaneEvents.Parent, "Lane C's parent should be the main (nil) lane") + + // 4. Run Agent B + eventB := newEvent("B") + runCtxB.Session.addEvent(eventB) + assert.Equal(t, []string{"A", "B"}, getEventNames(runCtxB.Session.getEvents()), "After B") + + // 5. Run Agent C (and Nested Fork for Par(D, E)) + eventC1 := newEvent("C1") + runCtxC.Session.addEvent(eventC1) + assert.Equal(t, []string{"A", "C1"}, getEventNames(runCtxC.Session.getEvents()), "After C1") + + ctxD := forkRunCtx(ctxC) + ctxE := forkRunCtx(ctxC) + + // Assertions for Nested Fork + runCtxD := getRunCtx(ctxD) + runCtxE := getRunCtx(ctxE) + assert.NotNil(t, runCtxD.Session.LaneEvents.Parent, "Lane D's parent should be Lane C") + assert.Same(t, runCtxC.Session.LaneEvents, runCtxD.Session.LaneEvents.Parent, "Lane D's parent must be Lane C's node") + assert.Same(t, runCtxC.Session.LaneEvents, runCtxE.Session.LaneEvents.Parent, "Lane E's parent must be Lane C's node") + + // 6. Run Agents D and E + eventD := newEvent("D") + runCtxD.Session.addEvent(eventD) + eventE := newEvent("E") + runCtxE.Session.addEvent(eventE) + + assert.Equal(t, []string{"A", "C1", "D"}, getEventNames(runCtxD.Session.getEvents()), "After D") + assert.Equal(t, []string{"A", "C1", "E"}, getEventNames(runCtxE.Session.getEvents()), "After E") + + // 7. Join Par(D, E) + joinRunCtxs(ctxC, ctxD, ctxE) + + // Assertions for Nested Join + // The events should now be committed to Lane C's event slice. + assert.Equal(t, []string{"A", "C1", "D", "E"}, getEventNames(runCtxC.Session.getEvents()), "After joining D and E") + + // 8. Join Par(B, C) + joinRunCtxs(mainCtx, ctxB, ctxC) + + // Assertions for Top-Level Join + // The events should now be committed to the main session's Events slice. + assert.Equal(t, []string{"A", "B", "C1", "D", "E"}, getEventNames(mainRunCtx.Session.getEvents()), "After joining B and C") + + // 9. Run Agent F + eventF := newEvent("F") + mainRunCtx.Session.addEvent(eventF) + assert.Equal(t, []string{"A", "B", "C1", "D", "E", "F"}, getEventNames(mainRunCtx.Session.getEvents()), "After F") +} + +// makeStreamingEventWrapper creates an agentEventWrapper with a streaming MessageOutput +// whose stream yields the given message then terminates with streamErr (or io.EOF if nil). +func makeStreamingEventWrapper(msg Message, streamErr error) *agentEventWrapper { + r, w := schema.Pipe[Message](2) + w.Send(msg, nil) + if streamErr != nil { + w.Send(nil, streamErr) + } + w.Close() + + return &agentEventWrapper{ + AgentEvent: &AgentEvent{ + AgentName: "test-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: r, + Role: schema.Assistant, + }, + }, + }, + } +} + +func TestGobEncodeStreamErrors(t *testing.T) { + t.Run("WillRetryError_unconsumed_stream_fails_GobEncode", func(t *testing.T) { + // An agentEventWrapper whose stream yields a message then WillRetryError. + // Without pre-consuming (no getMessageFromWrappedEvent call), GobEncode + // reaches MessageVariant.GobEncode which treats non-EOF errors as fatal. + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("partial", nil), + &WillRetryError{ErrStr: "model error", RetryAttempt: 1}, + ) + + _, err := wrapper.GobEncode() + assert.NoError(t, err, "GobEncode should handle WillRetryError streams gracefully") + }) + + t.Run("ErrStreamCanceled_unconsumed_stream_fails_GobEncode", func(t *testing.T) { + // Same scenario but with ErrStreamCanceled (*errors.errorString). + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("partial", nil), + ErrStreamCanceled, + ) + + _, err := wrapper.GobEncode() + assert.NoError(t, err, "GobEncode should handle ErrStreamCanceled streams gracefully") + }) + + t.Run("successful_stream_GobEncode_succeeds", func(t *testing.T) { + // Control: a clean stream (no error) should encode fine. + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("hello", nil), + nil, // no stream error + ) + + data, err := wrapper.GobEncode() + assert.NoError(t, err) + assert.NotEmpty(t, data) + + // Verify round-trip decode works. + decoded := &agentEventWrapper{AgentEvent: &AgentEvent{}} + err = decoded.GobDecode(data) + assert.NoError(t, err) + assert.Equal(t, "test-agent", decoded.AgentName) + }) + + t.Run("preconsumed_WillRetryError_GobEncode_succeeds", func(t *testing.T) { + // When getMessageFromWrappedEvent is called first, WillRetryError is + // cached in StreamErr and the stream is replaced with an error-free array. + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("partial", nil), + &WillRetryError{ErrStr: "model error", RetryAttempt: 1}, + ) + + _, consumeErr := getMessageFromWrappedEvent(wrapper) + assert.Error(t, consumeErr) + + data, err := wrapper.GobEncode() + assert.NoError(t, err, "GobEncode should succeed after pre-consuming WillRetryError stream") + assert.NotEmpty(t, data) + }) + + t.Run("preconsumed_ErrStreamCanceled_GobEncode_succeeds", func(t *testing.T) { + // ErrStreamCanceled is a *StreamCanceledError which IS gob-registered. + // After getMessageFromWrappedEvent, StreamErr = ErrStreamCanceled. + // Since it's registered, gob encoding succeeds. + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("partial", nil), + ErrStreamCanceled, + ) + + _, consumeErr := getMessageFromWrappedEvent(wrapper) + assert.Error(t, consumeErr) + + data, err := wrapper.GobEncode() + assert.NoError(t, err, "GobEncode should succeed; ErrStreamCanceled is gob-registered") + assert.NotEmpty(t, data) + }) + + t.Run("GobEncode_roundtrip_preserves_content", func(t *testing.T) { + // Verify that after GobEncode with a WillRetryError stream, + // the decoded wrapper has the partial message content and StreamErr intact. + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("partial response", nil), + &WillRetryError{ErrStr: "err", RetryAttempt: 1}, + ) + + data, err := wrapper.GobEncode() + assert.NoError(t, err) + + decoded := &agentEventWrapper{AgentEvent: &AgentEvent{}} + err = decoded.GobDecode(data) + assert.NoError(t, err) + assert.Equal(t, "test-agent", decoded.AgentName) + assert.True(t, decoded.Output.MessageOutput.IsStreaming) + // The stream should be consumable and yield the partial message. + msg, recvErr := decoded.Output.MessageOutput.MessageStream.Recv() + assert.NoError(t, recvErr) + assert.Contains(t, msg.Content, "partial response") + // StreamErr should be preserved for end-user visibility. + var willRetryErr *WillRetryError + assert.True(t, errors.As(decoded.StreamErr, &willRetryErr)) + assert.Equal(t, "err", willRetryErr.ErrStr) + }) + + t.Run("GobEncode_roundtrip_preserves_ErrStreamCanceled", func(t *testing.T) { + // ErrStreamCanceled (*StreamCanceledError) is gob-registered, so + // StreamErr should survive encoding/decoding. + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("partial", nil), + ErrStreamCanceled, + ) + + data, err := wrapper.GobEncode() + assert.NoError(t, err) + + decoded := &agentEventWrapper{AgentEvent: &AgentEvent{}} + err = decoded.GobDecode(data) + assert.NoError(t, err) + var streamCanceledErr *StreamCanceledError + assert.ErrorAs(t, decoded.StreamErr, &streamCanceledErr) + }) + + t.Run("GobEncode_idempotent", func(t *testing.T) { + // Calling GobEncode twice should succeed both times (stream replaced on first call). + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("hello", nil), + &WillRetryError{ErrStr: "err", RetryAttempt: 1}, + ) + + data1, err := wrapper.GobEncode() + assert.NoError(t, err) + + data2, err := wrapper.GobEncode() + assert.NoError(t, err) + + // Both should decode to equivalent content. + d1, d2 := &agentEventWrapper{AgentEvent: &AgentEvent{}}, &agentEventWrapper{AgentEvent: &AgentEvent{}} + assert.NoError(t, d1.GobDecode(data1)) + assert.NoError(t, d2.GobDecode(data2)) + assert.Equal(t, d1.AgentName, d2.AgentName) + }) + + t.Run("GobEncode_non_streaming_unaffected", func(t *testing.T) { + // Non-streaming events should encode/decode as before. + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + AgentName: "non-stream-agent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("direct", nil), + Role: schema.Assistant, + }, + }, + }, + } + + data, err := wrapper.GobEncode() + assert.NoError(t, err) + + decoded := &agentEventWrapper{AgentEvent: &AgentEvent{}} + assert.NoError(t, decoded.GobDecode(data)) + assert.Equal(t, "non-stream-agent", decoded.AgentName) + assert.False(t, decoded.Output.MessageOutput.IsStreaming) + }) + + t.Run("GobEncode_within_runSession", func(t *testing.T) { + // Simulate the real scenario: a runSession with a streaming event containing + // WillRetryError is gob-encoded (as happens during checkpoint save). + wrapper := makeStreamingEventWrapper( + schema.AssistantMessage("checkpoint content", nil), + &WillRetryError{ErrStr: "retry", RetryAttempt: 1}, + ) + + session := newRunSession() + session.Events = []*agentEventWrapper{wrapper} + + // Encode the entire session (the checkpoint path). + var buf bytes.Buffer + err := gob.NewEncoder(&buf).Encode(session) + assert.NoError(t, err, "encoding runSession with WillRetryError stream should succeed") + }) +} diff --git a/adk/runner.go b/adk/runner.go new file mode 100644 index 0000000..a7d722e --- /dev/null +++ b/adk/runner.go @@ -0,0 +1,342 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "runtime/debug" + "sync" + + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +func errorIterator[M MessageType](err error) *AsyncIterator[*TypedAgentEvent[M]] { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + gen.Send(&TypedAgentEvent[M]{Err: err}) + gen.Close() + return iter +} + +func newUserMessage[M MessageType](query string) (M, error) { + var zero M + switch any(zero).(type) { + case *schema.Message: + return any(schema.UserMessage(query)).(M), nil + case *schema.AgenticMessage: + return any(schema.UserAgenticMessage(query)).(M), nil + default: + return zero, fmt.Errorf("unsupported message type %T", zero) + } +} + +// TypedRunner is the primary entry point for executing an Agent. +// It manages the agent's lifecycle, including starting, resuming, and checkpointing. +// +// Execution always goes through the flowAgent pipeline, which handles +// multi-agent orchestration, callbacks, agent naming, run paths, and cancellation. +type TypedRunner[M MessageType] struct { + a TypedAgent[M] + enableStreaming bool + store CheckPointStore +} + +// Runner is the default runner type using *schema.Message. +type Runner = TypedRunner[*schema.Message] + +type CheckPointStore = core.CheckPointStore + +type CheckPointDeleter = core.CheckPointDeleter + +type TypedRunnerConfig[M MessageType] struct { + Agent TypedAgent[M] + EnableStreaming bool + + CheckPointStore CheckPointStore +} + +// RunnerConfig is the default runner config type using *schema.Message. +type RunnerConfig = TypedRunnerConfig[*schema.Message] + +// ResumeParams contains all parameters needed to resume an execution. +// This struct provides an extensible way to pass resume parameters without +// requiring breaking changes to method signatures. +type ResumeParams struct { + // Targets contains the addresses of components to be resumed as keys, + // with their corresponding resume data as values + Targets map[string]any + // Future extensible fields can be added here without breaking changes +} + +// NewRunner creates a new Runner with the given config. +func NewRunner(_ context.Context, conf RunnerConfig) *Runner { + return NewTypedRunner(conf) +} + +// NewTypedRunner creates a new TypedRunner with the given config. +func NewTypedRunner[M MessageType](conf TypedRunnerConfig[M]) *TypedRunner[M] { + return &TypedRunner[M]{ + enableStreaming: conf.EnableStreaming, + a: conf.Agent, + store: conf.CheckPointStore, + } +} + +func (r *TypedRunner[M]) Run(ctx context.Context, messages []M, + opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + return typedRunnerRunImpl(r.a, r.enableStreaming, r.store, ctx, messages, opts...) +} + +// Query is a convenience method that starts a new execution with a single user query string. +func (r *TypedRunner[M]) Query(ctx context.Context, + query string, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + msgs, err := newUserMessage[M](query) + if err != nil { + return errorIterator[M](err) + } + return r.Run(ctx, []M{msgs}, opts...) +} + +// Resume continues an interrupted execution from a checkpoint, using an "Implicit Resume All" strategy. +// This method is best for simpler use cases where the act of resuming implies that all previously +// interrupted points should proceed without specific data. +// +// When using this method, all interrupted agents will receive `isResumeFlow = false` when they +// call `GetResumeContext`, as no specific agent was targeted. This is suitable for the "Simple Confirmation" +// pattern where an agent only needs to know `wasInterrupted` is true to continue. +func (r *TypedRunner[M]) Resume(ctx context.Context, checkPointID string, opts ...AgentRunOption) ( + *AsyncIterator[*TypedAgentEvent[M]], error) { + return r.resumeInternal(ctx, checkPointID, nil, opts...) +} + +// ResumeWithParams continues an interrupted execution from a checkpoint with specific parameters. +// This is the most common and powerful way to resume, allowing you to target specific interrupt points +// (identified by their address/ID) and provide them with data. +// +// The params.Targets map should contain the addresses of the components to be resumed as keys. These addresses +// can point to any interruptible component in the entire execution graph, including ADK agents, compose +// graph nodes, or tools. The value can be the resume data for that component, or `nil` if no data is needed. +// +// When using this method: +// - Components whose addresses are in the params.Targets map will receive `isResumeFlow = true` when they +// call `GetResumeContext`. +// - Interrupted components whose addresses are NOT in the params.Targets map must decide how to proceed: +// -- "Leaf" components (the actual root causes of the original interrupt) MUST re-interrupt themselves +// to preserve their state. +// -- "Composite" agents (like SequentialAgent or ChatModelAgent) should generally proceed with their +// execution. They act as conduits, allowing the resume signal to flow to their children. They will +// naturally re-interrupt if one of their interrupted children re-interrupts, as they receive the +// new `CompositeInterrupt` signal from them. +func (r *TypedRunner[M]) ResumeWithParams(ctx context.Context, checkPointID string, params *ResumeParams, opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) { + return r.resumeInternal(ctx, checkPointID, params.Targets, opts...) +} + +func (r *TypedRunner[M]) resumeInternal(ctx context.Context, checkPointID string, resumeData map[string]any, + opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) { + return typedRunnerResumeInternalImpl(r.a, r.store, ctx, checkPointID, resumeData, opts...) +} + +func typedRunnerRunImpl[M MessageType](a TypedAgent[M], enableStreaming bool, store CheckPointStore, ctx context.Context, messages []M, opts ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[M]] { + o := getCommonOptions(nil, opts...) + + input := &TypedAgentInput[M]{ + Messages: messages, + EnableStreaming: enableStreaming, + } + + var zero M + if _, ok := any(zero).(*schema.Message); ok { + concreteAgent, _ := any(a).(Agent) + fa := toFlowAgent(ctx, concreteAgent) + if store != nil { + fa.checkPointStore = store + } + concreteInput := any(input).(*AgentInput) + ctx = ctxWithNewTypedRunCtx(ctx, input, o.sharedParentSession) + AddSessionValues(ctx, o.sessionValues) + + iter := fa.Run(ctx, concreteInput, opts...) + + if store == nil && o.cancelCtx == nil { + return any(iter).(*AsyncIterator[*TypedAgentEvent[M]]) + } + + niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(iter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, o.checkPointID, o.cancelCtx) + return niter + } + + fa := toTypedFlowAgent(a) + if store != nil { + fa.checkPointStore = store + } + + ctx = ctxWithNewTypedRunCtx(ctx, input, o.sharedParentSession) + AddSessionValues(ctx, o.sessionValues) + + iter := fa.Run(ctx, input, opts...) + + if store == nil && o.cancelCtx == nil { + return iter + } + + niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, iter, gen, o.checkPointID, o.cancelCtx) + return niter +} + +func typedRunnerResumeInternalImpl[M MessageType](a TypedAgent[M], store CheckPointStore, ctx context.Context, checkPointID string, resumeData map[string]any, //nolint:revive // argument-limit + opts ...AgentRunOption) (*AsyncIterator[*TypedAgentEvent[M]], error) { + if store == nil { + return nil, fmt.Errorf("failed to resume: store is nil") + } + + ctx, runCtx, resumeInfo, err := runnerLoadCheckPointImpl(store, ctx, checkPointID) + if err != nil { + return nil, fmt.Errorf("failed to load from checkpoint: %w", err) + } + + // Resume uses the streaming mode persisted in the checkpoint, not the value the + // caller passed when constructing the runner. This is the runner's own invariant: + // the checkpoint is the source of truth for what mode the original execution was + // running in, and any new checkpoint written during this resume must preserve it. + enableStreaming := resumeInfo.EnableStreaming + + o := getCommonOptions(nil, opts...) + if o.sharedParentSession { + parentSession := getSession(ctx) + if parentSession != nil { + runCtx.Session.Values = parentSession.Values + runCtx.Session.valuesMtx = parentSession.valuesMtx + } + } + if runCtx.Session.valuesMtx == nil { + runCtx.Session.valuesMtx = &sync.Mutex{} + } + if runCtx.Session.Values == nil { + runCtx.Session.Values = make(map[string]any) + } + + ctx = setRunCtx(ctx, runCtx) + AddSessionValues(ctx, o.sessionValues) + + if len(resumeData) > 0 { + ctx = core.BatchResumeWithData(ctx, resumeData) + } + + var zero M + if _, ok := any(zero).(*schema.Message); ok { + concreteAgent, _ := any(a).(Agent) + fa := toFlowAgent(ctx, concreteAgent) + ra, ok := Agent(fa).(ResumableAgent) + if !ok { + return nil, fmt.Errorf("agent %T does not support resume", a) + } + aIter := ra.Resume(ctx, resumeInfo, opts...) + + niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, any(aIter).(*AsyncIterator[*TypedAgentEvent[M]]), gen, &checkPointID, o.cancelCtx) + return niter, nil + } + + fa := toTypedFlowAgent(a) + ra, ok := TypedAgent[M](fa).(TypedResumableAgent[M]) + if !ok { + return nil, fmt.Errorf("agent %T does not support resume", a) + } + aIter := ra.Resume(ctx, resumeInfo, opts...) + + niter, gen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go typedRunnerHandleIterImpl(enableStreaming, store, ctx, aIter, gen, &checkPointID, o.cancelCtx) + return niter, nil +} + +func typedRunnerHandleIterImpl[M MessageType](enableStreaming bool, store CheckPointStore, ctx context.Context, aIter *AsyncIterator[*TypedAgentEvent[M]], //nolint:revive // argument-limit + gen *AsyncGenerator[*TypedAgentEvent[M]], checkPointID *string, cancelCtx *cancelContext) { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + gen.Send(&TypedAgentEvent[M]{Err: e}) + } + + gen.Close() + }() + var ( + interruptSignal *core.InterruptSignal + legacyData any + ) + for { + event, ok := aIter.Next() + if !ok { + break + } + + if event.Err != nil { + var cancelErr *CancelError + if errors.As(event.Err, &cancelErr) { + if cancelCtx != nil && cancelCtx.isRoot() && cancelCtx.shouldCancel() { + cancelCtx.markCancelHandled() + } + if cancelErr.interruptSignal != nil && checkPointID != nil { + cancelErr.InterruptContexts = core.ToInterruptContexts(cancelErr.interruptSignal, allowedAddressSegmentTypes) + err := runnerSaveCheckPointImpl(enableStreaming, store, ctx, *checkPointID, &InterruptInfo{}, cancelErr.interruptSignal) + if err != nil { + gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("failed to save checkpoint on cancel: %w", err)}) + } + } + gen.Send(event) + break + } + } + + if event.Action != nil && event.Action.internalInterrupted != nil { + if interruptSignal != nil { + panic("multiple interrupt actions should not happen in Runner") + } + interruptSignal = event.Action.internalInterrupted + interruptContexts := core.ToInterruptContexts(interruptSignal, allowedAddressSegmentTypes) + event = &TypedAgentEvent[M]{ + AgentName: event.AgentName, + RunPath: event.RunPath, + Output: event.Output, + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + Data: event.Action.Interrupted.Data, + InterruptContexts: interruptContexts, + }, + internalInterrupted: interruptSignal, + }, + } + legacyData = event.Action.Interrupted.Data + + if checkPointID != nil { + err := runnerSaveCheckPointImpl(enableStreaming, store, ctx, *checkPointID, &InterruptInfo{ + Data: legacyData, + }, interruptSignal) + if err != nil { + gen.Send(&TypedAgentEvent[M]{Err: fmt.Errorf("failed to save checkpoint: %w", err)}) + } + } + } + + gen.Send(event) + } +} diff --git a/adk/runner_test.go b/adk/runner_test.go new file mode 100644 index 0000000..cc59569 --- /dev/null +++ b/adk/runner_test.go @@ -0,0 +1,311 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +// mockRunnerAgent is a simple implementation of the Agent interface for testing Runner +type mockRunnerAgent struct { + name string + description string + responses []*AgentEvent + // Track calls to verify correct parameters were passed + callCount int + lastInput *AgentInput + enableStreaming bool +} + +func (a *mockRunnerAgent) Name(_ context.Context) string { + return a.name +} + +func (a *mockRunnerAgent) Description(_ context.Context) string { + return a.description +} + +func (a *mockRunnerAgent) Run(_ context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + // Record the call details for verification + a.callCount++ + a.lastInput = input + a.enableStreaming = input.EnableStreaming + + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + defer generator.Close() + + for _, event := range a.responses { + generator.Send(event) + + // If the event has an Exit action, stop sending events + if event.Action != nil && event.Action.Exit { + break + } + } + }() + + return iterator +} + +func newMockRunnerAgent(name, description string, responses []*AgentEvent) *mockRunnerAgent { + return &mockRunnerAgent{ + name: name, + description: description, + responses: responses, + } +} + +func TestNewRunner(t *testing.T) { + ctx := context.Background() + config := RunnerConfig{} + + runner := NewRunner(ctx, config) + + // Verify that a non-nil runner is returned + assert.NotNil(t, runner) +} + +func TestRunner_Run(t *testing.T) { + ctx := context.Background() + + // Create a mock agent with predefined responses + mockAgent_ := newMockRunnerAgent("TestAgent", "Test agent for Runner", []*AgentEvent{ + { + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from test agent", nil), + Role: schema.Assistant, + }, + }}, + }) + + // Create a runner + runner := NewRunner(ctx, RunnerConfig{Agent: mockAgent_}) + + // Create test messages + msgs := []Message{ + schema.UserMessage("Hello, agent!"), + } + + // Test Run method without streaming + iterator := runner.Run(ctx, msgs) + + // Verify that the agent's Run method was called with the correct parameters + assert.Equal(t, 1, mockAgent_.callCount) + assert.Equal(t, msgs, mockAgent_.lastInput.Messages) + assert.False(t, mockAgent_.enableStreaming) + + // Verify that we can get the expected response from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "TestAgent", event.AgentName) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.NotNil(t, event.Output.MessageOutput.Message) + assert.Equal(t, "Response from test agent", event.Output.MessageOutput.Message.Content) + + // Verify that the iterator is now closed + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestRunner_Run_WithStreaming(t *testing.T) { + ctx := context.Background() + + // Create a mock agent with predefined responses + mockAgent_ := newMockRunnerAgent("TestAgent", "Test agent for Runner", []*AgentEvent{ + { + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + Message: nil, + MessageStream: schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("Streaming response", nil)}), + Role: schema.Assistant, + }, + }}, + }) + + // Create a runner + runner := NewRunner(ctx, RunnerConfig{EnableStreaming: true, Agent: mockAgent_}) + + // Create test messages + msgs := []Message{ + schema.UserMessage("Hello, agent!"), + } + + // Test Run method with streaming enabled + iterator := runner.Run(ctx, msgs) + + // Verify that the agent's Run method was called with the correct parameters + assert.Equal(t, 1, mockAgent_.callCount) + assert.Equal(t, msgs, mockAgent_.lastInput.Messages) + assert.True(t, mockAgent_.enableStreaming) + + // Verify that we can get the expected response from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "TestAgent", event.AgentName) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + // Verify that the iterator is now closed + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestRunner_Query(t *testing.T) { + ctx := context.Background() + + // Create a mock agent with predefined responses + mockAgent_ := newMockRunnerAgent("TestAgent", "Test agent for Runner", []*AgentEvent{ + { + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response to query", nil), + Role: schema.Assistant, + }, + }}, + }) + + // Create a runner + runner := NewRunner(ctx, RunnerConfig{Agent: mockAgent_}) + + // Test Query method + iterator := runner.Query(ctx, "Test query") + + // Verify that the agent's Run method was called with the correct parameters + assert.Equal(t, 1, mockAgent_.callCount) + assert.Equal(t, 1, len(mockAgent_.lastInput.Messages)) + assert.Equal(t, "Test query", mockAgent_.lastInput.Messages[0].Content) + assert.False(t, mockAgent_.enableStreaming) + + // Verify that we can get the expected response from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "TestAgent", event.AgentName) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.NotNil(t, event.Output.MessageOutput.Message) + assert.Equal(t, "Response to query", event.Output.MessageOutput.Message.Content) + + // Verify that the iterator is now closed + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestRunner_Query_WithStreaming(t *testing.T) { + ctx := context.Background() + + // Create a mock agent with predefined responses + mockAgent_ := newMockRunnerAgent("TestAgent", "Test agent for Runner", []*AgentEvent{ + { + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + Message: nil, + MessageStream: schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("Streaming query response", nil)}), + Role: schema.Assistant, + }, + }}, + }) + + // Create a runner + runner := NewRunner(ctx, RunnerConfig{EnableStreaming: true, Agent: mockAgent_}) + + // Test Query method with streaming enabled + iterator := runner.Query(ctx, "Test query") + + // Verify that the agent's Run method was called with the correct parameters + assert.Equal(t, 1, mockAgent_.callCount) + assert.Equal(t, 1, len(mockAgent_.lastInput.Messages)) + assert.Equal(t, "Test query", mockAgent_.lastInput.Messages[0].Content) + assert.True(t, mockAgent_.enableStreaming) + + // Verify that we can get the expected response from the iterator + event, ok := iterator.Next() + assert.True(t, ok) + assert.Equal(t, "TestAgent", event.AgentName) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.True(t, event.Output.MessageOutput.IsStreaming) + + // Verify that the iterator is now closed + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestResumeWithMissingCheckpoint(t *testing.T) { + ctx := context.Background() + + agent := &myAgenticAgent{ + name: "resume-agent", + runFn: func(ctx context.Context, input *TypedAgentInput[*schema.AgenticMessage], options ...AgentRunOption) *AsyncIterator[*TypedAgentEvent[*schema.AgenticMessage]] { + iter, gen := NewAsyncIteratorPair[*TypedAgentEvent[*schema.AgenticMessage]]() + go func() { + defer gen.Close() + gen.Send(&TypedAgentEvent[*schema.AgenticMessage]{ + Output: &TypedAgentOutput[*schema.AgenticMessage]{ + MessageOutput: &TypedMessageVariant[*schema.AgenticMessage]{ + Message: agenticMsg("ok"), + }, + }, + }) + }() + return iter + }, + } + + store := newMyStore() + runner := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{ + Agent: agent, + CheckPointStore: store, + }) + + require.NotPanics(t, func() { + iter, err := runner.ResumeWithParams(ctx, "nonexistent-checkpoint", &ResumeParams{ + Targets: map[string]any{"fake-id": nil}, + }) + if err != nil { + t.Logf("Got expected error: %v", err) + return + } + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Err != nil { + t.Logf("Got error event: %v", event.Err) + } + } + }, "ResumeWithParams with nonexistent checkpoint should not panic") +} diff --git a/adk/turn_buffer.go b/adk/turn_buffer.go new file mode 100644 index 0000000..643c9bc --- /dev/null +++ b/adk/turn_buffer.go @@ -0,0 +1,134 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import "sync" + +type turnBuffer[T any] struct { + buf []T + mu sync.Mutex + notEmpty *sync.Cond + closed bool + woken bool +} + +func newTurnBuffer[T any]() *turnBuffer[T] { + tb := &turnBuffer[T]{} + tb.notEmpty = sync.NewCond(&tb.mu) + return tb +} + +func (tb *turnBuffer[T]) Send(value T) { + tb.mu.Lock() + defer tb.mu.Unlock() + + if tb.closed { + panic("turnBuffer: send on closed buffer") + } + + tb.buf = append(tb.buf, value) + tb.notEmpty.Signal() +} + +func (tb *turnBuffer[T]) TrySend(value T) bool { + tb.mu.Lock() + defer tb.mu.Unlock() + + if tb.closed { + return false + } + + tb.buf = append(tb.buf, value) + tb.notEmpty.Signal() + return true +} + +func (tb *turnBuffer[T]) Receive() (T, bool) { + tb.mu.Lock() + defer tb.mu.Unlock() + + for len(tb.buf) == 0 && !tb.closed && !tb.woken { + tb.notEmpty.Wait() + } + + tb.woken = false + + if len(tb.buf) == 0 { + var zero T + return zero, false + } + + val := tb.buf[0] + tb.buf = tb.buf[1:] + return val, true +} + +func (tb *turnBuffer[T]) Close() { + tb.mu.Lock() + defer tb.mu.Unlock() + + if !tb.closed { + tb.closed = true + tb.notEmpty.Broadcast() + } +} + +func (tb *turnBuffer[T]) IsClosed() bool { + tb.mu.Lock() + defer tb.mu.Unlock() + return tb.closed +} + +func (tb *turnBuffer[T]) TakeAll() []T { + tb.mu.Lock() + defer tb.mu.Unlock() + + if len(tb.buf) == 0 { + return nil + } + + values := tb.buf + tb.buf = nil + return values +} + +func (tb *turnBuffer[T]) PushFront(values []T) { + if len(values) == 0 { + return + } + + tb.mu.Lock() + defer tb.mu.Unlock() + + tb.buf = append(append([]T{}, values...), tb.buf...) + tb.notEmpty.Signal() +} + +func (tb *turnBuffer[T]) Wakeup() { + tb.mu.Lock() + defer tb.mu.Unlock() + + tb.woken = true + tb.notEmpty.Broadcast() +} + +func (tb *turnBuffer[T]) ClearWakeup() { + tb.mu.Lock() + defer tb.mu.Unlock() + + tb.woken = false +} diff --git a/adk/turn_loop.go b/adk/turn_loop.go new file mode 100644 index 0000000..c2fe185 --- /dev/null +++ b/adk/turn_loop.go @@ -0,0 +1,2056 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "fmt" + "runtime/debug" + "sync" + "sync/atomic" + "time" + + "github.com/cloudwego/eino/internal/safe" +) + +type stopPhase uint8 + +const ( + stopOpen stopPhase = iota + stopIdleWaiting + stopCommitted +) + +type preemptTurnPhase uint8 + +const ( + preemptTurnIdle preemptTurnPhase = iota + preemptTurnPlanning + preemptTurnActive +) + +func (p preemptTurnPhase) String() string { + switch p { + case preemptTurnIdle: + return "idle" + case preemptTurnPlanning: + return "planning" + case preemptTurnActive: + return "active" + default: + return fmt.Sprintf("unknown(%d)", p) + } +} + +type preemptTurnSnapshot struct { + hasTargetTurn bool + turnID uint64 + ctx context.Context + tc any +} + +type cancelRequestState struct { + cfg agentCancelConfig + timeoutDeadline *time.Time +} + +type preemptRequest struct { + cancel cancelRequestState + ackChans []chan struct{} +} + +func parseAgentCancelOptions(opts ...AgentCancelOption) agentCancelConfig { + cfg := agentCancelConfig{Mode: CancelImmediate} + for _, opt := range opts { + opt(&cfg) + } + return cfg +} + +func newCancelRequestState(opts []AgentCancelOption, now time.Time) cancelRequestState { + cfg := parseAgentCancelOptions(opts...) + var deadline *time.Time + if cfg.Timeout != nil && *cfg.Timeout > 0 && cfg.Mode != CancelImmediate { + d := now.Add(*cfg.Timeout) + deadline = &d + } + cfg.Timeout = nil + + return cancelRequestState{ + cfg: cfg, + timeoutDeadline: deadline, + } +} + +func (s *cancelRequestState) merge(opts []AgentCancelOption, now time.Time) { + if opts == nil { + return + } + + next := newCancelRequestState(opts, now) + if s.cfg.Mode == CancelImmediate || next.cfg.Mode == CancelImmediate { + s.cfg.Mode = CancelImmediate + s.timeoutDeadline = nil + } else { + s.cfg.Mode |= next.cfg.Mode + if next.timeoutDeadline != nil { + if s.timeoutDeadline == nil || next.timeoutDeadline.Before(*s.timeoutDeadline) { + deadline := *next.timeoutDeadline + s.timeoutDeadline = &deadline + } + } + } + if next.cfg.Recursive { + s.cfg.Recursive = true + } +} + +func (s cancelRequestState) cancelOptions(now time.Time) []AgentCancelOption { + cfg := s.cfg + if cfg.Mode != CancelImmediate && s.timeoutDeadline != nil { + remaining := s.timeoutDeadline.Sub(now) + if remaining <= 0 { + cfg.Mode = CancelImmediate + cfg.Timeout = nil + } else { + cfg.Timeout = &remaining + } + } + + opts := []AgentCancelOption{WithAgentCancelMode(cfg.Mode)} + if cfg.Recursive { + opts = append(opts, WithRecursive()) + } + if cfg.Timeout != nil { + opts = append(opts, WithAgentCancelTimeout(*cfg.Timeout)) + } + return opts +} + +func newPreemptRequest(ack chan struct{}, opts []AgentCancelOption, now time.Time) *preemptRequest { + req := &preemptRequest{cancel: newCancelRequestState(opts, now)} + if ack != nil { + req.ackChans = append(req.ackChans, ack) + } + return req +} + +func (r *preemptRequest) ack() { + if r == nil { + return + } + for _, ack := range r.ackChans { + close(ack) + } + r.ackChans = nil +} + +func (r *preemptRequest) merge(ack chan struct{}, opts []AgentCancelOption, now time.Time) { + if ack != nil { + r.ackChans = append(r.ackChans, ack) + } + r.cancel.merge(opts, now) +} + +func (r *preemptRequest) cancelOptions(now time.Time) []AgentCancelOption { + if r == nil { + return nil + } + return r.cancel.cancelOptions(now) +} + +// preemptController owns turn-targeted preempt requests and Push critical sections. +// +// Turn lifecycle: +// +// idle ──beginPlanningTurn──▶ planning ──beginActiveTurn──▶ active ──endActiveTurn──▶ idle +// │ ▲ +// └────────abortPlanningTurn─────────────────────────────┘ +// +// Push critical section (beginPush/endPush) overlaps with the turn lifecycle. The +// run loop calls waitForPushes before beginPlanningTurn to ensure no in-flight Push +// can observe stale turn state. +// +// Preempt request flow: +// - Push captures a snapshot (turnID + hasTargetTurn) via beginPush. +// - requestPreempt binds to the captured turnID; if the turn has moved on, the +// request is resolved as a no-op. +// - During active phase, receivePreempt transfers the pending request to the +// watcher, which submits cancel and then acks. +type preemptController struct { + mu sync.Mutex + cond *sync.Cond + + turnPhase preemptTurnPhase + turnID uint64 + currentTC any + currentRunCtx context.Context + + pushInFlight int + pending *preemptRequest + notify chan struct{} + closed bool +} + +func newPreemptController() *preemptController { + c := &preemptController{notify: make(chan struct{}, 1)} + c.cond = sync.NewCond(&c.mu) + return c +} + +func (c *preemptController) beginPlanningTurn() { + c.mu.Lock() + defer c.mu.Unlock() + + c.requirePhaseLocked(preemptTurnIdle, "beginPlanningTurn") + c.requireNoPendingLocked("beginPlanningTurn") + c.turnID++ + c.turnPhase = preemptTurnPlanning + c.currentRunCtx = nil + c.currentTC = nil +} + +func (c *preemptController) abortPlanningTurn() *preemptRequest { + c.mu.Lock() + defer c.mu.Unlock() + + c.requirePhaseLocked(preemptTurnPlanning, "abortPlanningTurn") + c.turnPhase = preemptTurnIdle + c.currentRunCtx = nil + c.currentTC = nil + req := c.pending + c.pending = nil + c.cond.Broadcast() + return req +} + +func (c *preemptController) beginActiveTurn(ctx context.Context, tc any) { + c.mu.Lock() + defer c.mu.Unlock() + + c.requirePhaseLocked(preemptTurnPlanning, "beginActiveTurn") + c.turnPhase = preemptTurnActive + c.currentRunCtx = ctx + c.currentTC = tc + if c.pending != nil { + c.notifyWatcherLocked() + } +} + +func (c *preemptController) endActiveTurn() *preemptRequest { + c.mu.Lock() + defer c.mu.Unlock() + + c.requirePhaseLocked(preemptTurnActive, "endActiveTurn") + c.turnPhase = preemptTurnIdle + c.currentRunCtx = nil + c.currentTC = nil + req := c.pending + c.pending = nil + c.cond.Broadcast() + return req +} + +func (c *preemptController) requirePhaseLocked(expected preemptTurnPhase, op string) { + if c.turnPhase != expected { + panic(fmt.Sprintf("adk: preemptController.%s called while turn phase is %s; expected %s", op, c.turnPhase, expected)) + } +} + +func (c *preemptController) requireNoPendingLocked(op string) { + if c.pending != nil { + panic(fmt.Sprintf("adk: preemptController.%s called with stale pending preempt request", op)) + } +} + +func (c *preemptController) beginPush() preemptTurnSnapshot { + c.mu.Lock() + defer c.mu.Unlock() + + c.pushInFlight++ + return preemptTurnSnapshot{ + hasTargetTurn: c.turnPhase == preemptTurnPlanning || c.turnPhase == preemptTurnActive, + turnID: c.turnID, + ctx: c.currentRunCtx, + tc: c.currentTC, + } +} + +func (c *preemptController) endPush() { + c.mu.Lock() + defer c.mu.Unlock() + + c.pushInFlight-- + if c.pushInFlight < 0 { + panic("adk: preemptController.endPush called without matching beginPush") + } + c.cond.Broadcast() +} + +func (c *preemptController) waitForPushes() { + c.mu.Lock() + defer c.mu.Unlock() + + for c.pushInFlight > 0 { + c.cond.Wait() + } +} + +func (c *preemptController) requestPreempt(target preemptTurnSnapshot, ack chan struct{}, opts ...AgentCancelOption) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.closed || !target.hasTargetTurn || c.turnPhase == preemptTurnIdle || c.turnID != target.turnID { + if ack != nil { + close(ack) + } + return + } + + now := time.Now() + if c.pending == nil { + c.pending = newPreemptRequest(ack, opts, now) + } else { + c.pending.merge(ack, opts, now) + } + if c.turnPhase == preemptTurnActive { + c.notifyWatcherLocked() + } +} + +func (c *preemptController) receivePreempt() (*preemptRequest, bool) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.turnPhase != preemptTurnActive || c.pending == nil { + return nil, false + } + req := c.pending + c.pending = nil + return req, true +} + +func (c *preemptController) closeForLoopExit() { + c.mu.Lock() + c.closed = true + c.turnPhase = preemptTurnIdle + c.currentRunCtx = nil + c.currentTC = nil + req := c.pending + c.pending = nil + select { + case <-c.notify: + default: + } + c.cond.Broadcast() + c.mu.Unlock() + + req.ack() +} + +func (c *preemptController) notifyWatcherLocked() { + select { + case c.notify <- struct{}{}: + default: + } +} + +type stopDecision struct { + commit bool + wakeIdle bool +} + +type stopCancelRequest struct { + cancel cancelRequestState +} + +func newStopCancelRequest(opts []AgentCancelOption, now time.Time) *stopCancelRequest { + return &stopCancelRequest{cancel: newCancelRequestState(opts, now)} +} + +func (r *stopCancelRequest) merge(opts []AgentCancelOption, now time.Time) { + if r == nil { + return + } + r.cancel.merge(opts, now) +} + +func (r *stopCancelRequest) cancelOptions(now time.Time) []AgentCancelOption { + if r == nil { + return nil + } + return r.cancel.cancelOptions(now) +} + +// stopController owns global Stop state and optional active-turn cancel requests. +// +// Stop has two independent layers: +// - terminal loop intent: committed Stop prevents future turns and closes the buffer; +// - optional active-turn cancel: cancel-capable Stop calls create a pending request +// consumed by the watcher if the current turn is still active. +// +// Unlike preempt, Stop is not bound to a turnID. It is global and terminal. +// A pending cancel request is consumed by the active turn or dropped when that +// turn ends before consumption. +type stopController struct { + mu sync.Mutex + + phase stopPhase + + hasActiveCancelTarget bool + pending *stopCancelRequest + notify chan struct{} + + idleFor time.Duration + skipCheckpoint bool + stopCause string + + closed bool +} + +func newStopController() *stopController { + return &stopController{notify: make(chan struct{}, 1)} +} + +func (c *stopController) requestStop(cfg *stopConfig) stopDecision { + c.mu.Lock() + defer c.mu.Unlock() + + if c.closed { + return stopDecision{} + } + if cfg.skipCheckpoint { + c.skipCheckpoint = true + } + if cfg.stopCause != "" && c.stopCause == "" { + c.stopCause = cfg.stopCause + } + if cfg.idleFor > 0 { + if c.phase != stopCommitted && c.idleFor == 0 { + c.phase = stopIdleWaiting + c.idleFor = cfg.idleFor + } + return stopDecision{wakeIdle: c.phase == stopIdleWaiting} + } + + committed := c.commitLocked() + if cfg.agentCancelOpts != nil { + now := time.Now() + if c.pending == nil { + c.pending = newStopCancelRequest(cfg.agentCancelOpts, now) + } else { + c.pending.merge(cfg.agentCancelOpts, now) + } + if c.hasActiveCancelTarget { + c.notifyWatcherLocked() + } + } + return stopDecision{commit: committed} +} + +func (c *stopController) commit() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.commitLocked() +} + +func (c *stopController) commitLocked() bool { + if c.closed || c.phase == stopCommitted { + return false + } + c.phase = stopCommitted + c.idleFor = 0 + return true +} + +func (c *stopController) isCommitted() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.phase == stopCommitted +} + +func (c *stopController) idleDuration() time.Duration { + c.mu.Lock() + defer c.mu.Unlock() + if c.phase != stopIdleWaiting { + return 0 + } + return c.idleFor +} + +func (c *stopController) skipCheckpointEnabled() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.skipCheckpoint +} + +func (c *stopController) cause() string { + c.mu.Lock() + defer c.mu.Unlock() + return c.stopCause +} + +func (c *stopController) beginActiveTurn() { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return + } + c.hasActiveCancelTarget = true + if c.pending != nil { + c.notifyWatcherLocked() + } +} + +func (c *stopController) endActiveTurn() *stopCancelRequest { + c.mu.Lock() + defer c.mu.Unlock() + c.hasActiveCancelTarget = false + req := c.pending + c.pending = nil + return req +} + +func (c *stopController) receiveCancel() (*stopCancelRequest, bool) { + c.mu.Lock() + defer c.mu.Unlock() + if !c.hasActiveCancelTarget || c.pending == nil { + return nil, false + } + req := c.pending + c.pending = nil + return req, true +} + +func (c *stopController) closeForLoopExit() { + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + c.hasActiveCancelTarget = false + c.pending = nil + select { + case <-c.notify: + default: + } +} + +func (c *stopController) notifyWatcherLocked() { + select { + case c.notify <- struct{}{}: + default: + } +} + +// TurnLoopConfig is the configuration for creating a TurnLoop. +type TurnLoopConfig[T any, M MessageType] struct { + // GenInput receives the TurnLoop instance and all buffered items, and decides what to process. + // It returns which items to consume now vs keep for later turns. + // The loop parameter allows calling Push() or Stop() directly from within the callback. + // Required. + GenInput func(ctx context.Context, loop *TurnLoop[T, M], items []T) (*GenInputResult[T, M], error) + + // GenResume is called at most once during Run(). When CheckpointID is + // configured, Run() queries Store for the checkpoint: + // - If the checkpoint contains runner state (i.e. an agent was interrupted + // or canceled mid-turn), Run() calls GenResume to plan a resume turn. + // - Otherwise (no checkpoint, or between-turns checkpoint), GenResume is + // never called and the loop proceeds via GenInput. + // + // It receives: + // - interruptedItems: the items being processed when the prior run was interrupted / canceled + // - unhandledItems: items buffered but not processed when the prior run exited + // - newItems: items that were Push()-ed before Run() was called + // + // It returns a GenResumeResult describing how to resume the interrupted agent + // turn (optional ResumeParams) and how to manipulate the buffer + // (Consumed/Remaining) before continuing. + GenResume func(ctx context.Context, loop *TurnLoop[T, M], interruptedItems, unhandledItems, newItems []T) (*GenResumeResult[T, M], error) + + // PrepareAgent returns an Agent configured to handle the consumed items. + // This callback should set up the agent with appropriate system prompt, + // tools, and middlewares based on what items are being processed. + // Called once per turn with the items that GenInput decided to consume. + // The loop parameter allows calling Push() or Stop() directly from within the callback. + // Required. + PrepareAgent func(ctx context.Context, loop *TurnLoop[T, M], consumed []T) (TypedAgent[M], error) + + // OnAgentEvents is called to handle events emitted by the agent. + // The TurnContext provides per-turn info and control: + // - tc.Consumed: items that triggered this agent execution + // - tc.Loop: allows calling Push() or Stop() directly from within the callback + // - tc.Preempted / tc.Stopped: signals while processing events + // + // Error handling: the returned error is only used when the callback itself + // wants to abort the TurnLoop. The callback should NEVER propagate + // CancelError — the framework handles it automatically: + // - Stop: the framework propagates CancelError as ExitReason, loop exits. + // - Preempt: the framework does not propagate CancelError; if the callback + // also returns nil, the loop continues with the next turn. + // In practice, return a non-nil error only for callback-internal failures + // that should terminate the loop. + // + // Optional. If not provided, events are drained and the first error + // (including CancelError from Stop) is returned as ExitReason. + OnAgentEvents func(ctx context.Context, tc *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error + + // Store is the checkpoint store for persistence and resume. Optional. + // When set together with CheckpointID, enables automatic checkpoint-based resume. + // The TurnLoop always persists both runner checkpoint bytes and item bookkeeping + // (InterruptedItems, UnhandledItems) via gob encoding, so T must be gob-encodable + // when Store is used. + Store CheckPointStore + + // CheckpointID, when set together with Store, enables automatic + // checkpoint-based resume. On Run(), the TurnLoop queries Store for this ID: + // - If a checkpoint exists with runner state (mid-turn interrupt / cancel), + // GenResume is called to plan the resume turn. + // - If a checkpoint exists without runner state (between-turns), + // the stored unhandled items are buffered and the loop proceeds + // normally via GenInput. + // - If no checkpoint exists, the loop starts fresh. + // + // On exit, if the TurnLoop saved a new checkpoint, it is saved under this + // same CheckpointID. On clean exit (no checkpoint saved), the existing + // checkpoint under CheckpointID is deleted to prevent stale resumption. + CheckpointID string +} + +// GenInputResult contains the result of GenInput processing. +type GenInputResult[T any, M MessageType] struct { + // RunCtx, if non-nil, overrides the context for this turn's execution + // (PrepareAgent, agent run, OnAgentEvents). + // + // Must be derived from the ctx passed to GenInput to preserve the + // TurnLoop's cancellation semantics and inherited values. For example: + // + // runCtx := context.WithValue(ctx, traceKey{}, extractTraceID(items)) + // return &GenInputResult[T]{RunCtx: runCtx, ...}, nil + // + // If nil, the TurnLoop's context is used unchanged. + RunCtx context.Context + + // Input is the agent input to execute + Input *TypedAgentInput[M] + + // RunOpts are the options for this agent run. + // Note: do not pass WithCheckPointID here; the TurnLoop automatically + // injects the checkpointID into the Runner. + RunOpts []AgentRunOption + + // Consumed are the items selected for this turn. + // They are removed from the buffer and passed to PrepareAgent. + Consumed []T + + // Remaining are the items to keep in the buffer for a future turn. + // TurnLoop pushes Remaining back into the buffer before running the agent. + // + // Items from the GenInput input slice that are in neither Consumed nor Remaining + // are dropped by the loop. + Remaining []T +} + +// GenResumeResult contains the result of GenResume processing. +type GenResumeResult[T any, M MessageType] struct { + // RunCtx, if non-nil, overrides the context for this resumed turn's execution + // (PrepareAgent, agent resume, OnAgentEvents). + RunCtx context.Context + + // RunOpts are the options for this agent resume run. + // Note: do not pass WithCheckPointID here; the TurnLoop automatically + // injects the checkpointID into the Runner. + RunOpts []AgentRunOption + + // ResumeParams are optional parameters for resuming an interrupted agent. + ResumeParams *ResumeParams + + // Consumed are the items selected for this resumed turn. + // They are removed from the buffer and passed to PrepareAgent. + Consumed []T + + // Remaining are the items to keep in the buffer for a future turn. + // TurnLoop pushes Remaining back into the buffer before resuming the agent. + // + // Items from (interruptedItems, unhandledItems, newItems) that are in neither Consumed + // nor Remaining are dropped by the loop. + Remaining []T +} + +type turnRunSpec[T any, M MessageType] struct { + runCtx context.Context + input *TypedAgentInput[M] + runOpts []AgentRunOption + resumeParams *ResumeParams + isResume bool + consumed []T + resumeBytes []byte +} + +type turnPlan[T any, M MessageType] struct { + turnCtx context.Context + remaining []T + spec *turnRunSpec[T, M] +} + +func (l *TurnLoop[T, M]) planTurn( + ctx context.Context, + isResume bool, + items []T, + pr *turnLoopPendingResume[T], +) (*turnPlan[T, M], error) { + if !isResume { + result, err := l.config.GenInput(ctx, l, items) + if err != nil { + return nil, err + } + if result == nil { + return nil, errors.New("GenInputResult is nil") + } + if result.Input == nil { + return nil, errors.New("agent input is nil") + } + turnCtx := ctx + if result.RunCtx != nil { + turnCtx = result.RunCtx + } + return &turnPlan[T, M]{ + turnCtx: turnCtx, + remaining: result.Remaining, + spec: &turnRunSpec[T, M]{ + runCtx: result.RunCtx, + input: result.Input, + runOpts: result.RunOpts, + consumed: result.Consumed, + }, + }, nil + } + if pr == nil { + return nil, errors.New("resume payload is nil") + } + if l.config.GenResume == nil { + return nil, errors.New("GenResume is required for resume") + } + resumeResult, err := l.config.GenResume(ctx, l, pr.interrupted, pr.unhandled, pr.newItems) + if err != nil { + return nil, err + } + if resumeResult == nil { + return nil, errors.New("GenResumeResult is nil") + } + turnCtx := ctx + if resumeResult.RunCtx != nil { + turnCtx = resumeResult.RunCtx + } + return &turnPlan[T, M]{ + turnCtx: turnCtx, + remaining: resumeResult.Remaining, + spec: &turnRunSpec[T, M]{ + runCtx: resumeResult.RunCtx, + runOpts: resumeResult.RunOpts, + resumeParams: resumeResult.ResumeParams, + isResume: true, + consumed: resumeResult.Consumed, + resumeBytes: pr.resumeBytes, + }, + }, nil +} + +// InterruptError is the ExitReason when the TurnLoop exits due to a business +// interrupt (AgentAction.Interrupted). It carries InterruptContexts needed for +// targeted resumption via ResumeParams, parallel to CancelError. +// +// Unlike CancelError (which indicates forceful cancellation), InterruptError +// indicates the agent voluntarily paused execution at a business-defined point. +type InterruptError struct { + // InterruptContexts provides the interrupt contexts needed for targeted + // resumption via ResumeParams. Each context represents a step in the agent + // hierarchy that was interrupted. Use each InterruptCtx.ID as a key in + // ResumeParams.Targets. + InterruptContexts []*InterruptCtx +} + +func (e *InterruptError) Error() string { + return fmt.Sprintf("agent interrupted: %d context(s)", len(e.InterruptContexts)) +} + +// TurnLoopExitState is returned when TurnLoop exits, containing the exit reason +// and any items that were not processed. +type TurnLoopExitState[T any, M MessageType] struct { + // ExitReason indicates why the loop exited. + // nil means clean exit (Stop() was called without cancel options, or the + // agent completed normally before Stop took effect). + // Non-nil values include context errors, callback errors, *CancelError, etc. + // When Stop(WithImmediate()) or Stop(WithGraceful()) cancels a running + // agent, ExitReason will be a *CancelError. + // This never contains checkpoint errors — see CheckpointErr for those. + ExitReason error + + // UnhandledItems contains items that were buffered but not processed. + // These are items for which Push returned true but were never consumed by a turn. + // This is always valid regardless of ExitReason. + UnhandledItems []T + + // InterruptedItems contains the items whose turn was interrupted — either by + // a cancel (Stop with cancel options → *CancelError) or by a business + // interrupt (AgentAction.Interrupted → *InterruptError). + // On resume, these are passed to GenResume's interruptedItems parameter. + InterruptedItems []T + + // StopCause is the business-supplied reason passed via WithStopCause. + // Empty if Stop was not called or no cause was provided. + StopCause string + + // CheckpointAttempted indicates whether a checkpoint save was attempted when the loop exited. + // True when Store is configured, CheckpointID is set, the loop was not idle + // at exit time, WithSkipCheckpoint was not used, and the exit was caused by + // Stop() (clean or cancel) or a business interrupt (*InterruptError). + CheckpointAttempted bool + + // CheckpointErr is the error from checkpoint save, if any. + // nil when CheckpointAttempted is false (no attempt was made) or when the save succeeded. + CheckpointErr error + + // TakeLateItems returns items that were pushed after the loop stopped + // (i.e., Push returned false for these items). These items are NOT included + // in the checkpoint. + // + // This function is idempotent: the first call computes and caches the result; + // subsequent calls return the same slice. + // + // After TakeLateItems is called, any subsequent Push() will panic to + // prevent items from being silently lost. + // + // It is safe to call TakeLateItems from any goroutine after Wait() returns. + // If TakeLateItems is never called, late items are simply garbage collected. + TakeLateItems func() []T +} + +// TurnContext provides per-turn context to the OnAgentEvents callback. +type TurnContext[T any, M MessageType] struct { + // Loop is the TurnLoop instance, allowing Push() or Stop() calls. + Loop *TurnLoop[T, M] + + // Consumed contains items that triggered this agent execution. + Consumed []T + + // Preempted is closed when a preempt signal fires for the current turn + // (via Push with WithPreempt/WithPreemptTimeout) and at least one + // preemptive Push contributed to the CancelError for the current turn. + // "Contributed" means the preempt's cancel options were included in the + // CancelError before it was finalized. Remains open if no preempt contributed. + // Use in a select to detect preemption while processing events. + // + // Both Preempted and Stopped may be closed within the same turn if both + // signals arrive while the agent is still being cancelled. Whichever + // arrives after the cancel is fully handled will not contribute. + Preempted <-chan struct{} + + // Stopped is closed when a Stop() call contributed to the CancelError for the + // current turn. + // "Contributed" means Stop's cancel options were included in the CancelError + // before it was finalized. Remains open if Stop did not contribute. + // Use in a select to detect stop while processing events. + // + // See Preempted for the relationship between the two channels. + Stopped <-chan struct{} + + // StopCause returns the business-supplied reason from WithStopCause. + // This value is only meaningful after the Stopped channel is closed. + // Before that, it returns an empty string. + StopCause func() string +} + +// TurnLoop is a push-based event loop for agent execution. +// Users push items via Push() and the loop processes them through the agent. +// +// Create with NewTurnLoop, then start with Run: +// +// loop := NewTurnLoop(cfg) +// // pass loop to other components, push initial items, etc. +// loop.Run(ctx) +// +// # Permissive API +// +// All methods are valid on a not-yet-running loop: +// - Push: items are buffered and will be processed once Run is called. +// - Stop: sets the stopped flag; a subsequent Run will exit immediately. +// - Wait: blocks until Run is called AND the loop exits. If Run is never +// called, Wait blocks forever (this is a programming error, analogous +// to reading from a channel that nobody writes to). +type TurnLoop[T any, M MessageType] struct { + config TurnLoopConfig[T, M] + + buffer *turnBuffer[T] + + stopped int32 + started int32 + + done chan struct{} + + result *TurnLoopExitState[T, M] + + runOnce sync.Once + + stopCtrl *stopController + + preemptCtrl *preemptController + + runErr error + + interruptedItems []T + + checkPointRunnerBytes []byte + interruptContexts []*InterruptCtx + capturedCancelErr *CancelError + + pendingResume *turnLoopPendingResume[T] + + loadCheckpointID string + + onAgentEvents func(ctx context.Context, tc *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error + + lateMu sync.Mutex + lateItems []T + lateSealed bool +} + +func (l *TurnLoop[T, M]) appendLate(item T) { + l.lateMu.Lock() + defer l.lateMu.Unlock() + if l.lateSealed { + panic("TurnLoop: Push called after TakeLateItems") + } + l.lateItems = append(l.lateItems, item) +} + +type turnLoopCheckpoint[T any] struct { + RunnerCheckpoint []byte + // HasRunnerState reports whether RunnerCheckpoint contains resumable runner state. + // It is false for "between turns" checkpoints where no agent execution was + // interrupted (e.g. Stop() before the first turn or between turns). + HasRunnerState bool + UnhandledItems []T + CanceledItems []T // gob-compat: kept as CanceledItems for deserialization of existing checkpoints +} + +func marshalTurnLoopCheckpoint[T any](c *turnLoopCheckpoint[T]) ([]byte, error) { + buf := new(bytes.Buffer) + if err := gob.NewEncoder(buf).Encode(c); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func unmarshalTurnLoopCheckpoint[T any](data []byte) (*turnLoopCheckpoint[T], error) { + var c turnLoopCheckpoint[T] + if err := gob.NewDecoder(bytes.NewReader(data)).Decode(&c); err != nil { + return nil, err + } + return &c, nil +} + +func (l *TurnLoop[T, M]) saveTurnLoopCheckpoint(ctx context.Context, checkPointID string, c *turnLoopCheckpoint[T]) error { + if l.config.Store == nil { + return errors.New("checkpoint store is nil") + } + data, err := marshalTurnLoopCheckpoint(c) + if err != nil { + return err + } + return l.config.Store.Set(ctx, checkPointID, data) +} + +func (l *TurnLoop[T, M]) deleteTurnLoopCheckpoint(ctx context.Context, checkPointID string) error { + if l.config.Store == nil { + return nil + } + if deleter, ok := l.config.Store.(CheckPointDeleter); ok { + return deleter.Delete(ctx, checkPointID) + } + return nil +} + +func (l *TurnLoop[T, M]) tryLoadCheckpoint(ctx context.Context) error { + checkPointID := l.config.CheckpointID + if checkPointID == "" || l.config.Store == nil { + return nil + } + + l.loadCheckpointID = checkPointID + + data, existed, err := l.config.Store.Get(ctx, checkPointID) + if err != nil { + return fmt.Errorf("failed to load checkpoint[%s]: %w", checkPointID, err) + } + if !existed { + return nil + } + + var cp *turnLoopCheckpoint[T] + if len(data) == 0 { + return nil + } + cp, err = unmarshalTurnLoopCheckpoint[T](data) + if err != nil { + return fmt.Errorf("failed to unmarshal checkpoint[%s]: %w", checkPointID, err) + } + + newItems := l.buffer.TakeAll() + + if cp.HasRunnerState { + if len(cp.RunnerCheckpoint) == 0 { + l.buffer.PushFront(newItems) + return fmt.Errorf("checkpoint[%s] has runner state but bytes are empty", checkPointID) + } + l.pendingResume = &turnLoopPendingResume[T]{ + interrupted: append([]T{}, cp.CanceledItems...), + unhandled: append([]T{}, cp.UnhandledItems...), + newItems: append([]T{}, newItems...), + resumeBytes: append([]byte{}, cp.RunnerCheckpoint...), + } + } else { + items := make([]T, 0, len(cp.UnhandledItems)+len(newItems)) + items = append(items, cp.UnhandledItems...) + items = append(items, newItems...) + l.buffer.PushFront(items) + } + + return nil +} + +type turnLoopPendingResume[T any] struct { + interrupted []T + unhandled []T + newItems []T + resumeBytes []byte +} + +// SafePoint describes at which boundary the agent may be cancelled. +// It is a bitmask: values can be combined with bitwise OR to accept multiple +// safe points (e.g. AfterToolCalls | AfterChatModel). Internally, SafePoint +// is translated to CancelMode via toCancelMode(). +// +// SafePoint is used only in the preemption API (WithPreempt/WithPreemptTimeout). +// A key design constraint: preemption always targets a safe point — the user's +// intent is to cancel at a well-defined boundary, never to abort immediately. +// Immediate cancellation is only reachable as an automatic timeout escalation +// (via WithPreemptTimeout), not as a direct user choice. This is why SafePoint +// has no "immediate" value and why WithPreempt requires a non-zero SafePoint +// (panics otherwise). +type SafePoint int + +const ( + // AfterChatModel allows the agent to finish the current chat-model + // call before being cancelled. + AfterChatModel SafePoint = 1 << iota + // AfterToolCalls allows the agent to finish the current tool-call round + // before being cancelled. + AfterToolCalls + // AnySafePoint is shorthand for AfterChatModel | AfterToolCalls. + AnySafePoint = AfterChatModel | AfterToolCalls +) + +func (sp SafePoint) toCancelMode() CancelMode { + var mode CancelMode + if sp&AfterToolCalls != 0 { + mode |= CancelAfterToolCalls + } + if sp&AfterChatModel != 0 { + mode |= CancelAfterChatModel + } + return mode +} + +type stopConfig struct { + agentCancelOpts []AgentCancelOption + skipCheckpoint bool + stopCause string + idleFor time.Duration +} + +// StopOption is an option for Stop(). +type StopOption func(*stopConfig) + +// WithGraceful requests a graceful stop that waits at the nearest safe point +// (after tool calls or after a chat-model call) and propagates recursively to +// nested agents. It does not impose a time limit; use WithGracefulTimeout to +// add a grace period after which the stop escalates to immediate cancellation. +// +// WithGraceful and WithGracefulTimeout are mutually exclusive; if both are +// passed to the same Stop call, the last one wins. +func WithGraceful() StopOption { + return func(cfg *stopConfig) { + cfg.agentCancelOpts = []AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel | CancelAfterToolCalls), + WithRecursive(), + } + } +} + +// WithImmediate aborts the running agent turn as soon as possible. +// The agent is cancelled immediately without waiting for any safe point. +// Nested agents inside AgentTools will also receive the cancel signal +// and be torn down. +// +// This is the most aggressive stop mode — typically used when the caller +// wants to shut down the TurnLoop with no intention of resuming. +func WithImmediate() StopOption { + return func(cfg *stopConfig) { + cfg.agentCancelOpts = []AgentCancelOption{ + WithRecursive(), + } + } +} + +// WithGracefulTimeout is like WithGraceful but adds a grace period. +// If the agent has not reached a safe point within gracePeriod, the stop +// escalates to immediate cancellation. +// +// gracePeriod must be positive; passing a zero or negative duration panics. +// +// WithGraceful and WithGracefulTimeout are mutually exclusive; if both are +// passed to the same Stop call, the last one wins. +func WithGracefulTimeout(gracePeriod time.Duration) StopOption { + if gracePeriod <= 0 { + panic("adk: WithGracefulTimeout: gracePeriod must be positive") + } + return func(cfg *stopConfig) { + cfg.agentCancelOpts = []AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel | CancelAfterToolCalls), + WithRecursive(), + WithAgentCancelTimeout(gracePeriod), + } + } +} + +// WithSkipCheckpoint tells the TurnLoop not to persist a checkpoint for this +// Stop call. Use this when the caller does not intend to resume in the future. +// The flag is sticky: once any Stop() call sets it, subsequent calls cannot undo it. +func WithSkipCheckpoint() StopOption { + return func(cfg *stopConfig) { + cfg.skipCheckpoint = true + } +} + +// WithStopCause attaches a business-supplied reason string to this Stop call. +// The cause is surfaced in TurnLoopExitState.StopCause and, after the Stopped +// channel closes, via TurnContext.StopCause(). +// If multiple Stop() calls provide a cause, the first non-empty value wins. +func WithStopCause(cause string) StopOption { + return func(cfg *stopConfig) { + cfg.stopCause = cause + } +} + +// UntilIdleFor defers the stop until the TurnLoop has been continuously idle +// (blocked between turns with no pending items) for at least the given +// duration. Each time a new item arrives the timer resets from zero. +// +// This is useful when business code monitors agent activity externally and +// wants to shut down the loop once there has been no work for a while, without +// racing with concurrent Push calls. +// +// UntilIdleFor does not impact a running agent. It only takes effect when the +// loop is idle between turns. Cancel options (WithImmediate, WithGraceful, +// WithGracefulTimeout) in the same Stop call are silently ignored — they are +// meaningless alongside UntilIdleFor. +// +// To escalate after a prior UntilIdleFor, issue a separate Stop call: +// +// loop.Stop(UntilIdleFor(30 * time.Second)) // wait for idle +// // ... later, if you need to abort immediately: +// loop.Stop(WithImmediate()) // overrides the idle wait +// +// Only the first UntilIdleFor duration takes effect; subsequent calls with +// a different duration are ignored. A Stop() call without UntilIdleFor always +// shuts down the loop immediately regardless of any pending idle timer. +// +// UntilIdleFor is combinable with non-cancel StopOptions (WithSkipCheckpoint, +// WithStopCause) in the same call. +// +// duration must be positive; passing a zero or negative value panics. +func UntilIdleFor(duration time.Duration) StopOption { + if duration <= 0 { + panic("adk: UntilIdleFor: duration must be positive") + } + return func(cfg *stopConfig) { + cfg.idleFor = duration + } +} + +type pushConfig[T any, M MessageType] struct { + preempt bool + preemptDelay time.Duration + agentCancelOpts []AgentCancelOption + pushStrategy func(context.Context, *TurnContext[T, M]) []PushOption[T, M] +} + +// PushOption is an option for Push(). +type PushOption[T any, M MessageType] func(*pushConfig[T, M]) + +// WithPreempt signals that the current agent turn should be cancelled at the +// specified safePoint after pushing the new item. The loop cancels the current +// turn and starts a new one, where GenInput will see all buffered items +// including the newly pushed one. +// Use WithPreemptTimeout to add a timeout that escalates to immediate abort. +// +// Because safe points fire at turn-level boundaries (after the chat model +// returns or after all tool calls complete), no nested agent is running at +// the moment of cancellation — nested agents within AgentTools have either +// not started yet (AfterChatModel) or already finished (AfterToolCalls). +// Note: WithPreempt does NOT include WithRecursive (no escalation path exists). +// WithPreemptTimeout DOES include WithRecursive so that on timeout escalation, +// nested agents are properly torn down. +// +// WithPreempt and WithPreemptTimeout are mutually exclusive; if both are +// passed to the same Push call, the last one wins. +// +// safePoint must not be zero; passing SafePoint(0) panics. +func WithPreempt[T any, M MessageType](safePoint SafePoint) PushOption[T, M] { + if safePoint == 0 { + panic("adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint") + } + return func(cfg *pushConfig[T, M]) { + cfg.preempt = true + cfg.agentCancelOpts = []AgentCancelOption{ + WithAgentCancelMode(safePoint.toCancelMode()), + } + } +} + +// WithPreemptTimeout is like WithPreempt but adds a timeout. If the agent has +// not reached the safe point within timeout, the preemption escalates to +// immediate cancellation. On escalation, nested agents inside AgentTools will +// also receive the cancel signal and be torn down. +// +// safePoint must not be zero; passing SafePoint(0) panics. +func WithPreemptTimeout[T any, M MessageType](safePoint SafePoint, timeout time.Duration) PushOption[T, M] { + if safePoint == 0 { + panic("adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint") + } + return func(cfg *pushConfig[T, M]) { + cfg.preempt = true + cfg.agentCancelOpts = []AgentCancelOption{ + WithAgentCancelMode(safePoint.toCancelMode()), + WithAgentCancelTimeout(timeout), + WithRecursive(), + } + } +} + +// WithPreemptDelay sets a delay duration before resolving a preemptive Push. +// When used with WithPreempt or WithPreemptTimeout, the pushed item is buffered +// immediately, while the preempt request is resolved after the delay against the +// turn observed by Push. If that captured turn has already ended, the request is +// resolved as a no-op and must not cancel a later turn. +func WithPreemptDelay[T any, M MessageType](delay time.Duration) PushOption[T, M] { + return func(cfg *pushConfig[T, M]) { + cfg.preemptDelay = delay + } +} + +// WithPushStrategy provides dynamic push option resolution based on the current turn state. +// The callback receives the current turn's context and TurnContext (nil if no turn is active) +// and returns the actual PushOptions to apply. When WithPushStrategy is used, all other +// PushOptions passed to the same Push call are ignored. +// +// The returned options must not contain another WithPushStrategy; any nested +// strategy is silently stripped. +// +// Example: preempt only if the current turn is processing low-priority items: +// +// loop.Push(urgentItem, WithPushStrategy(func(ctx context.Context, tc *TurnContext[MyItem, *schema.Message]) []PushOption[MyItem, *schema.Message] { +// if tc == nil { +// return nil // between turns, plain push +// } +// if isLowPriority(tc.Consumed) { +// return []PushOption[MyItem, *schema.Message]{WithPreempt[MyItem, *schema.Message](AnySafePoint)} +// } +// return nil // don't preempt high-priority work +// })) +func WithPushStrategy[T any, M MessageType](fn func(ctx context.Context, tc *TurnContext[T, M]) []PushOption[T, M]) PushOption[T, M] { + return func(cfg *pushConfig[T, M]) { + cfg.pushStrategy = fn + } +} + +func defaultTurnLoopOnAgentEvents[T any, M MessageType](_ context.Context, _ *TurnContext[T, M], events *AsyncIterator[*TypedAgentEvent[M]]) error { + for { + event, ok := events.Next() + if !ok { + break + } + if event.Err != nil { + return event.Err + } + } + return nil +} + +// NewTurnLoop creates a new TurnLoop without starting it. +// The returned loop accepts Push and Stop calls immediately; pushed items +// are buffered until Run is called. +// Call Run to start the processing goroutine. +// +// NewTurnLoop panics if GenInput or PrepareAgent is nil. +func NewTurnLoop[T any, M MessageType](cfg TurnLoopConfig[T, M]) *TurnLoop[T, M] { + if cfg.GenInput == nil { + panic("adk: NewTurnLoop: GenInput is required") + } + if cfg.PrepareAgent == nil { + panic("adk: NewTurnLoop: PrepareAgent is required") + } + + l := &TurnLoop[T, M]{ + config: cfg, + buffer: newTurnBuffer[T](), + done: make(chan struct{}), + stopCtrl: newStopController(), + preemptCtrl: newPreemptController(), + } + if cfg.OnAgentEvents != nil { + l.onAgentEvents = cfg.OnAgentEvents + } else { + l.onAgentEvents = defaultTurnLoopOnAgentEvents[T, M] + } + return l +} + +func (l *TurnLoop[T, M]) start(ctx context.Context) { + l.runOnce.Do(func() { + atomic.StoreInt32(&l.started, 1) + go l.run(ctx) + }) +} + +// Run starts the loop's processing goroutine. It is non-blocking: the loop +// runs in the background and results are obtained via Wait. +// +// If CheckpointID is configured in TurnLoopConfig and a matching checkpoint +// exists in Store, the loop automatically resumes from that checkpoint. +// Otherwise it starts fresh with whatever items were Push()-ed. +// +// Calling Run more than once is a no-op: only the first call starts the loop. +func (l *TurnLoop[T, M]) Run(ctx context.Context) { + l.start(ctx) +} + +// Push adds an item to the loop's buffer for processing. +// This method is non-blocking and thread-safe. +// Returns false if the loop has stopped, true otherwise. If a preemptive push +// succeeds, the second return value is a channel that callers can wait on to +// confirm the preempt request has been resolved. Specifically: +// - If Push observes a planning or active turn that is still the target when +// the request resolves, the channel closes after TurnLoop attempts to submit +// cancel for that target turn. +// - If Push observes no target turn, the loop has not started, the preempt +// subsystem is closed, or a delayed target is already gone, the channel +// closes as a no-op resolution. +// +// If the loop has not been started yet (Run not called), items are buffered +// and will be processed once Run is called. +// After Wait() returns, failed pushes can be recovered via TurnLoopExitState.TakeLateItems(). +// Once TakeLateItems() has been called, any subsequent push that would become a +// late item will panic instead of being silently dropped. +// +// Use WithPreempt() or WithPreemptTimeout() to atomically push an item and signal +// preemption of the current agent. This is useful for urgent items that should +// interrupt the current processing. +// The returned channel may be waited on if the caller needs to ensure the preempt +// signal has been observed. +// +// Use WithPreemptDelay() together with WithPreempt()/WithPreemptTimeout() to delay +// request resolution. Push returns immediately after the item is buffered, and +// the delayed request remains bound to the turn observed by Push. +func (l *TurnLoop[T, M]) Push(item T, opts ...PushOption[T, M]) (bool, <-chan struct{}) { + cfg := &pushConfig[T, M]{} + for _, opt := range opts { + opt(cfg) + } + + if cfg.pushStrategy != nil { + return l.pushWithStrategy(item, cfg) + } + + return l.pushWithConfig(item, cfg) +} + +// pushWithStrategy snapshots the current target turn while the strategy decides +// how to enqueue the item. If it requests preempt, that request is bound to the +// captured turn identity, including delayed preempt requests. +func (l *TurnLoop[T, M]) pushWithStrategy(item T, cfg *pushConfig[T, M]) (bool, <-chan struct{}) { + strategy := cfg.pushStrategy + + snapshot := l.preemptCtrl.beginPush() + defer l.preemptCtrl.endPush() + + runCtx := snapshot.ctx + if runCtx == nil { + runCtx = context.Background() + } + var tc *TurnContext[T, M] + if snapshot.tc != nil { + tc = snapshot.tc.(*TurnContext[T, M]) + } + realOpts := strategy(runCtx, tc) + cfg = &pushConfig[T, M]{} + for _, opt := range realOpts { + opt(cfg) + } + cfg.pushStrategy = nil + + if !cfg.preempt { + if !l.buffer.TrySend(item) { + l.appendLate(item) + return false, nil + } + return true, nil + } + + if atomic.LoadInt32(&l.stopped) != 0 { + l.appendLate(item) + return false, nil + } + + if !l.buffer.TrySend(item) { + l.appendLate(item) + return false, nil + } + + ack := make(chan struct{}) + if atomic.LoadInt32(&l.started) == 0 { + close(ack) + return true, ack + } + + if cfg.preemptDelay > 0 { + go func() { + select { + case <-time.After(cfg.preemptDelay): + l.preemptCtrl.requestPreempt(snapshot, ack, cfg.agentCancelOpts...) + case <-l.done: + close(ack) + } + }() + } else { + l.preemptCtrl.requestPreempt(snapshot, ack, cfg.agentCancelOpts...) + } + return true, ack +} + +func (l *TurnLoop[T, M]) pushWithConfig(item T, cfg *pushConfig[T, M]) (bool, <-chan struct{}) { + if atomic.LoadInt32(&l.stopped) != 0 { + l.appendLate(item) + return false, nil + } + + if cfg.preempt { + snapshot := l.preemptCtrl.beginPush() + defer l.preemptCtrl.endPush() + + if !l.buffer.TrySend(item) { + l.appendLate(item) + return false, nil + } + + ack := make(chan struct{}) + if atomic.LoadInt32(&l.started) == 0 { + close(ack) + return true, ack + } + + if cfg.preemptDelay > 0 { + go func() { + select { + case <-time.After(cfg.preemptDelay): + l.preemptCtrl.requestPreempt(snapshot, ack, cfg.agentCancelOpts...) + case <-l.done: + close(ack) + } + }() + } else { + l.preemptCtrl.requestPreempt(snapshot, ack, cfg.agentCancelOpts...) + } + return true, ack + } + + if !l.buffer.TrySend(item) { + l.appendLate(item) + return false, nil + } + return true, nil +} + +// Stop signals the loop to stop and returns immediately (non-blocking). +// Without options, the current agent turn runs to completion and the loop +// exits at the turn boundary without starting a new turn. ExitReason is nil. +// +// Use WithImmediate() to abort the running agent turn immediately. +// Use WithGraceful() to cancel at the nearest safe point with recursive +// propagation to nested agents. +// Use WithGracefulTimeout() for safe-point cancel with an escalation deadline. +// Use UntilIdleFor() to defer the stop until the loop has been continuously +// idle for a given duration; the loop shuts down automatically once the idle +// timer fires. +// +// This method may be called multiple times; subsequent calls update cancel options. +// A Stop() call without UntilIdleFor shuts down the loop immediately, even if +// a prior UntilIdleFor is still waiting. +// Call Wait() to block until the loop has fully exited and get the result. +// +// Stop may be called before Run. In that case, the stopped flag is set and +// a subsequent Run will exit the loop immediately. +// +// If the running agent does not support the WithCancel AgentRunOption, +// all cancel-related options (WithImmediate, WithGraceful, WithGracefulTimeout) +// degrade to "exit the loop on entering the next iteration" — the current +// agent turn runs to completion before the loop exits. +func (l *TurnLoop[T, M]) Stop(opts ...StopOption) { + cfg := &stopConfig{} + for _, opt := range opts { + opt(cfg) + } + + // UntilIdleFor is incompatible with cancel options (WithImmediate, + // WithGraceful, WithGracefulTimeout) in the same call. Cancel opts only + // make sense for an immediate or escalated stop; UntilIdleFor defers the + // stop until idle, and must not impact a running agent. Drop them silently. + if cfg.idleFor > 0 { + cfg.agentCancelOpts = nil + } + + decision := l.stopCtrl.requestStop(cfg) + if decision.wakeIdle { + l.buffer.Wakeup() + } + if decision.commit { + l.finishStopCommit() + } +} + +func (l *TurnLoop[T, M]) commitStop() { + if !l.stopCtrl.commit() { + return + } + l.finishStopCommit() +} + +func (l *TurnLoop[T, M]) finishStopCommit() { + atomic.StoreInt32(&l.stopped, 1) + l.buffer.Close() +} + +// Wait blocks until the loop exits and returns the result. +// This method is safe to call from multiple goroutines. +// All callers will receive the same result. +// +// Wait blocks until Run is called AND the loop exits. If Run is +// never called, Wait blocks forever. +func (l *TurnLoop[T, M]) Wait() *TurnLoopExitState[T, M] { + <-l.done + return l.result +} + +func (l *TurnLoop[T, M]) run(ctx context.Context) { + defer l.cleanup(ctx) + + if err := l.tryLoadCheckpoint(ctx); err != nil { + l.runErr = err + return + } + + // Monitor context cancellation: close the buffer so that a blocking + // Receive() unblocks. The loop will then check ctx.Err() and exit. + go func() { + select { + case <-ctx.Done(): + l.buffer.Close() + case <-l.done: + } + }() + + for { + if l.stopCtrl.isCommitted() { + return + } + + isResume := false + var pr *turnLoopPendingResume[T] + var items []T + var pushBack []T + + if l.pendingResume != nil { + isResume = true + pr = l.pendingResume + l.pendingResume = nil + + l.preemptCtrl.waitForPushes() + pr.newItems = append(pr.newItems, l.buffer.TakeAll()...) + + pushBack = make([]T, 0, len(pr.interrupted)+len(pr.unhandled)+len(pr.newItems)) + pushBack = append(pushBack, pr.interrupted...) + pushBack = append(pushBack, pr.unhandled...) + pushBack = append(pushBack, pr.newItems...) + } else { + var first T + var ok bool + + if idleFor := l.stopCtrl.idleDuration(); idleFor > 0 { + l.buffer.ClearWakeup() + idleTimer := time.NewTimer(idleFor) + cancelIdle := make(chan struct{}) + // When the idle timer fires, commitStop closes the buffer via + // buffer.Close(), which broadcasts to unblock the pending + // Receive() call below. + go func() { + select { + case <-idleTimer.C: + l.commitStop() + case <-cancelIdle: + } + }() + + first, ok = l.buffer.Receive() + + idleTimer.Stop() + close(cancelIdle) + + // A spurious wakeup can occur if Stop(UntilIdleFor) called + // buffer.Wakeup() after ClearWakeup() above but before + // Receive() entered its wait. In that case, Receive returns + // !ok from the woken flag, not from buffer closure. + // Re-enter the loop so the idle timer restarts cleanly. + if !ok && !l.buffer.IsClosed() { + continue + } + } else { + first, ok = l.buffer.Receive() + // Woken up by Stop(UntilIdleFor); re-enter loop to start the idle timer. + if !ok && l.stopCtrl.idleDuration() > 0 { + continue + } + } + + if !ok { + if err := ctx.Err(); err != nil { + l.runErr = err + } + return + } + + if err := ctx.Err(); err != nil { + l.buffer.PushFront([]T{first}) + l.runErr = err + return + } + + if l.stopCtrl.isCommitted() { + l.buffer.PushFront([]T{first}) + return + } + + l.preemptCtrl.waitForPushes() + rest := l.buffer.TakeAll() + items = append([]T{first}, rest...) + pushBack = items + } + + l.preemptCtrl.beginPlanningTurn() + abortPlanning := func() { + l.preemptCtrl.abortPlanningTurn().ack() + } + + plan, err := l.planTurn(ctx, isResume, items, pr) + if err != nil { + abortPlanning() + if len(pushBack) > 0 { + l.buffer.PushFront(pushBack) + } + l.runErr = err + return + } + + if l.stopCtrl.isCommitted() { + abortPlanning() + if len(pushBack) > 0 { + l.buffer.PushFront(pushBack) + } + return + } + + agent, err := l.config.PrepareAgent(plan.turnCtx, l, plan.spec.consumed) + if err != nil { + abortPlanning() + if len(pushBack) > 0 { + l.buffer.PushFront(pushBack) + } + l.runErr = err + return + } + + if l.stopCtrl.isCommitted() { + abortPlanning() + if len(pushBack) > 0 { + l.buffer.PushFront(pushBack) + } + return + } + + l.buffer.PushFront(plan.remaining) + + runErr := l.runAgentAndHandleEvents(plan.turnCtx, agent, plan.spec) + + if runErr != nil { + // Set interruptedItems when a cancel or interrupt was captured from the + // event stream, regardless of what the user's callback returned. The items + // were factually mid-execution when the signal arrived. + if l.capturedCancelErr != nil || l.interruptContexts != nil { + l.interruptedItems = append([]T{}, plan.spec.consumed...) + } + l.runErr = runErr + return + } + + // Business interrupt: agent produced an Interrupted action, exit to persist checkpoint. + if l.interruptContexts != nil { + l.interruptedItems = append([]T{}, plan.spec.consumed...) + l.runErr = &InterruptError{InterruptContexts: l.interruptContexts} + return + } + } +} + +func (l *TurnLoop[T, M]) setupBridgeStore(spec *turnRunSpec[T, M], runOpts []AgentRunOption) ([]AgentRunOption, *bridgeStore, error) { + store := l.config.Store + if store == nil && spec.isResume { + return nil, nil, fmt.Errorf("failed to resume agent: checkpoint store is nil") + } + if store == nil { + return runOpts, nil, nil + } + runOpts = append(runOpts, WithCheckPointID(bridgeCheckpointID)) + if spec.isResume { + if len(spec.resumeBytes) == 0 { + return nil, nil, fmt.Errorf("resume checkpoint is empty") + } + return runOpts, newResumeBridgeStore(bridgeCheckpointID, spec.resumeBytes), nil + } + return runOpts, newBridgeStore(), nil +} + +// watchPreempt runs for the lifetime of a single active turn. It consumes +// pending preempt requests exactly once and submits cancel for that turn. +func (l *TurnLoop[T, M]) watchPreempt(done <-chan struct{}, agentCancelFunc AgentCancelFunc, preemptDone chan struct{}) { + preemptDoneClosed := false + for { + select { + case <-done: + return + case <-l.preemptCtrl.notify: + req, ok := l.preemptCtrl.receivePreempt() + if !ok { + continue + } + // CancelHandle is intentionally not awaited here: agentCancelFunc commits the cancel signal synchronously, + // while waiting would block until the turn finishes and can deadlock this watcher against the done signal. + _, contributed := agentCancelFunc(req.cancelOptions(time.Now())...) + if contributed && !preemptDoneClosed { + close(preemptDone) + preemptDoneClosed = true + } + req.ack() + } + } +} + +// watchStop runs for the lifetime of a single active turn. It consumes pending +// Stop cancel requests exactly once and submits them to that turn. +func (l *TurnLoop[T, M]) watchStop(done <-chan struct{}, agentCancelFunc AgentCancelFunc, stoppedDone chan struct{}) { + stoppedClosed := false + + submit := func(req *stopCancelRequest) { + _, contributed := agentCancelFunc(req.cancelOptions(time.Now())...) + if contributed && !stoppedClosed { + close(stoppedDone) + stoppedClosed = true + } + } + + for { + if req, ok := l.stopCtrl.receiveCancel(); ok { + submit(req) + continue + } + + select { + case <-done: + return + case <-l.stopCtrl.notify: + } + } +} + +func (l *TurnLoop[T, M]) runAgentAndHandleEvents( + ctx context.Context, + agent TypedAgent[M], + spec *turnRunSpec[T, M], +) error { + l.interruptContexts = nil + l.capturedCancelErr = nil + l.checkPointRunnerBytes = nil + + var iter *AsyncIterator[*TypedAgentEvent[M]] + + runOpts, ms, err := l.setupBridgeStore(spec, spec.runOpts) + if err != nil { + l.preemptCtrl.abortPlanningTurn().ack() + return err + } + store := l.config.Store + cancelOpt, agentCancelFunc := WithCancel() + runOpts = append(runOpts, cancelOpt) + + // For Run path the streaming mode comes from the input. For Resume path the + // runner reads the streaming mode persisted in the checkpoint, so the value we + // pass here is irrelevant. + enableStreaming := false + if spec.input != nil { + enableStreaming = spec.input.EnableStreaming + } + runner := NewTypedRunner(TypedRunnerConfig[M]{ + EnableStreaming: enableStreaming, + Agent: agent, + CheckPointStore: ms, + }) + + preemptDone := make(chan struct{}) + stoppedDone := make(chan struct{}) + + tc := &TurnContext[T, M]{ + Loop: l, + Consumed: spec.consumed, + Preempted: preemptDone, + Stopped: stoppedDone, + StopCause: l.stopCtrl.cause, + } + l.preemptCtrl.beginActiveTurn(ctx, tc) + l.stopCtrl.beginActiveTurn() + defer func() { + l.stopCtrl.endActiveTurn() + l.preemptCtrl.endActiveTurn().ack() + }() + + if spec.isResume { + var err error + if spec.resumeParams != nil { + iter, err = runner.ResumeWithParams(ctx, bridgeCheckpointID, spec.resumeParams, runOpts...) + } else { + iter, err = runner.Resume(ctx, bridgeCheckpointID, runOpts...) + } + if err != nil { + return fmt.Errorf("failed to resume agent: %w", err) + } + } else { + iter = runner.Run(ctx, spec.input.Messages, runOpts...) + } + + // Wrap iterator to capture framework-level signals (CancelError, InterruptContexts) + // from events before they flow to OnAgentEvents. This ensures the framework can + // track these signals independently of what the user's callback returns. + srcIter := iter + proxyIter, proxyGen := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + go func() { + defer proxyGen.Close() + for { + event, ok := srcIter.Next() + if !ok { + break + } + if event != nil { + if event.Err != nil { + var cancelErr *CancelError + if errors.As(event.Err, &cancelErr) { + l.capturedCancelErr = cancelErr + } + } + if event.Action != nil && event.Action.Interrupted != nil { + l.interruptContexts = event.Action.Interrupted.InterruptContexts + } + } + proxyGen.Send(event) + } + }() + iter = proxyIter + + handleEvents := func() error { + return l.onAgentEvents(ctx, tc, iter) + } + + done := make(chan struct{}) + var handleErr error + + go func() { + defer func() { + panicErr := recover() + if panicErr != nil { + handleErr = safe.NewPanicErr(panicErr, debug.Stack()) + } + close(done) + }() + handleErr = handleEvents() + }() + go l.watchPreempt(done, agentCancelFunc, preemptDone) + go l.watchStop(done, agentCancelFunc, stoppedDone) + + finalizeCheckpoint := func() error { + if store != nil && ms != nil { + data, ok, err := ms.Get(ctx, bridgeCheckpointID) + if err != nil { + return fmt.Errorf("failed to read runner checkpoint: %w", err) + } + if ok { + l.checkPointRunnerBytes = append([]byte{}, data...) + } + } + return nil + } + + // Wait for the turn to end. Three outcomes: + // + // done: Events fully handled (normal or error). If Stop() was + // called, save checkpoint so the caller can resume later. + // Also handle the select race: if preemptDone is closed + // too, treat as a preempt (return nil) instead of leaking + // the CancelError. + // + // preemptDone: A preemptive Push successfully cancelled the agent. + // Wait for the handleEvents goroutine to drain, then + // return nil — the run loop will start a new turn. + // + // stoppedDone: Stop() cancelled the agent. Save checkpoint so the + // caller can resume later. + select { + case <-done: + select { + case <-preemptDone: + return nil + default: + } + if err := finalizeCheckpoint(); err != nil { + if handleErr != nil { + handleErr = fmt.Errorf("%w; checkpoint error: %v", handleErr, err) + } else { + handleErr = err + } + } + return l.applyFrameworkCapturedError(handleErr) + case <-preemptDone: + <-done + return nil + case <-stoppedDone: + <-done + if err := finalizeCheckpoint(); err != nil { + if handleErr != nil { + handleErr = fmt.Errorf("%w; checkpoint error: %v", handleErr, err) + } else { + handleErr = err + } + } + return l.applyFrameworkCapturedError(handleErr) + } +} + +// applyFrameworkCapturedError resolves the final error for runAgentAndHandleEvents. +// Priority scheme: +// - If handleErr != nil: the user's callback error wins (framework does not overwrite). +// - If handleErr == nil and a CancelError was captured: use the captured CancelError. +// - If handleErr == nil and interrupt contexts were captured: this is handled by the +// caller (run loop) via l.interruptContexts, so return nil here. +// +// In all cases, the caller uses l.capturedCancelErr and l.interruptContexts to +// determine interruptedItems independently of the returned error. +func (l *TurnLoop[T, M]) applyFrameworkCapturedError(handleErr error) error { + if handleErr != nil { + return handleErr + } + if l.capturedCancelErr != nil { + return l.capturedCancelErr + } + return nil +} + +func (l *TurnLoop[T, M]) cleanup(ctx context.Context) { + atomic.StoreInt32(&l.stopped, 1) + + unhandled := l.buffer.TakeAll() + checkpointID := l.config.CheckpointID + isIdle := len(l.checkPointRunnerBytes) == 0 && len(unhandled) == 0 && len(l.interruptedItems) == 0 + + // Only save checkpoint when the loop exited due to an explicit Stop(), + // a CancelError, or a business interrupt (InterruptError). + // Also checkpoint when a cancel/interrupt was captured from the event stream + // but the user's callback returned a custom error (the items were still in-flight). + exitCausedByStop := l.runErr == nil || errors.As(l.runErr, new(*CancelError)) || l.capturedCancelErr != nil + businessInterrupt := errors.As(l.runErr, new(*InterruptError)) || l.interruptContexts != nil + shouldSaveCheckpoint := l.config.Store != nil && checkpointID != "" && + ((l.stopCtrl.isCommitted() && exitCausedByStop) || businessInterrupt) && + !isIdle && !l.stopCtrl.skipCheckpointEnabled() + + var checkpointed bool + var checkpointErr error + + if shouldSaveCheckpoint { + cp := &turnLoopCheckpoint[T]{ + RunnerCheckpoint: l.checkPointRunnerBytes, + HasRunnerState: len(l.checkPointRunnerBytes) > 0, + UnhandledItems: unhandled, + CanceledItems: l.interruptedItems, + } + checkpointed = true + checkpointErr = l.saveTurnLoopCheckpoint(ctx, checkpointID, cp) + } else if l.loadCheckpointID != "" { + _ = l.deleteTurnLoopCheckpoint(ctx, l.loadCheckpointID) + } + + var takeLateOnce sync.Once + var takeLateResult []T + + l.result = &TurnLoopExitState[T, M]{ + ExitReason: l.runErr, + UnhandledItems: unhandled, + InterruptedItems: l.interruptedItems, + StopCause: l.stopCtrl.cause(), + CheckpointAttempted: checkpointed, + CheckpointErr: checkpointErr, + TakeLateItems: func() []T { + takeLateOnce.Do(func() { + l.lateMu.Lock() + takeLateResult = append([]T{}, l.lateItems...) + l.lateSealed = true + l.lateMu.Unlock() + }) + return takeLateResult + }, + } + + l.stopCtrl.closeForLoopExit() + l.preemptCtrl.closeForLoopExit() + l.buffer.Close() + close(l.done) +} diff --git a/adk/turn_loop_test.go b/adk/turn_loop_test.go new file mode 100644 index 0000000..2fb903b --- /dev/null +++ b/adk/turn_loop_test.go @@ -0,0 +1,6024 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +type turnLoopMockAgent struct { + name string + events []*AgentEvent + runFunc func(ctx context.Context, input *AgentInput) (*AgentOutput, error) +} + +func (a *turnLoopMockAgent) Name(_ context.Context) string { return a.name } +func (a *turnLoopMockAgent) Description(_ context.Context) string { return "mock agent" } +func (a *turnLoopMockAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + + if a.runFunc != nil { + go func() { + defer gen.Close() + output, err := a.runFunc(ctx, input) + if err != nil { + gen.Send(&AgentEvent{Err: err}) + return + } + gen.Send(&AgentEvent{Output: output}) + }() + return iter + } + + go func() { + defer gen.Close() + for _, e := range a.events { + gen.Send(e) + } + }() + return iter +} + +type turnLoopCheckpointStore struct { + m map[string][]byte + mu sync.Mutex +} + +func (s *turnLoopCheckpointStore) Set(_ context.Context, key string, value []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + s.m[key] = value + return nil +} + +func (s *turnLoopCheckpointStore) Get(_ context.Context, key string) ([]byte, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.m[key] + return v, ok, nil +} + +type turnLoopCancellableMockAgent struct { + name string + runFunc func(ctx context.Context, input *AgentInput) (*AgentOutput, error) + onCancel func(cc *cancelContext) + cancel context.CancelFunc + mu sync.Mutex +} + +func (a *turnLoopCancellableMockAgent) Name(_ context.Context) string { return a.name } +func (a *turnLoopCancellableMockAgent) Description(_ context.Context) string { return "mock agent" } + +func (a *turnLoopCancellableMockAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + + o := getCommonOptions(nil, opts...) + cc := o.cancelCtx + + a.mu.Lock() + var cancelCtx context.Context + cancelCtx, a.cancel = context.WithCancel(ctx) + a.mu.Unlock() + + go func() { + defer gen.Close() + if cc != nil { + go func() { + <-cc.cancelChan + // CRITICAL: call onCancel BEFORE cancel() to avoid race condition. + // If cancel() fires first, the runFunc returns immediately, + // flowAgent's defer calls markDone(), and doneChan closes + // before onCancel can read cc.config. + if a.onCancel != nil { + a.onCancel(cc) + } + a.mu.Lock() + if a.cancel != nil { + a.cancel() + } + a.mu.Unlock() + }() + } + + output, err := a.runFunc(cancelCtx, input) + if err != nil { + gen.Send(&AgentEvent{Err: err}) + return + } + gen.Send(&AgentEvent{Output: output}) + }() + return iter +} + +type turnLoopStopModeProbeAgent struct { + ccCh chan *cancelContext +} + +func (a *turnLoopStopModeProbeAgent) Name(_ context.Context) string { return "probe" } +func (a *turnLoopStopModeProbeAgent) Description(_ context.Context) string { return "probe" } +func (a *turnLoopStopModeProbeAgent) Run(_ context.Context, _ *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + o := getCommonOptions(nil, opts...) + cc := o.cancelCtx + a.ccCh <- cc + go func() { + defer gen.Close() + <-cc.cancelChan + for { + if cc.getMode() == CancelImmediate { + gen.Send(&AgentEvent{Err: cc.createCancelError()}) + return + } + time.Sleep(1 * time.Millisecond) + } + }() + return iter +} + +func newAndRunTurnLoop[T any, M MessageType](ctx context.Context, cfg TurnLoopConfig[T, M]) *TurnLoop[T, M] { + l := NewTurnLoop(cfg) + l.Run(ctx) + return l +} + +func genInputConsumeAll(_ context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil +} + +func genInputConsumeAllWithMsg(_ context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil +} + +func genInputConsumeFirst(_ context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: []string{items[0]}, + Remaining: items[1:], + }, nil +} + +var prepareTestAgent = func(_ context.Context, _ *TurnLoop[string, *schema.Message], _ []string) (Agent, error) { + return &turnLoopMockAgent{name: "test"}, nil +} + +func prepareAgent(a Agent) func(context.Context, *TurnLoop[string, *schema.Message], []string) (Agent, error) { + return func(_ context.Context, _ *TurnLoop[string, *schema.Message], _ []string) (Agent, error) { + return a, nil + } +} + +func waitOrFail(t *testing.T, ch <-chan struct{}, msg string) { + t.Helper() + select { + case <-ch: + case <-time.After(2 * time.Second): + t.Fatal(msg) + } +} + +func newTestStore() *turnLoopCheckpointStore { + return &turnLoopCheckpointStore{m: make(map[string][]byte)} +} + +func newPreemptTestLoop(t *testing.T, agent *turnLoopCancellableMockAgent) *TurnLoop[string, *schema.Message] { + t.Helper() + + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + + originalRunFunc := agent.runFunc + agent.runFunc = func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { close(agentStarted) }) + return originalRunFunc(ctx, input) + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: genInputConsumeFirst, + }) + + loop.Push("first") + + waitOrFail(t, agentStarted, "agent did not start") + + return loop +} + +func TestTurnLoop_StopIsIdempotent(t *testing.T) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Stop() + loop.Stop() + loop.Stop() + + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_WaitMultipleGoroutines(t *testing.T) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Stop() + + var wg sync.WaitGroup + results := make([]*TurnLoopExitState[string, *schema.Message], 3) + + for i := 0; i < 3; i++ { + i := i + wg.Add(1) + go func() { + defer wg.Done() + results[i] = loop.Wait() + }() + } + + wg.Wait() + + assert.Equal(t, results[0], results[1]) + assert.Equal(t, results[1], results[2]) +} + +func TestTurnLoop_UnhandledItemsOnStop(t *testing.T) { + started := make(chan struct{}) + blocked := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + close(started) + <-blocked + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items[:1], + Remaining: items[1:], + }, nil + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + loop.Push("msg2") + loop.Push("msg3") + + <-started + + loop.Stop() + close(blocked) + + result := loop.Wait() + assert.NotEmpty(t, result.UnhandledItems, "should return unhandled items") +} + +func TestTurnLoop_GenInputError(t *testing.T) { + genErr := errors.New("gen input error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return nil, genErr + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, genErr) +} + +func TestTurnLoop_GetAgentError(t *testing.T) { + agentErr := errors.New("get agent error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return nil, agentErr + }, + }) + + loop.Push("msg1") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, agentErr) +} + +func TestTurnLoop_BatchProcessing(t *testing.T) { + var batches [][]string + var mu sync.Mutex + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + mu.Lock() + batches = append(batches, items) + mu.Unlock() + + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items[:1], + Remaining: items[1:], + }, nil + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + loop.Push("msg2") + loop.Push("msg3") + + time.Sleep(200 * time.Millisecond) + + loop.Stop() + loop.Wait() + + mu.Lock() + defer mu.Unlock() + + assert.NotEmpty(t, batches, "should have processed at least one batch") +} + +func TestTurnLoop_StopWithMode(t *testing.T) { + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Stop(WithGraceful()) + + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_Preempt_CancelsCurrentAgent(t *testing.T) { + agentStarted := make(chan struct{}) + agentCancelled := make(chan struct{}) + agentStartedOnce := sync.Once{} + agentCancelledOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + agentCancelledOnce.Do(func() { + close(agentCancelled) + }) + return &AgentOutput{}, nil + }, + } + + genInputCalls := int32(0) + secondGenInputCalled := make(chan struct{}) + secondGenInputOnce := sync.Once{} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + count := atomic.AddInt32(&genInputCalls, 1) + if count >= 2 { + secondGenInputOnce.Do(func() { + close(secondGenInputCalled) + }) + } + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: []string{items[0]}, + Remaining: items[1:], + }, nil + }, + }) + + loop.Push("first") + + waitOrFail(t, agentStarted, "agent did not start") + + loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) + + waitOrFail(t, agentCancelled, "agent was not cancelled by preempt") + + waitOrFail(t, secondGenInputCalled, "second GenInput was not called after preempt") + + loop.Stop(WithImmediate()) + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.GreaterOrEqual(t, atomic.LoadInt32(&genInputCalls), int32(2)) +} + +func TestTurnLoop_Preempt_DiscardsConsumedItems(t *testing.T) { + agentStarted := make(chan struct{}) + agentDone := make(chan struct{}) + agentStartedOnce := sync.Once{} + agentDoneOnce := sync.Once{} + firstAgentRun := true + var firstRunMu sync.Mutex + + genInputResults := make([][]string, 0) + var mu sync.Mutex + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + firstRunMu.Lock() + isFirst := firstAgentRun + firstAgentRun = false + firstRunMu.Unlock() + + if isFirst { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + } else { + agentDoneOnce.Do(func() { + close(agentDone) + }) + } + return &AgentOutput{}, nil + }, + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + mu.Lock() + genInputResults = append(genInputResults, items) + mu.Unlock() + + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: []string{items[0]}, + Remaining: items[1:], + }, nil + }, + }) + + loop.Push("first") + + waitOrFail(t, agentStarted, "agent did not start") + + loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) + + waitOrFail(t, agentDone, "second agent run did not complete") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + + mu.Lock() + defer mu.Unlock() + require.GreaterOrEqual(t, len(genInputResults), 2) + assert.NotContains(t, genInputResults[1], "first") + assert.Contains(t, genInputResults[1], "urgent") +} + +func TestTurnLoop_Preempt_WithAgentCancelMode(t *testing.T) { + cancelFuncCalled := make(chan struct{}) + cancelFuncCalledOnce := sync.Once{} + firstCancelModeUsed := CancelImmediate + var cancelModeMu sync.Mutex + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return &AgentOutput{}, nil + }, + onCancel: func(cc *cancelContext) { + cancelModeMu.Lock() + cancelFuncCalledOnce.Do(func() { + firstCancelModeUsed = cc.getMode() + close(cancelFuncCalled) + }) + cancelModeMu.Unlock() + }, + } + + loop := newPreemptTestLoop(t, agent) + + loop.Push("urgent", WithPreempt[string, *schema.Message](AfterToolCalls)) + + waitOrFail(t, cancelFuncCalled, "cancelFunc was not called by preempt") + + loop.Stop(WithImmediate()) + result := loop.Wait() + assert.NoError(t, result.ExitReason) + cancelModeMu.Lock() + actualMode := firstCancelModeUsed + cancelModeMu.Unlock() + assert.Equal(t, CancelAfterToolCalls, actualMode) +} + +func TestTurnLoop_PreemptAck_ClosesAfterCancelIsInitiated(t *testing.T) { + cancelObserved := make(chan struct{}) + agentFinishGate := make(chan struct{}) + cancelObservedOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + <-agentFinishGate + return &AgentOutput{}, nil + }, + onCancel: func(cc *cancelContext) { + cancelObservedOnce.Do(func() { close(cancelObserved) }) + }, + } + + loop := newPreemptTestLoop(t, agent) + + ok, ack := loop.Push("urgent", WithPreempt[string, *schema.Message](AfterToolCalls)) + assert.True(t, ok) + assert.NotNil(t, ack) + + waitOrFail(t, ack, "preempt ack was not closed") + + waitOrFail(t, cancelObserved, "cancel was not initiated") + + close(agentFinishGate) + + loop.Stop(WithImmediate()) + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_PreemptAck_ClosesImmediatelyIfLoopNotStarted(t *testing.T) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + ok, ack := loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) + assert.True(t, ok) + assert.NotNil(t, ack) + + waitOrFail(t, ack, "preempt ack was not closed") +} + +func TestTurnLoop_Preempt_EscalatesOnSecondPreempt(t *testing.T) { + firstCancelSeen := make(chan struct{}) + agentFinishGate := make(chan struct{}) + firstCancelOnce := sync.Once{} + + var ccPtr atomic.Value + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + <-agentFinishGate + return &AgentOutput{}, nil + }, + onCancel: func(cc *cancelContext) { + ccPtr.Store(cc) + firstCancelOnce.Do(func() { close(firstCancelSeen) }) + }, + } + + loop := newPreemptTestLoop(t, agent) + + loop.Push("urgent1", WithPreempt[string, *schema.Message](AfterChatModel)) + waitOrFail(t, firstCancelSeen, "first preempt did not trigger cancel") + + loop.Push("urgent2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) + + wantMode := CancelAfterChatModel | CancelAfterToolCalls + deadline := time.Now().Add(1 * time.Second) + for time.Now().Before(deadline) { + v := ccPtr.Load() + if v == nil { + time.Sleep(5 * time.Millisecond) + continue + } + cc := v.(*cancelContext) + if cc.getMode() == wantMode && atomic.LoadInt32(&cc.escalated) == 1 { + break + } + time.Sleep(5 * time.Millisecond) + } + + v := ccPtr.Load() + if v == nil { + t.Fatal("cancel context was not captured") + } + cc := v.(*cancelContext) + assert.Equal(t, wantMode, cc.getMode()) + assert.Equal(t, int32(1), atomic.LoadInt32(&cc.escalated)) + + close(agentFinishGate) + + loop.Stop(WithImmediate()) + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_Preempt_JoinsSafePointModesOnSecondPreempt(t *testing.T) { + firstCancelSeen := make(chan struct{}) + agentFinishGate := make(chan struct{}) + firstCancelOnce := sync.Once{} + + var ccPtr atomic.Value + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + <-agentFinishGate + return &AgentOutput{}, nil + }, + onCancel: func(cc *cancelContext) { + ccPtr.Store(cc) + firstCancelOnce.Do(func() { close(firstCancelSeen) }) + }, + } + + loop := newPreemptTestLoop(t, agent) + + loop.Push("urgent1", WithPreempt[string, *schema.Message](AfterChatModel)) + waitOrFail(t, firstCancelSeen, "first preempt did not trigger cancel") + + loop.Push("urgent2", WithPreempt[string, *schema.Message](AfterToolCalls)) + + want := CancelAfterChatModel | CancelAfterToolCalls + deadline := time.Now().Add(1 * time.Second) + for time.Now().Before(deadline) { + v := ccPtr.Load() + if v == nil { + time.Sleep(5 * time.Millisecond) + continue + } + cc := v.(*cancelContext) + if cc.getMode() == want { + break + } + time.Sleep(5 * time.Millisecond) + } + + v := ccPtr.Load() + if v == nil { + t.Fatal("cancel context was not captured") + } + cc := v.(*cancelContext) + assert.Equal(t, want, cc.getMode()) + + close(agentFinishGate) + + loop.Stop(WithImmediate()) + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_Push_WithoutPreempt_DoesNotCancel(t *testing.T) { + agentRunCount := 0 + agentDone := make(chan struct{}) + + agent := &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentRunCount++ + if agentRunCount == 1 { + time.Sleep(100 * time.Millisecond) + } + if agentRunCount == 2 { + close(agentDone) + } + return &AgentOutput{}, nil + }, + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: genInputConsumeFirst, + }) + + loop.Push("first") + time.Sleep(20 * time.Millisecond) + loop.Push("second") + + waitOrFail(t, agentDone, "second agent run did not complete") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.Equal(t, 2, agentRunCount) +} + +func TestTurnLoop_PreemptDelay_NoMispreemptOnNaturalCompletion(t *testing.T) { + agent1Started := make(chan struct{}) + allowAgent1Done := make(chan struct{}) + agent1Done := make(chan struct{}) + agent2Started := make(chan struct{}) + allowAgent2Done := make(chan struct{}) + agent2Done := make(chan struct{}) + agent1StartedOnce := sync.Once{} + agent1DoneOnce := sync.Once{} + agent2StartedOnce := sync.Once{} + agent2DoneOnce := sync.Once{} + + var agentRunCount int32 + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + count := atomic.AddInt32(&agentRunCount, 1) + switch count { + case 1: + agent1StartedOnce.Do(func() { close(agent1Started) }) + <-allowAgent1Done + agent1DoneOnce.Do(func() { close(agent1Done) }) + case 2: + agent2StartedOnce.Do(func() { close(agent2Started) }) + select { + case <-allowAgent2Done: + case <-ctx.Done(): + t.Error("Agent2 was unexpectedly cancelled") + return nil, ctx.Err() + } + agent2DoneOnce.Do(func() { close(agent2Done) }) + } + return &AgentOutput{}, nil + }, + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: genInputConsumeFirst, + }) + + loop.Push("first") + + waitOrFail(t, agent1Started, "agent1 did not start") + + _, ack := loop.Push("second", WithPreempt[string, *schema.Message](AnySafePoint), WithPreemptDelay[string, *schema.Message](500*time.Millisecond)) + require.NotNil(t, ack) + close(allowAgent1Done) + + waitOrFail(t, agent1Done, "agent1 did not complete naturally") + + waitOrFail(t, agent2Started, "agent2 did not start") + + requireAckClosed(t, ack) + close(allowAgent2Done) + + waitOrFail(t, agent2Done, "agent2 did not complete - may have been incorrectly preempted") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.Equal(t, int32(2), atomic.LoadInt32(&agentRunCount)) +} + +func TestTurnLoop_PreemptDuringPlanningCancelsUpcomingAgent(t *testing.T) { + genInputEntered := make(chan struct{}) + allowGenInput := make(chan struct{}) + cancelSubmitted := make(chan struct{}) + agentCancelled := make(chan struct{}) + genInputOnce := sync.Once{} + cancelOnce := sync.Once{} + agentCancelledOnce := sync.Once{} + + var agentRunCount int32 + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + count := atomic.AddInt32(&agentRunCount, 1) + if count == 1 { + <-ctx.Done() + agentCancelledOnce.Do(func() { close(agentCancelled) }) + return &AgentOutput{}, nil + } + return &AgentOutput{}, nil + }, + onCancel: func(cc *cancelContext) { + cancelOnce.Do(func() { close(cancelSubmitted) }) + }, + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + genInputOnce.Do(func() { + close(genInputEntered) + <-allowGenInput + }) + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: []string{items[0]}, + Remaining: items[1:], + }, nil + }, + }) + + loop.Push("first") + + waitOrFail(t, genInputEntered, "GenInput did not enter planning phase") + + ok, ack := loop.Push("urgent", WithPreempt[string, *schema.Message](AnySafePoint)) + require.True(t, ok) + require.NotNil(t, ack) + requireAckOpen(t, ack) + + close(allowGenInput) + + waitOrFail(t, cancelSubmitted, "planning-phase preempt did not submit cancel") + requireAckClosed(t, ack) + + waitOrFail(t, agentCancelled, "agent was not cancelled by planning-phase preempt") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.GreaterOrEqual(t, atomic.LoadInt32(&agentRunCount), int32(1)) +} + +func TestTurnLoop_ConcurrentPush(t *testing.T) { + var count int32 + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + atomic.AddInt32(&count, int32(len(items))) + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + }) + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + for j := 0; j < 10; j++ { + _, _ = loop.Push(fmt.Sprintf("msg-%d-%d", i, j)) + } + }(i) + } + + wg.Wait() + time.Sleep(200 * time.Millisecond) + + loop.Stop() + result := loop.Wait() + + processed := atomic.LoadInt32(&count) + unhandled := len(result.UnhandledItems) + + assert.True(t, processed > 0, "should have processed some items") + assert.True(t, int(processed)+unhandled <= 100, "total should not exceed pushed amount") +} + +func TestTurnLoop_StopAfterReceive_RecoverItem(t *testing.T) { + receiveStarted := make(chan struct{}) + cancelDone := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + close(receiveStarted) + <-cancelDone + time.Sleep(50 * time.Millisecond) + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + <-receiveStarted + + loop.Stop() + close(cancelDone) + + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_StopAfterGenInput_RecoverConsumed(t *testing.T) { + genInputDone := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + close(genInputDone) + time.Sleep(50 * time.Millisecond) + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items[:1], + Remaining: items[1:], + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + time.Sleep(100 * time.Millisecond) + return &turnLoopMockAgent{name: "test"}, nil + }, + }) + + loop.Push("msg1") + loop.Push("msg2") + + <-genInputDone + + time.Sleep(60 * time.Millisecond) + loop.Stop() + + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_GetAgentError_RecoverConsumed(t *testing.T) { + agentErr := errors.New("get agent error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items[:1], + Remaining: items[1:], + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { + return nil, agentErr + }, + }) + + loop.Push("msg1") + loop.Push("msg2") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, agentErr) + assert.NotEmpty(t, result.UnhandledItems, "should recover at least the consumed item and remaining") +} + +func TestTurnLoop_GenInputError_RecoverItems(t *testing.T) { + genErr := errors.New("gen input error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return nil, genErr + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + loop.Push("msg2") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, genErr) + assert.Len(t, result.UnhandledItems, 2, "should recover all items when GenInput fails") + assert.Contains(t, result.UnhandledItems, "msg1") + assert.Contains(t, result.UnhandledItems, "msg2") +} + +func TestTurnLoop_PrepareAgentError_RecoverItemsInOrder(t *testing.T) { + agentErr := errors.New("prepare agent error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + var urgent string + remaining := make([]string, 0, len(items)) + for _, item := range items { + if item == "urgent" { + urgent = item + } else { + remaining = append(remaining, item) + } + } + if urgent != "" { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: []string{urgent}, + Remaining: remaining, + }, nil + } + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items[:1], + Remaining: items[1:], + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return nil, agentErr + }, + }) + + loop.Push("msg1") + loop.Push("urgent") + loop.Push("msg2") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, agentErr) + assert.Len(t, result.UnhandledItems, 3, "should recover all items") + assert.Equal(t, []string{"msg1", "urgent", "msg2"}, result.UnhandledItems, + "should preserve original push order even when GenInput selects non-prefix items") +} + +// Context cancel tests: the TurnLoop monitors context cancellation by closing +// the internal buffer when ctx.Done() fires, which unblocks the blocking +// Receive() call. The loop then checks ctx.Err() and exits with the context error. + +func TestTurnLoop_ContextCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + genInputStarted := make(chan struct{}) + genInputDone := make(chan struct{}) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + close(genInputStarted) + <-genInputDone + if err := ctx.Err(); err != nil { + return nil, err + } + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + + <-genInputStarted + cancel() + close(genInputDone) + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, context.Canceled) +} + +func TestTurnLoop_ContextDeadlineExceeded(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + select { + case <-time.After(100 * time.Millisecond): + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, context.DeadlineExceeded) +} + +func TestTurnLoop_ContextCancelBeforeReceive(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + // Push before Run to guarantee the item is buffered before the + // context-monitoring goroutine can close the buffer. + _, _ = loop.Push("msg1") + loop.Run(ctx) + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, context.Canceled) + assert.Len(t, result.UnhandledItems, 1) +} + +func TestTurnLoop_ContextCancelDuringBlockingReceive(t *testing.T) { + // When context is cancelled while Receive() is blocking (no items in buffer), + // the context monitoring goroutine closes the buffer, which unblocks Receive(). + ctx, cancel := context.WithCancel(context.Background()) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + // Don't push any items — let Receive() block + time.Sleep(50 * time.Millisecond) + cancel() + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, context.Canceled) +} + +func TestTurnLoop_ContextCancelAfterGenInput_RecoverItems(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + + genInputCount := 0 + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + genInputCount++ + if genInputCount == 1 { + cancel() + } + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items[:1], + Remaining: items[1:], + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], c []string) (Agent, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + return &turnLoopMockAgent{name: "test"}, nil + }, + }) + + loop.Push("msg1") + loop.Push("msg2") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, context.Canceled) + assert.NotEmpty(t, result.UnhandledItems, "should recover consumed and remaining items") +} + +func TestTurnLoop_OnAgentEventsReceivesEvents(t *testing.T) { + var receivedEvents []*AgentEvent + var receivedConsumed []string + var mu sync.Mutex + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + mu.Lock() + receivedConsumed = append(receivedConsumed, tc.Consumed...) + mu.Unlock() + + for { + event, ok := events.Next() + if !ok { + break + } + mu.Lock() + receivedEvents = append(receivedEvents, event) + mu.Unlock() + } + return nil + }, + }) + + loop.Push("msg1") + + time.Sleep(100 * time.Millisecond) + + loop.Stop() + result := loop.Wait() + + assert.NoError(t, result.ExitReason) + + mu.Lock() + defer mu.Unlock() + assert.NotEmpty(t, receivedConsumed, "should have received consumed items") +} + +// TestTurnLoop_BareStop_AgentRunsToCompletion verifies the core contract of +// bare Stop(): the running agent finishes naturally with an uncanceled context, +// the loop exits cleanly (ExitReason == nil), and no new turn starts even when +// additional items are buffered. +func TestTurnLoop_BareStop_AgentRunsToCompletion(t *testing.T) { + const agentWorkDuration = 200 * time.Millisecond + + agentStarted := make(chan struct{}) + agentCtxErr := make(chan error, 1) + agentOutput := make(chan string, 1) + + turnsExecuted := int32(0) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeFirst, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "worker", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + atomic.AddInt32(&turnsExecuted, 1) + close(agentStarted) + + // Simulate real work (NOT blocking on <-ctx.Done()) + time.Sleep(agentWorkDuration) + + // Record context state AFTER work completes + agentCtxErr <- ctx.Err() + agentOutput <- "work-done" + + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + // Push two items so the loop has a reason to start a second turn. + loop.Push("task1") + loop.Push("task2") + + // Wait for the agent to start processing task1. + waitOrFail(t, agentStarted, "agent did not start") + + // Call bare Stop() while the agent is doing work. + loop.Stop() + + result := loop.Wait() + + // 1. Agent's context was NOT canceled. + select { + case err := <-agentCtxErr: + assert.NoError(t, err, "bare Stop must not cancel the agent's context") + default: + t.Fatal("agent never reported context state") + } + + // 2. Agent completed its work. + select { + case out := <-agentOutput: + assert.Equal(t, "work-done", out) + default: + t.Fatal("agent never produced output") + } + + // 3. ExitReason is nil (clean exit, not a CancelError). + assert.NoError(t, result.ExitReason) + + // 4. InterruptedItems is empty (agent was not interrupted). + assert.Empty(t, result.InterruptedItems) + + // 5. Only one turn executed; the second item is unhandled. + assert.Equal(t, int32(1), atomic.LoadInt32(&turnsExecuted), + "bare Stop must prevent new turns from starting after the current one completes") + assert.Equal(t, []string{"task2"}, result.UnhandledItems, + "the second item should appear in UnhandledItems") +} + +func TestTurnLoop_StopCheckPointIDInCancelError(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + checkpointID := "turn-loop-cancel-ckpt-1" + store := newTestStore() + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: checkpointID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + + loop.Push("msg1") + + <-modelStarted + loop.Stop(WithImmediate()) + + result := loop.Wait() + + var cancelErr *CancelError + assert.True(t, errors.As(result.ExitReason, &cancelErr), "ExitReason should be a *CancelError") + + store.mu.Lock() + defer store.mu.Unlock() + _, ok := store.m[checkpointID] + assert.True(t, ok, "checkpoint should be saved under the configured CheckpointID") +} + +// TestTurnLoop_CancelError_CapturedIndependentlyOfCallback verifies that the TurnLoop +// correctly reports *CancelError as ExitReason and populates InterruptedItems even when +// the user's custom OnAgentEvents callback swallows the CancelError (returns nil). +// This tests the documented guarantee: "the callback should NEVER propagate CancelError +// — the framework handles it automatically." +func TestTurnLoop_CancelError_CapturedIndependentlyOfCallback(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + checkpointID := "cancel-capture-independent-1" + store := newTestStore() + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: checkpointID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + // Custom OnAgentEvents that deliberately swallows all errors including CancelError. + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*TypedAgentEvent[*schema.Message]]) error { + for { + _, ok := events.Next() + if !ok { + break + } + // Deliberately ignore event.Err — do NOT propagate CancelError. + } + return nil // swallow everything + }, + }) + + loop.Push("msg1") + + <-modelStarted + loop.Stop(WithImmediate()) + + result := loop.Wait() + + // The framework should capture CancelError independently of the callback's return value. + var cancelErr *CancelError + assert.True(t, errors.As(result.ExitReason, &cancelErr), + "ExitReason should be *CancelError even when OnAgentEvents swallows it, got: %v", result.ExitReason) + + // InterruptedItems should be populated. + assert.Equal(t, []string{"msg1"}, result.InterruptedItems, + "InterruptedItems should contain the items that were being processed") + + // Checkpoint should be saved. + store.mu.Lock() + defer store.mu.Unlock() + _, ok := store.m[checkpointID] + assert.True(t, ok, "checkpoint should be saved under the configured CheckpointID") +} + +// TestTurnLoop_CancelError_CustomErrorWins_InterruptedItemsStillSet verifies that when +// the user's OnAgentEvents callback returns a custom error during a cancel, the custom +// error becomes ExitReason (not overwritten by CancelError), but InterruptedItems is still +// populated because the items were factually mid-execution when the cancel signal arrived. +func TestTurnLoop_CancelError_CustomErrorWins_InterruptedItemsStillSet(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + checkpointID := "cancel-custom-error-wins-1" + store := newTestStore() + customErr := fmt.Errorf("user callback encountered a problem") + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: checkpointID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + // Custom OnAgentEvents that returns a custom error instead of the CancelError. + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*TypedAgentEvent[*schema.Message]]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + return customErr + }, + }) + + loop.Push("msg1") + + <-modelStarted + loop.Stop(WithImmediate()) + + result := loop.Wait() + + // User's custom error should win as ExitReason. + assert.ErrorIs(t, result.ExitReason, customErr, + "ExitReason should be the user's custom error, not CancelError") + + // But InterruptedItems should still be populated (items were factually in-flight). + assert.Equal(t, []string{"msg1"}, result.InterruptedItems, + "InterruptedItems should contain the items that were being processed") + + // Checkpoint should be saved (cancel was captured, items were in-flight). + store.mu.Lock() + defer store.mu.Unlock() + _, ok := store.m[checkpointID] + assert.True(t, ok, "checkpoint should be saved even when user returns custom error") +} + +func TestTurnLoop_StopWithoutCheckpointIDDoesNotPersist(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + store := newTestStore() + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + + loop.Push("msg1") + + <-modelStarted + loop.Stop(WithImmediate()) + + result := loop.Wait() + + var cancelErr *CancelError + assert.True(t, errors.As(result.ExitReason, &cancelErr), "ExitReason should be a *CancelError") + + store.mu.Lock() + defer store.mu.Unlock() + assert.Empty(t, store.m, "no checkpoint should be saved when CheckpointID is not configured") +} + +func TestTurnLoop_StopWhileIdle_SkipsCheckpoint(t *testing.T) { + ctx := context.Background() + store := &deletableCheckpointStore{ + turnLoopCheckpointStore: turnLoopCheckpointStore{m: make(map[string][]byte)}, + } + cpID := "idle-session" + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Stop() + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + + store.mu.Lock() + defer store.mu.Unlock() + _, exists := store.m[cpID] + assert.False(t, exists, "no checkpoint should be saved when TurnLoop is idle") +} + +func TestTurnLoop_StopBetweenTurnsAndResume(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "between-turns-session" + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("a") + loop.Push("b") + loop.Stop() + loop.Run(ctx) + + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + + var seen []string + var mu sync.Mutex + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + mu.Lock() + seen = append([]string{}, items...) + mu.Unlock() + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + + loop2.Push("c") + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.NoError(t, exit2.ExitReason) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{"a", "b", "c"}, seen) +} + +// TestTurnLoop_StopCancel_InterruptedItems_PersistAndRestore verifies the full +// lifecycle: Stop(WithImmediate()) cancels a running agent, the exit state +// reports the correct InterruptedItems, the checkpoint persists them, and a new +// TurnLoop restoring from that checkpoint passes the exact same items to +// GenResume's interruptedItems parameter. +func TestTurnLoop_StopCancel_InterruptedItems_PersistAndRestore(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + store := newTestStore() + cpID := "inflight-persist-session" + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + require.NoError(t, err) + + // Phase 1: Run, cancel mid-turn, verify exit state. + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + + loop.Push("msg1") + <-modelStarted + loop.Stop(WithImmediate()) + exit := loop.Wait() + + // ExitReason must be CancelError. + var cancelErr *CancelError + require.True(t, errors.As(exit.ExitReason, &cancelErr), "expected *CancelError, got: %v", exit.ExitReason) + + // InterruptedItems must contain the consumed items. + assert.Equal(t, []string{"msg1"}, exit.InterruptedItems, + "InterruptedItems in exit state should contain the items that were mid-execution") + + // Checkpoint must be saved. + assert.True(t, exit.CheckpointAttempted, "checkpoint should be attempted") + assert.NoError(t, exit.CheckpointErr, "checkpoint save should succeed") + + store.mu.Lock() + _, cpExists := store.m[cpID] + store.mu.Unlock() + require.True(t, cpExists, "checkpoint should exist in store") + + // Phase 2: Restore from checkpoint, verify GenResume receives the exact InterruptedItems. + slowModel.setDelay(10 * time.Millisecond) + + var genResumeCalled bool + var resumeInterruptedItems []string + var resumeUnhandledItems []string + + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], interruptedItems, unhandledItems, newItems []string) (*GenResumeResult[string, *schema.Message], error) { + genResumeCalled = true + resumeInterruptedItems = append([]string{}, interruptedItems...) + resumeUnhandledItems = append([]string{}, unhandledItems...) + return &GenResumeResult[string, *schema.Message]{ + Consumed: interruptedItems, + Remaining: append(append([]string{}, unhandledItems...), newItems...), + }, nil + }, + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + t.Fatal("GenInput should not be called when resuming from a cancel checkpoint") + return nil, nil + }, + PrepareAgent: prepareAgent(agent), + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + + loop2.Run(ctx) + exit2 := loop2.Wait() + + assert.NoError(t, exit2.ExitReason) + assert.True(t, genResumeCalled, "GenResume should be called when restoring from a cancel checkpoint") + assert.Equal(t, []string{"msg1"}, resumeInterruptedItems, + "GenResume's interruptedItems should match the original exit state's InterruptedItems") + assert.Empty(t, resumeUnhandledItems, + "unhandledItems should be empty (all items were consumed before cancel)") +} + +// TestTurnLoop_PreemptThenStop_InterruptedItems_ReflectsStoppedTurn verifies that +// when a preempt interrupts turn 1 and then Stop(WithImmediate()) cancels turn 2, +// the reported InterruptedItems correspond to turn 2's consumed items (not turn 1's). +func TestTurnLoop_PreemptThenStop_InterruptedItems_ReflectsStoppedTurn(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "preempt-then-stop-session" + + turnCount := int32(0) + modelStarted := make(chan struct{}, 10) + + slowModel := &cancelTestChatModel{ + delayNs: int64(5 * time.Second), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 10), + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + require.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + n := atomic.AddInt32(&turnCount, 1) + _ = n + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: prepareAgent(agent), + }) + + // Turn 1: push "a", wait for agent to start, then preempt with "b". + loop.Push("a") + <-modelStarted + + _, ack := loop.Push("b", WithPreemptTimeout[string, *schema.Message](AnySafePoint, 10*time.Millisecond)) + // Wait for the preempt to be acknowledged. + waitOrFail(t, ack, "preempt ack timed out") + + // Turn 2: the loop restarts with all buffered items. Wait for agent to start. + <-modelStarted + + // Now stop with immediate cancel on turn 2. + loop.Stop(WithImmediate()) + exit := loop.Wait() + + // ExitReason must be CancelError. + var cancelErr *CancelError + require.True(t, errors.As(exit.ExitReason, &cancelErr), "expected *CancelError, got: %v", exit.ExitReason) + + // KEY ASSERTION: InterruptedItems must reflect turn 2's consumed items (which + // includes both "a" and "b" since GenInput consumes all items), NOT turn 1's ["a"]. + // Turn 2 consumed whatever GenInput received (the re-buffered items from preempt + "b"). + assert.NotEmpty(t, exit.InterruptedItems, "InterruptedItems should not be empty after Stop cancel") + assert.Contains(t, exit.InterruptedItems, "b", + "InterruptedItems should contain 'b' which was part of turn 2's consumed items") +} + +func TestTurnLoop_BusinessInterrupt_PersistAndResume(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "interrupt-session" + + // Agent that produces a business interrupt via Interrupt() call. + interruptAgent := &turnLoopInterruptAgent{interruptInfo: "approval_needed"} + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return interruptAgent, nil + }, + }) + + loop.Push("msg1") + exit := loop.Wait() + + // 1. ExitReason is an *InterruptError (not nil, not *CancelError). + var intErr *InterruptError + require.True(t, errors.As(exit.ExitReason, &intErr), "expected *InterruptError, got: %v", exit.ExitReason) + + // 2. InterruptContexts is populated. + require.NotEmpty(t, intErr.InterruptContexts) + + // 3. InterruptedItems contains the items being processed. + assert.Equal(t, []string{"msg1"}, exit.InterruptedItems) + + // 4. Checkpoint was persisted. + assert.True(t, exit.CheckpointAttempted) + assert.NoError(t, exit.CheckpointErr) + + store.mu.Lock() + _, cpExists := store.m[cpID] + store.mu.Unlock() + assert.True(t, cpExists, "checkpoint should exist in store") + + // 5. Resume: new TurnLoop with same CheckpointID gets GenResume called. + var genResumeCalled bool + var resumeInterruptedItems []string + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], interruptedItems []string, unhandledItems []string, newItems []string) (*GenResumeResult[string, *schema.Message], error) { + genResumeCalled = true + resumeInterruptedItems = append([]string{}, interruptedItems...) + return &GenResumeResult[string, *schema.Message]{ + Consumed: interruptedItems, + Remaining: append(append([]string{}, unhandledItems...), newItems...), + }, nil + }, + GenInput: genInputConsumeAll, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + // On resume, agent completes normally. + return &turnLoopMockAgent{ + name: "ResumeAgent", + events: []*AgentEvent{{Output: &AgentOutput{}}}, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.NoError(t, exit2.ExitReason) + assert.True(t, genResumeCalled, "GenResume should be called on checkpoint resume") + assert.Equal(t, []string{"msg1"}, resumeInterruptedItems, "interruptedItems should contain the original items") +} + +// turnLoopInterruptAgent is a test agent that produces a business interrupt event. +type turnLoopInterruptAgent struct { + interruptInfo any +} + +func (a *turnLoopInterruptAgent) Name(_ context.Context) string { return "InterruptAgent" } +func (a *turnLoopInterruptAgent) Description(_ context.Context) string { + return "agent that interrupts" +} +func (a *turnLoopInterruptAgent) Run(ctx context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + event := Interrupt(ctx, a.interruptInfo) + gen.Send(event) + }() + return iter +} + +func TestTurnLoop_CheckpointIDWithoutStore_FreshStart(t *testing.T) { + ctx := context.Background() + var genInputCalled bool + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + CheckpointID: "some-id", + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + genInputCalled = true + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + loop.Push("a") + loop.Run(ctx) + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + assert.True(t, genInputCalled) +} + +func TestTurnLoop_CheckpointNotFound_FreshStart(t *testing.T) { + ctx := context.Background() + store := newTestStore() + var genInputCalled bool + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: "nonexistent-id", + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + genInputCalled = true + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + loop.Push("a") + loop.Run(ctx) + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + assert.True(t, genInputCalled) +} + +func TestTurnLoop_CheckpointEmptyData_TreatedAsNoCheckpoint(t *testing.T) { + ctx := context.Background() + store := newTestStore() + store.m["cp-empty"] = nil + + var genInputCalled bool + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: "cp-empty", + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + genInputCalled = true + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + loop.Push("a") + loop.Run(ctx) + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + assert.True(t, genInputCalled) +} + +type errorCheckpointStore struct { + getErr error + setErr error +} + +func (s *errorCheckpointStore) Get(_ context.Context, _ string) ([]byte, bool, error) { + return nil, false, s.getErr +} + +func (s *errorCheckpointStore) Set(_ context.Context, _ string, _ []byte) error { + return s.setErr +} + +func TestTurnLoop_CheckpointLoadError_ReturnsError(t *testing.T) { + ctx := context.Background() + store := &errorCheckpointStore{getErr: fmt.Errorf("store unavailable")} + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: "cp-1", + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Run(ctx) + exit := loop.Wait() + assert.Error(t, exit.ExitReason) + assert.Contains(t, exit.ExitReason.Error(), "store unavailable") +} + +func TestTurnLoop_CheckpointCorruptData_ReturnsError(t *testing.T) { + ctx := context.Background() + store := newTestStore() + store.m["cp-corrupt"] = []byte("not-valid-gob-data") + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: "cp-corrupt", + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Run(ctx) + exit := loop.Wait() + assert.Error(t, exit.ExitReason) + assert.Contains(t, exit.ExitReason.Error(), "failed to unmarshal checkpoint") +} + +func TestTurnLoop_CheckpointSaveError_ReturnsError(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + saveStore := &errorCheckpointStore{setErr: fmt.Errorf("write failed")} + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: saveStore, + CheckpointID: "cp-1", + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + loop.Push("msg1") + <-modelStarted + loop.Stop(WithImmediate()) + exit := loop.Wait() + assert.Error(t, exit.ExitReason) + assert.True(t, exit.CheckpointAttempted) + assert.Error(t, exit.CheckpointErr) + assert.Contains(t, exit.CheckpointErr.Error(), "write failed") +} + +func TestTurnLoop_StaleCheckpointDeletion_OnCleanResume(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "stale-session" + + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop1.Push("a") + loop1.Stop() + loop1.Run(ctx) + loop1.Wait() + + store.mu.Lock() + _, exists := store.m[cpID] + store.mu.Unlock() + assert.True(t, exists, "checkpoint should exist after first loop saves it") + + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + loop2.Push("b") + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.NoError(t, exit2.ExitReason) + + store.mu.Lock() + _, exists = store.m[cpID] + store.mu.Unlock() + assert.True(t, exists, "checkpoint should still exist because loop2 was stopped and saved a new one") +} + +type deletableCheckpointStore struct { + turnLoopCheckpointStore + deleteCalled bool + deletedKey string +} + +func (s *deletableCheckpointStore) Delete(_ context.Context, key string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.deleteCalled = true + s.deletedKey = key + delete(s.m, key) + return nil +} + +func TestTurnLoop_CheckpointDeleter_CalledOnContextCancel(t *testing.T) { + ctx := context.Background() + store := &deletableCheckpointStore{ + turnLoopCheckpointStore: turnLoopCheckpointStore{m: make(map[string][]byte)}, + } + cpID := "deleter-session" + + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop1.Push("a") + loop1.Stop() + loop1.Run(ctx) + loop1.Wait() + + store.mu.Lock() + _, exists := store.m[cpID] + store.mu.Unlock() + assert.True(t, exists, "checkpoint saved after loop1") + + ctx2, cancel2 := context.WithCancel(ctx) + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + cancel2() + return nil + }, + }) + loop2.Push("b") + loop2.Run(ctx2) + exit2 := loop2.Wait() + assert.ErrorIs(t, exit2.ExitReason, context.Canceled) + + store.mu.Lock() + defer store.mu.Unlock() + assert.True(t, store.deleteCalled, "CheckPointDeleter.Delete should be called") + assert.Equal(t, cpID, store.deletedKey) + _, exists = store.m[cpID] + assert.False(t, exists, "checkpoint should be removed from store") +} + +func TestTurnLoop_GenResumeNil_Error(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "resume-nil-session" + modelStarted := make(chan struct{}, 1) + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + loop1.Push("msg1") + <-modelStarted + loop1.Stop(WithImmediate()) + loop1.Wait() + + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.Error(t, exit2.ExitReason) + assert.Contains(t, exit2.ExitReason.Error(), "GenResume is required") +} + +func TestTurnLoop_SameCheckpointID_OverwritePattern(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "overwrite-session" + + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop1.Push("a") + loop1.Push("b") + loop1.Stop() + loop1.Run(ctx) + loop1.Wait() + + store.mu.Lock() + data1 := append([]byte{}, store.m[cpID]...) + store.mu.Unlock() + assert.NotEmpty(t, data1) + + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop2.Push("c") + loop2.Stop() + loop2.Run(ctx) + loop2.Wait() + + store.mu.Lock() + data2 := append([]byte{}, store.m[cpID]...) + store.mu.Unlock() + assert.NotEmpty(t, data2) + assert.NotEqual(t, data1, data2, "checkpoint data should change because items are different") + + var seen []string + var mu sync.Mutex + loop3 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + mu.Lock() + seen = append([]string{}, items...) + mu.Unlock() + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + loop3.Push("d") + loop3.Run(ctx) + exit3 := loop3.Wait() + assert.NoError(t, exit3.ExitReason) + + mu.Lock() + defer mu.Unlock() + assert.Equal(t, []string{"a", "b", "c", "d"}, seen, "should see loop2's unhandled items (a,b,c from loop2's checkpoint) plus new d") +} + +func TestTurnLoop_CheckpointHasRunnerStateButEmptyBytes(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "empty-runner-bytes" + + cp := &turnLoopCheckpoint[string]{ + HasRunnerState: true, + RunnerCheckpoint: nil, + UnhandledItems: []string{"x"}, + } + data, err := marshalTurnLoopCheckpoint(cp) + assert.NoError(t, err) + store.m[cpID] = data + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Run(ctx) + exit := loop.Wait() + assert.Error(t, exit.ExitReason) + assert.Contains(t, exit.ExitReason.Error(), "has runner state but bytes are empty") +} + +func TestTurnLoop_GenResumeReturnsError(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "resume-err-session" + modelStarted := make(chan struct{}, 1) + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + loop1.Push("msg1") + <-modelStarted + loop1.Stop(WithImmediate()) + loop1.Wait() + + genResumeErr := fmt.Errorf("resume callback failed") + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], canceled, unhandled, newItems []string) (*GenResumeResult[string, *schema.Message], error) { + return nil, genResumeErr + }, + PrepareAgent: prepareTestAgent, + }) + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.Error(t, exit2.ExitReason) + assert.ErrorIs(t, exit2.ExitReason, genResumeErr) +} + +func TestTurnLoop_ResumeWaitsForInFlightPushBeforePlanning(t *testing.T) { + ctx := context.Background() + resumeErr := errors.New("stop after observing resume inputs") + strategyEntered := make(chan struct{}) + allowStrategy := make(chan struct{}) + pushDone := make(chan struct{}) + genResumeCalled := make(chan struct{}) + + var resumeNewItems []string + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], interruptedItems, unhandledItems, newItems []string) (*GenResumeResult[string, *schema.Message], error) { + resumeNewItems = append([]string{}, newItems...) + close(genResumeCalled) + return nil, resumeErr + }, + PrepareAgent: prepareTestAgent, + }) + loop.pendingResume = &turnLoopPendingResume[string]{ + interrupted: []string{"interrupted"}, + newItems: []string{"pre-existing"}, + } + + go func() { + defer close(pushDone) + ok, ack := loop.Push("during-resume", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + close(strategyEntered) + <-allowStrategy + return nil + })) + assert.True(t, ok) + assert.Nil(t, ack) + }() + + waitOrFail(t, strategyEntered, "strategy did not enter") + + loop.Run(ctx) + + select { + case <-genResumeCalled: + t.Fatal("GenResume should wait for in-flight PushStrategy to finish") + default: + } + + close(allowStrategy) + waitOrFail(t, pushDone, "push did not finish") + + exit := loop.Wait() + assert.ErrorIs(t, exit.ExitReason, resumeErr) + assert.Equal(t, []string{"pre-existing", "during-resume"}, resumeNewItems) +} + +func TestTurnLoop_CheckpointSaveError_MergesWithExistingError(t *testing.T) { + ctx := context.Background() + modelStarted := make(chan struct{}, 1) + saveStore := &errorCheckpointStore{setErr: fmt.Errorf("disk full")} + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: saveStore, + CheckpointID: "cp-merge-err", + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + loop.Push("msg1") + <-modelStarted + loop.Stop(WithImmediate()) + exit := loop.Wait() + assert.Error(t, exit.ExitReason) + var ce *CancelError + assert.True(t, errors.As(exit.ExitReason, &ce), "ExitReason should be CancelError, not merged with checkpoint error") + assert.True(t, exit.CheckpointAttempted) + assert.Error(t, exit.CheckpointErr) + assert.Contains(t, exit.CheckpointErr.Error(), "disk full") +} + +func TestTurnLoop_ResumeWithParams(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "resume-params-session" + modelStarted := make(chan struct{}, 1) + + slowModel := &cancelTestChatModel{ + delayNs: int64(500 * time.Millisecond), + response: &schema.Message{ + Role: schema.Assistant, + Content: "Hello", + }, + startedChan: modelStarted, + doneChan: make(chan struct{}, 1), + } + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Instruction: "You are a test assistant", + Model: slowModel, + }) + assert.NoError(t, err) + + loop1 := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareAgent(agent), + }) + loop1.Push("msg1") + <-modelStarted + loop1.Stop(WithImmediate()) + exit1 := loop1.Wait() + var ce *CancelError + assert.True(t, errors.As(exit1.ExitReason, &ce)) + + var resumeParamsUsed *ResumeParams + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + GenResume: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], canceled, unhandled, newItems []string) (*GenResumeResult[string, *schema.Message], error) { + params := &ResumeParams{ + Targets: map[string]any{"some-address": "user-data"}, + } + resumeParamsUsed = params + return &GenResumeResult[string, *schema.Message]{ + ResumeParams: params, + Consumed: append(append(canceled, unhandled...), newItems...), + }, nil + }, + PrepareAgent: prepareAgent(agent), + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + }) + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.NotNil(t, resumeParamsUsed, "GenResume should have been called with ResumeParams") + assert.Contains(t, resumeParamsUsed.Targets, "some-address") + _ = exit2 +} + +func TestTurnLoop_ResumeInterruptAgain_PreservesEnableStreamingCheckpoint(t *testing.T) { + for _, enableStreaming := range []bool{true, false} { + t.Run(fmt.Sprintf("enable_streaming_%t", enableStreaming), func(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := fmt.Sprintf("streaming-resume-%t", enableStreaming) + originalMessage := "msg1" + + firstAgent := &myAgent{ + runFn: func(ctx context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.Equal(t, enableStreaming, input.EnableStreaming) + assert.Len(t, input.Messages, 1) + assert.Equal(t, originalMessage, input.Messages[0].Content) + + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(Interrupt(ctx, "first_interrupt")) + }() + return iter + }, + } + + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: func(_ context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{ + Messages: []Message{schema.UserMessage(items[0])}, + EnableStreaming: enableStreaming, + }, + Consumed: items, + }, nil + }, + PrepareAgent: prepareAgent(firstAgent), + }) + loop1.Push(originalMessage) + loop1.Run(ctx) + exit1 := loop1.Wait() + require.ErrorAs(t, exit1.ExitReason, new(*InterruptError)) + require.NoError(t, exit1.CheckpointErr) + + secondAgent := &myAgent{ + resumeFn: func(ctx context.Context, info *ResumeInfo, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.Equal(t, enableStreaming, info.EnableStreaming) + + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer gen.Close() + gen.Send(Interrupt(ctx, "second_interrupt")) + }() + return iter + }, + } + + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + GenResume: func(_ context.Context, _ *TurnLoop[string, *schema.Message], interrupted, unhandled, newItems []string) (*GenResumeResult[string, *schema.Message], error) { + return &GenResumeResult[string, *schema.Message]{ + Consumed: interrupted, + Remaining: append(append([]string{}, unhandled...), newItems...), + }, nil + }, + PrepareAgent: prepareAgent(secondAgent), + }) + loop2.Run(ctx) + exit2 := loop2.Wait() + require.ErrorAs(t, exit2.ExitReason, new(*InterruptError)) + require.NoError(t, exit2.CheckpointErr) + + // Verify the runner-level checkpoint persisted by the second interrupt + // still encodes the original streaming mode. This is the invariant the PR + // fixes: even though loop2's TypedRunner was constructed with the + // resume-path placeholder (false), runner uses resumeInfo.EnableStreaming + // from the previous checkpoint when re-saving. + store.mu.Lock() + data, ok := store.m[cpID] + store.mu.Unlock() + require.True(t, ok) + cp, err := unmarshalTurnLoopCheckpoint[string](data) + require.NoError(t, err) + require.True(t, cp.HasRunnerState) + _, _, info2, err := runnerLoadCheckPointImpl(newResumeBridgeStore(bridgeCheckpointID, cp.RunnerCheckpoint), context.Background(), bridgeCheckpointID) + require.NoError(t, err) + assert.Equal(t, enableStreaming, info2.EnableStreaming) + }) + } +} + +func TestTurnLoop_Stop_EscalatesCancelMode(t *testing.T) { + ctx := context.Background() + agentStarted := make(chan *cancelContext, 1) + probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return probe, nil + }, + }) + + loop.Push("msg1") + cc := <-agentStarted + + loop.Stop(WithGracefulTimeout(10 * time.Second)) + loop.Stop(WithImmediate()) + + deadline := time.After(1 * time.Second) + for { + if cc.getMode() == CancelImmediate { + break + } + select { + case <-deadline: + t.Fatal("cancel mode did not escalate to CancelImmediate") + default: + } + time.Sleep(1 * time.Millisecond) + } + + exit := loop.Wait() + var ce *CancelError + require.True(t, errors.As(exit.ExitReason, &ce)) + assert.Equal(t, CancelImmediate, ce.Info.Mode) +} + +func TestTurnLoop_DefaultOnAgentEvents_ErrorPropagation(t *testing.T) { + agentErr := errors.New("agent execution error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + return nil, agentErr + }, + }, nil + }, + // No OnAgentEvents — use default handler + }) + + loop.Push("msg1") + + result := loop.Wait() + // The default handler should propagate the agent error as ExitReason + assert.Error(t, result.ExitReason) +} + +func TestTurnLoop_OnAgentEventsError(t *testing.T) { + handlerErr := errors.New("event handler error") + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + // Drain events then return error + for { + _, ok := events.Next() + if !ok { + break + } + } + return handlerErr + }, + }) + + loop.Push("msg1") + + result := loop.Wait() + assert.ErrorIs(t, result.ExitReason, handlerErr) +} + +func TestTurnLoop_StopCallFromGenInput(t *testing.T) { + // Test that calling Stop() from within GenInput works correctly + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, loop *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + loop.Stop() + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + + result := loop.Wait() + assert.NoError(t, result.ExitReason) +} + +func TestTurnLoop_PushFromOnAgentEvents(t *testing.T) { + // Test that calling Push() from within OnAgentEvents works + pushCount := int32(0) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeFirst, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + count := atomic.AddInt32(&pushCount, 1) + if count == 1 { + // Push a follow-up item from the callback + _, _ = tc.Loop.Push("follow-up") + } else { + tc.Loop.Stop() + } + return nil + }, + }) + + loop.Push("initial") + + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.Equal(t, int32(2), atomic.LoadInt32(&pushCount)) +} + +// Tests for NewTurnLoop: the permissive API where Push, Stop, and Wait are +// all valid on a not-yet-running loop. + +func TestNewTurnLoop_PushBeforeRun(t *testing.T) { + // Items pushed before Run are buffered and processed after Run starts. + var processedItems []string + var mu sync.Mutex + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + mu.Lock() + processedItems = append(processedItems, items...) + mu.Unlock() + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: prepareTestAgent, + }) + + // Push before Run — items should be buffered. + ok, _ := loop.Push("msg1") + assert.True(t, ok) + ok, _ = loop.Push("msg2") + assert.True(t, ok) + + loop.Run(context.Background()) + + time.Sleep(100 * time.Millisecond) + + loop.Stop() + result := loop.Wait() + + mu.Lock() + defer mu.Unlock() + + assert.NoError(t, result.ExitReason) + assert.Contains(t, processedItems, "msg1") + assert.Contains(t, processedItems, "msg2") +} + +func TestNewTurnLoop_WaitBeforeRun(t *testing.T) { + // Wait blocks until Run is called AND the loop exits. + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + waitDone := make(chan *TurnLoopExitState[string, *schema.Message], 1) + go func() { + waitDone <- loop.Wait() + }() + + // Wait should not return yet since Run hasn't been called. + select { + case <-waitDone: + t.Fatal("Wait returned before Run was called") + case <-time.After(50 * time.Millisecond): + // expected + } + + loop.Push("msg1") + loop.Stop() + loop.Run(context.Background()) + + select { + case result := <-waitDone: + assert.NoError(t, result.ExitReason) + assert.Equal(t, []string{"msg1"}, result.UnhandledItems) + case <-time.After(1 * time.Second): + t.Fatal("Wait did not return after Run + Stop") + } +} + +func TestNewTurnLoop_RunIsIdempotent(t *testing.T) { + var genInputCalls int32 + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + atomic.AddInt32(&genInputCalls, 1) + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("msg1") + loop.Run(context.Background()) + loop.Run(context.Background()) + loop.Run(context.Background()) + + time.Sleep(100 * time.Millisecond) + + loop.Stop() + result := loop.Wait() + + assert.NoError(t, result.ExitReason) + assert.True(t, atomic.LoadInt32(&genInputCalls) >= 1) +} + +func TestNewTurnLoop_ConcurrentPushAndRun(t *testing.T) { + // Concurrent Push and Run should not race. + for i := 0; i < 100; i++ { + var count int32 + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + atomic.AddInt32(&count, int32(len(items))) + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{name: "test"}, nil + }, + }) + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + _, _ = loop.Push("item") + }() + + go func() { + defer wg.Done() + loop.Run(context.Background()) + }() + + wg.Wait() + + time.Sleep(50 * time.Millisecond) + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + + processed := atomic.LoadInt32(&count) + unhandled := len(result.UnhandledItems) + assert.True(t, int(processed)+unhandled <= 1, + "total should not exceed pushed amount") + } +} + +type turnCtxKey struct{} + +func TestTurnLoop_RunCtx_Propagation(t *testing.T) { + // Verify that GenInputResult.RunCtx is propagated to PrepareAgent, + // the agent run, and OnAgentEvents. + + const traceVal = "trace-123" + var prepareCtxVal, agentCtxVal, eventsCtxVal string + + cfg := TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, loop *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + // Derive a new context with per-item trace data + runCtx := context.WithValue(ctx, turnCtxKey{}, traceVal) + return &GenInputResult[string, *schema.Message]{ + RunCtx: runCtx, + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, loop *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + if v, ok := ctx.Value(turnCtxKey{}).(string); ok { + prepareCtxVal = v + } + return &turnLoopMockAgent{ + name: "trace-agent", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + if v, ok := ctx.Value(turnCtxKey{}).(string); ok { + agentCtxVal = v + } + return &AgentOutput{}, nil + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + if v, ok := ctx.Value(turnCtxKey{}).(string); ok { + eventsCtxVal = v + } + for { + if _, ok := events.Next(); !ok { + break + } + } + tc.Loop.Stop() + return nil + }, + } + + loop := NewTurnLoop(cfg) + loop.Push("hello") + loop.Run(context.Background()) + result := loop.Wait() + + assert.Nil(t, result.ExitReason) + assert.Equal(t, traceVal, prepareCtxVal, "PrepareAgent should receive RunCtx") + assert.Equal(t, traceVal, agentCtxVal, "Agent run should receive RunCtx") + assert.Equal(t, traceVal, eventsCtxVal, "OnAgentEvents should receive RunCtx") +} + +func TestTurnLoop_TurnContext_PreemptedChannel(t *testing.T) { + preemptedSeen := make(chan struct{}) + agentStarted := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + close(agentStarted) + select { + case <-tc.Preempted: + close(preemptedSeen) + case <-time.After(5 * time.Second): + t.Error("timed out waiting for Preempted channel") + } + // Drain events + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Push("msg2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) + + select { + case <-preemptedSeen: + // success + case <-time.After(5 * time.Second): + t.Fatal("preempted channel was never observed in OnAgentEvents") + } + + loop.Stop() + loop.Wait() +} + +// ============================================================================= +// preemptController unit tests +// ============================================================================= + +func requireAckClosed(t *testing.T, ack <-chan struct{}) { + t.Helper() + waitOrFail(t, ack, "ack should be closed") +} + +func requireAckOpen(t *testing.T, ack <-chan struct{}) { + t.Helper() + select { + case <-ack: + t.Fatal("ack should still be open") + default: + } +} + +func requirePreemptPhase(t *testing.T, c *preemptController, phase preemptTurnPhase) { + t.Helper() + require.Eventually(t, func() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.turnPhase == phase + }, time.Second, time.Millisecond) +} + +func TestPreemptController_BeginPushSnapshotsPlanningTurn(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + + snapshot := c.beginPush() + c.endPush() + + assert.True(t, snapshot.hasTargetTurn) + assert.NotZero(t, snapshot.turnID) + assert.Nil(t, snapshot.ctx) + assert.Nil(t, snapshot.tc) +} + +type testContextKey struct{} + +func TestPreemptController_BeginPushSnapshotsActiveTurn(t *testing.T) { + c := newPreemptController() + ctx := context.WithValue(context.Background(), testContextKey{}, "value") + tc := "turn-context" + + c.beginPlanningTurn() + c.beginActiveTurn(ctx, tc) + snapshot := c.beginPush() + c.endPush() + + assert.True(t, snapshot.hasTargetTurn) + assert.NotZero(t, snapshot.turnID) + assert.Equal(t, ctx, snapshot.ctx) + assert.Equal(t, tc, snapshot.tc) +} + +func TestPreemptController_RequestPreemptIdleTurnAcksImmediately(t *testing.T) { + c := newPreemptController() + snapshot := c.beginPush() + c.endPush() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + requireAckClosed(t, ack) + _, ok := c.receivePreempt() + assert.False(t, ok) +} + +func TestPreemptController_RejectsInvalidTurnPhaseTransitions(t *testing.T) { + c := newPreemptController() + + assert.PanicsWithValue(t, "adk: preemptController.beginActiveTurn called while turn phase is idle; expected planning", func() { + c.beginActiveTurn(context.Background(), "tc") + }) + assert.PanicsWithValue(t, "adk: preemptController.abortPlanningTurn called while turn phase is idle; expected planning", func() { + c.abortPlanningTurn() + }) + assert.PanicsWithValue(t, "adk: preemptController.endActiveTurn called while turn phase is idle; expected active", func() { + c.endActiveTurn() + }) + + c.beginPlanningTurn() + assert.PanicsWithValue(t, "adk: preemptController.beginPlanningTurn called while turn phase is planning; expected idle", func() { + c.beginPlanningTurn() + }) + c.abortPlanningTurn() + + c.pending = newPreemptRequest(nil, nil, time.Now()) + assert.PanicsWithValue(t, "adk: preemptController.beginPlanningTurn called with stale pending preempt request", func() { + c.beginPlanningTurn() + }) +} + +func TestPreemptController_RequestPreemptForPlanningTurnIsConsumedAfterActivation(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + snapshot := c.beginPush() + c.endPush() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + _, ok := c.receivePreempt() + assert.False(t, ok) + requireAckOpen(t, ack) + + c.beginActiveTurn(context.Background(), "tc") + req, ok := c.receivePreempt() + require.True(t, ok) + requireAckOpen(t, ack) + + req.ack() + requireAckClosed(t, ack) +} + +func TestPreemptController_RequestPreemptForActiveTurnIsConsumedOnce(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + req, ok := c.receivePreempt() + require.True(t, ok) + opts := req.cancelOptions(time.Now()) + cfg := parseAgentCancelOptions(opts...) + assert.Equal(t, CancelAfterChatModel, cfg.Mode) + requireAckOpen(t, ack) + + _, ok = c.receivePreempt() + assert.False(t, ok) + + req.ack() + requireAckClosed(t, ack) +} + +func TestPreemptController_AbortPlanningTurnAcksUnconsumedRequest(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + snapshot := c.beginPush() + c.endPush() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + req := c.abortPlanningTurn() + require.NotNil(t, req) + req.ack() + requireAckClosed(t, ack) + + _, ok := c.receivePreempt() + assert.False(t, ok) +} + +func TestPreemptController_EndActiveTurnAcksUnconsumedRequest(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + req := c.endActiveTurn() + require.NotNil(t, req) + req.ack() + requireAckClosed(t, ack) + + _, ok := c.receivePreempt() + assert.False(t, ok) +} + +func TestPreemptController_EndActiveTurnDoesNotAckConsumedRequest(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + req, ok := c.receivePreempt() + require.True(t, ok) + + assert.Nil(t, c.endActiveTurn()) + requireAckOpen(t, ack) + + req.ack() + requireAckClosed(t, ack) +} + +func TestPreemptController_TargetTurnMismatchAcksImmediately(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + snapshot := c.beginPush() + c.endPush() + c.abortPlanningTurn() + c.beginPlanningTurn() + + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + requireAckClosed(t, ack) + _, ok := c.receivePreempt() + assert.False(t, ok) +} + +func TestPreemptController_WaitForPushesBlocksUntilPushEnds(t *testing.T) { + c := newPreemptController() + c.beginPush() + + waitDone := make(chan struct{}) + go func() { + c.waitForPushes() + close(waitDone) + }() + + select { + case <-waitDone: + t.Fatal("waitForPushes should block while a push is in flight") + default: + } + + c.endPush() + waitOrFail(t, waitDone, "waitForPushes should unblock after endPush") +} + +func TestPreemptController_MultipleRequestsBeforeReceiveMergeAcksAndUseMergedOpts(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + ack1 := make(chan struct{}) + ack2 := make(chan struct{}) + c.requestPreempt(snapshot, ack1, WithAgentCancelMode(CancelAfterChatModel), WithAgentCancelTimeout(time.Minute)) + c.requestPreempt(snapshot, ack2, WithAgentCancelMode(CancelAfterToolCalls), WithRecursive(), WithAgentCancelTimeout(time.Second)) + + req, ok := c.receivePreempt() + require.True(t, ok) + opts := req.cancelOptions(time.Now()) + cfg := parseAgentCancelOptions(opts...) + assert.Equal(t, CancelAfterChatModel|CancelAfterToolCalls, cfg.Mode) + assert.True(t, cfg.Recursive) + require.NotNil(t, cfg.Timeout) + assert.LessOrEqual(t, *cfg.Timeout, time.Second) + + requireAckOpen(t, ack1) + requireAckOpen(t, ack2) + req.ack() + requireAckClosed(t, ack1) + requireAckClosed(t, ack2) +} + +func TestPreemptController_ConcurrentPreemptRequestsMergeAndAck(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + const requestCount = 50 + start := make(chan struct{}) + acks := make([]chan struct{}, requestCount) + var wg sync.WaitGroup + for i := 0; i < requestCount; i++ { + acks[i] = make(chan struct{}) + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + c.requestPreempt(snapshot, acks[i], WithAgentCancelMode(CancelAfterChatModel)) + }(i) + } + + close(start) + wg.Wait() + + req, ok := c.receivePreempt() + require.True(t, ok) + req.ack() + for _, ack := range acks { + requireAckClosed(t, ack) + } +} + +func TestPreemptController_CloseForLoopExitDuringDelayedPreempt(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + + // Capture a snapshot while the turn is active. + snapshot := c.beginPush() + c.endPush() + + // closeForLoopExit tears down controller state during TurnLoop cleanup. + c.closeForLoopExit() + + // A delayed goroutine fires requestPreempt AFTER closeForLoopExit. + // This must not panic or deadlock; ack should be closed immediately + // because the controller is now closed. + ack := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + }() + + waitOrFail(t, done, "requestPreempt after closeForLoopExit must not deadlock") + requireAckClosed(t, ack) +} + +func TestPreemptController_RequestPreemptAfterCloseForLoopExit(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + c.closeForLoopExit() + + // requestPreempt on a closed controller should close ack immediately. + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterToolCalls)) + requireAckClosed(t, ack) + + // receivePreempt should return nothing. + _, ok := c.receivePreempt() + assert.False(t, ok) +} + +func TestPreemptController_ConcurrentBeginPushAndWaitForPushes(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + + const pushCount = 100 + var wg sync.WaitGroup + + // Launch many goroutines doing beginPush / endPush concurrently. + for i := 0; i < pushCount; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = c.beginPush() + // Simulate some work. + time.Sleep(time.Microsecond) + c.endPush() + }() + } + + // Meanwhile, waitForPushes should eventually return once all are done. + waitDone := make(chan struct{}) + go func() { + c.waitForPushes() + close(waitDone) + }() + + wg.Wait() // all pushes complete + + waitOrFail(t, waitDone, "waitForPushes deadlocked with concurrent beginPush/endPush") +} + +func TestPreemptController_RequestPreemptWithNilAck(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + // requestPreempt with nil ack channel must not panic. + assert.NotPanics(t, func() { + c.requestPreempt(snapshot, nil, WithAgentCancelMode(CancelAfterChatModel)) + }) + + // The request should still be stored and consumable. + req, ok := c.receivePreempt() + require.True(t, ok) + // ack() with nil channels in the list must not panic. + assert.NotPanics(t, func() { + req.ack() + }) +} + +func TestPreemptController_MergeImmediateOverridesTimeout(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc") + snapshot := c.beginPush() + c.endPush() + + // First request: graceful with timeout. + ack1 := make(chan struct{}) + c.requestPreempt(snapshot, ack1, + WithAgentCancelMode(CancelAfterToolCalls), + WithAgentCancelTimeout(10*time.Second)) + + // Second request: CancelImmediate (no timeout). + ack2 := make(chan struct{}) + c.requestPreempt(snapshot, ack2, WithAgentCancelMode(CancelImmediate)) + + req, ok := c.receivePreempt() + require.True(t, ok) + + opts := req.cancelOptions(time.Now()) + cfg := parseAgentCancelOptions(opts...) + + // CancelImmediate should win and timeout should be nil. + assert.Equal(t, CancelImmediate, cfg.Mode) + assert.Nil(t, cfg.Timeout, "CancelImmediate merge should clear timeout") + + req.ack() + requireAckClosed(t, ack1) + requireAckClosed(t, ack2) +} + +func TestPreemptController_DelayedPreemptTargetGoneBetweenTurns(t *testing.T) { + c := newPreemptController() + + // Turn 1: planning → active → end + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc1") + oldSnapshot := c.beginPush() + c.endPush() + req := c.endActiveTurn() + assert.Nil(t, req) // no pending request + + // Turn 2: start a new turn + c.beginPlanningTurn() + c.beginActiveTurn(context.Background(), "tc2") + + // A delayed preempt from Turn 1 fires with stale snapshot. + // It should resolve as no-op (ack immediately) because turnID doesn't match. + ack := make(chan struct{}) + c.requestPreempt(oldSnapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + requireAckClosed(t, ack) + + // The new turn should have no pending preempt. + _, ok := c.receivePreempt() + assert.False(t, ok, "stale preempt must not affect new turn") + + c.endActiveTurn() +} + +func TestPreemptController_EndPushWithoutBeginPushPanics(t *testing.T) { + c := newPreemptController() + + // endPush without a matching beginPush should panic with the new invariant. + assert.PanicsWithValue(t, + "adk: preemptController.endPush called without matching beginPush", + func() { + c.endPush() + }, + ) +} + +func TestPreemptController_BeginActiveTurnNotifiesExistingPending(t *testing.T) { + c := newPreemptController() + c.beginPlanningTurn() + snapshot := c.beginPush() + c.endPush() + + // Send a preempt during planning phase. + ack := make(chan struct{}) + c.requestPreempt(snapshot, ack, WithAgentCancelMode(CancelAfterChatModel)) + + // During planning, receivePreempt returns nothing. + _, ok := c.receivePreempt() + assert.False(t, ok) + + // beginActiveTurn should notify the watcher via the notify channel. + c.beginActiveTurn(context.Background(), "tc") + + // The notify channel should have a message. + select { + case <-c.notify: + // Expected: watcher notification was sent. + case <-time.After(1 * time.Second): + t.Fatal("beginActiveTurn should notify watcher when there is a pending request") + } + + // Now receivePreempt should return the pending request. + req, ok := c.receivePreempt() + require.True(t, ok) + req.ack() + requireAckClosed(t, ack) +} + +// ============================================================================= +// Integration tests for race-prone preempt scenarios +// ============================================================================= + +func TestTurnLoop_ConcurrentPreemptsDuringTurn(t *testing.T) { + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + return &AgentOutput{}, nil + }, + } + + var genInputCount int32 + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + atomic.AddInt32(&genInputCount, 1) + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items, + }, nil + }, + }) + + loop.Push("seed") + + waitOrFail(t, agentStarted, "agent did not start") + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + ok, ack := loop.Push(fmt.Sprintf("urgent-%d", i), WithPreemptTimeout[string, *schema.Message](AnySafePoint, 10*time.Millisecond)) + if ok && ack != nil { + select { + case <-ack: + case <-time.After(5 * time.Second): + t.Error("ack channel not closed within timeout") + } + } + }(i) + } + + // Stop the loop concurrently. The run loop may be blocked on + // buffer.Receive after processing all preempts; Stop unblocks it + // and triggers closeForLoopExit which closes any orphaned ack channels. + go func() { + time.Sleep(500 * time.Millisecond) + loop.Stop(WithImmediate()) + }() + + wg.Wait() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.True(t, atomic.LoadInt32(&genInputCount) >= 2, "should have had at least the initial turn + one preempted turn") +} + +func TestTurnLoop_PreemptBetweenTurnsAcksImmediately(t *testing.T) { + var cancelCount int32 + var turnCount int32 + firstTurnDone := make(chan struct{}) + secondTurnDone := make(chan struct{}) + firstTurnOnce := sync.Once{} + secondTurnOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "fast", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + return &AgentOutput{}, nil + }, + onCancel: func(cc *cancelContext) { + atomic.AddInt32(&cancelCount, 1) + }, + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + atomic.AddInt32(&turnCount, 1) + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + switch atomic.LoadInt32(&turnCount) { + case 1: + firstTurnOnce.Do(func() { close(firstTurnDone) }) + case 2: + secondTurnOnce.Do(func() { close(secondTurnDone) }) + } + return nil + }, + }) + + loop.Push("first") + waitOrFail(t, firstTurnDone, "first turn did not complete") + requirePreemptPhase(t, loop.preemptCtrl, preemptTurnIdle) + + ok, ack := loop.Push("between-turns", WithPreempt[string, *schema.Message](AnySafePoint)) + require.True(t, ok) + require.NotNil(t, ack) + requireAckClosed(t, ack) + assert.Equal(t, int32(0), atomic.LoadInt32(&cancelCount), "between-turn preempt must not submit cancel") + + waitOrFail(t, secondTurnDone, "between-turn item was not processed") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.Equal(t, int32(2), atomic.LoadInt32(&turnCount)) +} + +func TestTurnLoop_PushStrategy_DuringTurnTransition(t *testing.T) { + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + allowFinish := make(chan struct{}) + strategyEntered := make(chan struct{}) + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + select { + case <-allowFinish: + return &AgentOutput{}, nil + case <-ctx.Done(): + return &AgentOutput{}, nil + } + }, + } + + var genInputCount int32 + secondTurnDone := make(chan struct{}) + secondTurnOnce := sync.Once{} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + count := atomic.AddInt32(&genInputCount, 1) + if count >= 2 { + secondTurnOnce.Do(func() { + close(secondTurnDone) + }) + } + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items, + }, nil + }, + }) + + loop.Push("first") + + waitOrFail(t, agentStarted, "agent did not start") + + strategyBlocker := make(chan struct{}) + var strategyTCNotNil int32 + + go func() { + loop.Push("strategic-item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + if tc != nil { + atomic.StoreInt32(&strategyTCNotNil, 1) + } + close(strategyEntered) + <-strategyBlocker + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} + })) + }() + + waitOrFail(t, strategyEntered, "strategy did not enter") + + close(allowFinish) + + select { + case <-secondTurnDone: + t.Fatal("second turn should not be planned before strategy Push finishes") + default: + } + + close(strategyBlocker) + + waitOrFail(t, secondTurnDone, "second turn should eventually run after strategy resolves") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.True(t, atomic.LoadInt32(&genInputCount) >= 2) + assert.Equal(t, int32(1), atomic.LoadInt32(&strategyTCNotNil)) +} + +func TestTurnLoop_ConcurrentPreemptAndStop(t *testing.T) { + for iter := 0; iter < 20; iter++ { + t.Run(fmt.Sprintf("iter_%d", iter), func(t *testing.T) { + ctx := context.Background() + + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + return &AgentOutput{}, nil + }, + } + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return agent, nil + }, + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items, + }, nil + }, + }) + + loop.Push("seed") + + select { + case <-agentStarted: + case <-time.After(1 * time.Second): + t.Fatal("agent did not start") + } + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + _, ack := loop.Push("preempt-item", WithPreempt[string, *schema.Message](AnySafePoint)) + if ack != nil { + <-ack + } + }() + + go func() { + defer wg.Done() + loop.Stop(WithImmediate()) + }() + + wg.Wait() + loop.Wait() + }) + } +} + +func TestTurnLoop_ConcurrentPushStrategyAndStop(t *testing.T) { + for iter := 0; iter < 20; iter++ { + t.Run(fmt.Sprintf("iter_%d", iter), func(t *testing.T) { + ctx := context.Background() + + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + return &AgentOutput{}, nil + }, + } + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return agent, nil + }, + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{}, + Consumed: items, + }, nil + }, + }) + + loop.Push("seed") + + select { + case <-agentStarted: + case <-time.After(1 * time.Second): + t.Fatal("agent did not start") + } + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + _, ack := loop.Push("strategic-item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} + })) + if ack != nil { + <-ack + } + }() + + go func() { + defer wg.Done() + loop.Stop(WithImmediate()) + }() + + wg.Wait() + loop.Wait() + }) + } +} + +func TestTurnLoop_TurnContext_StoppedChannel(t *testing.T) { + stoppedSeen := make(chan struct{}) + agentStarted := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + close(agentStarted) + select { + case <-tc.Stopped: + close(stoppedSeen) + case <-time.After(5 * time.Second): + t.Error("timed out waiting for Stopped channel") + } + // Drain events + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Stop(WithImmediate()) + + select { + case <-stoppedSeen: + // success + case <-time.After(5 * time.Second): + t.Fatal("stopped channel was never observed in OnAgentEvents") + } + + loop.Wait() +} + +func TestTurnLoop_TurnContext_BothPreemptedAndStopped(t *testing.T) { + t.Run("PreemptThenStop_OnlyPreemptContributes", func(t *testing.T) { + preemptedSeen := make(chan struct{}) + agentStarted := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*TypedAgentEvent[*schema.Message]]) error { + close(agentStarted) + select { + case <-tc.Preempted: + close(preemptedSeen) + case <-time.After(5 * time.Second): + t.Error("timed out waiting for Preempted") + } + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Push("msg2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) + + select { + case <-preemptedSeen: + case <-time.After(5 * time.Second): + t.Fatal("Preempted channel was never closed") + } + + loop.Stop(WithImmediate()) + loop.Wait() + }) + + t.Run("StopThenPreempt_OnlyStopContributes", func(t *testing.T) { + stoppedSeen := make(chan struct{}) + agentStarted := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*TypedAgentEvent[*schema.Message]]) error { + close(agentStarted) + select { + case <-tc.Stopped: + close(stoppedSeen) + case <-time.After(5 * time.Second): + t.Error("timed out waiting for Stopped") + } + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Stop(WithImmediate()) + + select { + case <-stoppedSeen: + case <-time.After(5 * time.Second): + t.Fatal("Stopped channel was never closed") + } + + loop.Push("msg2", WithPreemptTimeout[string, *schema.Message](AnySafePoint, time.Millisecond)) + loop.Wait() + }) +} + +func TestTurnLoop_PushStrategy_DuringTurn(t *testing.T) { + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + agentCancelled := make(chan struct{}) + agentCancelledOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + agentCancelledOnce.Do(func() { + close(agentCancelled) + }) + return &AgentOutput{}, nil + }, + } + + genInputCalls := int32(0) + secondGenInputCalled := make(chan struct{}) + secondGenInputOnce := sync.Once{} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + count := atomic.AddInt32(&genInputCalls, 1) + if count >= 2 { + secondGenInputOnce.Do(func() { + close(secondGenInputCalled) + }) + } + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: []string{items[0]}, + Remaining: items[1:], + }, nil + }, + }) + + loop.Push("first") + + waitOrFail(t, agentStarted, "agent did not start") + + // Strategy inspects TurnContext during a running turn and decides to preempt. + var strategyCalled int32 + var strategyTC *TurnContext[string, *schema.Message] + loop.Push("urgent", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + atomic.AddInt32(&strategyCalled, 1) + strategyTC = tc + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} + })) + + waitOrFail(t, agentCancelled, "agent was not cancelled by strategy-returned preempt") + + waitOrFail(t, secondGenInputCalled, "second GenInput was not called after preempt") + + loop.Stop(WithImmediate()) + loop.Wait() + + assert.Equal(t, int32(1), atomic.LoadInt32(&strategyCalled)) + assert.NotNil(t, strategyTC, "strategy should receive non-nil TurnContext during a turn") + assert.Equal(t, []string{"first"}, strategyTC.Consumed) +} + +func TestTurnLoop_PushStrategy_BetweenTurns(t *testing.T) { + // Push with strategy before Run() — TurnContext should be nil. + var strategyCalled int32 + var strategyTCWasNil bool + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + return &AgentOutput{}, nil + }, + } + + agentDone := make(chan struct{}) + agentDoneOnce := sync.Once{} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + Remaining: nil, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + agentDoneOnce.Do(func() { + close(agentDone) + }) + return nil + }, + }) + + // Push with strategy — no turn is active yet, so tc should be nil. + loop.Push("item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + atomic.AddInt32(&strategyCalled, 1) + strategyTCWasNil = tc == nil + return nil // plain push, no preempt + })) + + waitOrFail(t, agentDone, "agent did not complete") + + loop.Stop() + loop.Wait() + + assert.Equal(t, int32(1), atomic.LoadInt32(&strategyCalled)) + assert.True(t, strategyTCWasNil, "strategy should receive nil TurnContext between turns") +} + +func TestTurnLoop_PushStrategy_OverridesOtherOptions(t *testing.T) { + // Push with both WithPreempt and WithPushStrategy — only strategy's result applies. + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + return &AgentOutput{}, nil + }, + } + + agentDone := make(chan struct{}) + agentDoneOnce := sync.Once{} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + Remaining: nil, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + agentDoneOnce.Do(func() { + close(agentDone) + }) + return nil + }, + }) + + // Strategy returns nil (no preempt), even though WithPreempt is also passed. + // The strategy should override — so the agent should NOT be preempted. + ok, ack := loop.Push("item", WithPreempt[string, *schema.Message](AnySafePoint), WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + return nil // no preempt + })) + assert.True(t, ok) + assert.Nil(t, ack, "ack should be nil since strategy returned no preempt") + + waitOrFail(t, agentDone, "agent did not complete normally") + + loop.Stop() + loop.Wait() +} + +func TestTurnLoop_PushStrategy_NestedStrategyStripped(t *testing.T) { + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + return &AgentOutput{}, nil + }, + } + + agentDone := make(chan struct{}) + agentDoneOnce := sync.Once{} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + Remaining: nil, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + _, ok := events.Next() + if !ok { + break + } + } + agentDoneOnce.Do(func() { + close(agentDone) + }) + return nil + }, + }) + + // Strategy returns another WithPushStrategy — the nested one should be stripped. + innerCalled := int32(0) + ok, ack := loop.Push("item", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + return []PushOption[string, *schema.Message]{ + WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + atomic.AddInt32(&innerCalled, 1) + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} + }), + } + })) + assert.True(t, ok) + assert.Nil(t, ack, "ack should be nil since nested strategy was stripped (no preempt)") + + waitOrFail(t, agentDone, "agent did not complete normally") + + loop.Stop() + loop.Wait() + + assert.Equal(t, int32(0), atomic.LoadInt32(&innerCalled), "nested strategy should not be called") +} + +func TestTurnLoop_PushStrategy_ConsumedInspection(t *testing.T) { + // Strategy preempts only when current turn is processing "low-priority" items. + agentStarted := make(chan struct{}) + agentStartedOnce := sync.Once{} + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + agentStartedOnce.Do(func() { + close(agentStarted) + }) + <-ctx.Done() + return &AgentOutput{}, nil + }, + } + + genInputCalls := int32(0) + secondGenInputItems := make(chan []string, 1) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + count := atomic.AddInt32(&genInputCalls, 1) + if count >= 2 { + select { + case secondGenInputItems <- append([]string{}, items...): + default: + } + } + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: []string{items[0]}, + Remaining: items[1:], + }, nil + }, + }) + + loop.Push("low-priority-task") + + waitOrFail(t, agentStarted, "agent did not start") + + // Strategy checks Consumed and preempts because current turn has "low-priority" items. + loop.Push("urgent-task", WithPushStrategy(func(ctx context.Context, tc *TurnContext[string, *schema.Message]) []PushOption[string, *schema.Message] { + if tc != nil && len(tc.Consumed) > 0 && tc.Consumed[0] == "low-priority-task" { + return []PushOption[string, *schema.Message]{WithPreempt[string, *schema.Message](AnySafePoint)} + } + return nil + })) + + select { + case items := <-secondGenInputItems: + assert.Contains(t, items, "urgent-task") + case <-time.After(2 * time.Second): + t.Fatal("second GenInput was not called after strategy-driven preempt") + } + + loop.Stop(WithImmediate()) + loop.Wait() +} + +func TestTurnLoop_PushAfterStop_BufferedAsLateItems(t *testing.T) { + ctx := context.Background() + processed := make(chan string, 10) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + processed <- tc.Consumed[0] + return nil + }, + }) + + loop.Push("msg1") + <-processed + loop.Stop() + result := loop.Wait() + + // Push after stop — should be buffered as late items + ok1, _ := loop.Push("late1") + ok2, _ := loop.Push("late2") + ok3, _ := loop.Push("late3") + assert.False(t, ok1) + assert.False(t, ok2) + assert.False(t, ok3) + + late := result.TakeLateItems() + assert.Equal(t, []string{"late1", "late2", "late3"}, late) +} + +func TestTurnLoop_TakeLateItems_Idempotent(t *testing.T) { + ctx := context.Background() + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Stop() + loop.Run(ctx) + result := loop.Wait() + + loop.Push("late1") + + first := result.TakeLateItems() + second := result.TakeLateItems() + third := result.TakeLateItems() + + assert.Equal(t, []string{"late1"}, first) + assert.Equal(t, first, second, "subsequent calls should return the same slice") + assert.Equal(t, first, third, "subsequent calls should return the same slice") +} + +func TestTurnLoop_PushAfterTakeLateItems_Panics(t *testing.T) { + ctx := context.Background() + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Stop() + loop.Run(ctx) + result := loop.Wait() + + result.TakeLateItems() + + assert.PanicsWithValue(t, "TurnLoop: Push called after TakeLateItems", func() { + loop.Push("too-late") + }) +} + +func TestTurnLoop_TakeLateItems_NeverCalled_NoImpact(t *testing.T) { + ctx := context.Background() + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Push("b") + loop.Stop() + loop.Run(ctx) + result := loop.Wait() + + // Don't call TakeLateItems — verify UnhandledItems works normally + assert.Contains(t, result.UnhandledItems, "b") + assert.Nil(t, result.ExitReason) +} + +func TestTurnLoop_CheckpointErr_SeparateFromExitReason(t *testing.T) { + ctx := context.Background() + saveStore := &errorCheckpointStore{setErr: fmt.Errorf("storage unavailable")} + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: saveStore, + CheckpointID: "cp-separate-err", + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Stop() + loop.Run(ctx) + result := loop.Wait() + + // ExitReason should be nil (clean stop), checkpoint error should be separate + assert.Nil(t, result.ExitReason) + assert.True(t, result.CheckpointAttempted) + assert.Error(t, result.CheckpointErr) + assert.Contains(t, result.CheckpointErr.Error(), "storage unavailable") +} + +func TestTurnLoop_CheckpointAttempted_FalseWhenNoStore(t *testing.T) { + ctx := context.Background() + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop.Push("a") + loop.Stop() + loop.Run(ctx) + result := loop.Wait() + + assert.False(t, result.CheckpointAttempted) + assert.Nil(t, result.CheckpointErr) +} + +func TestTurnLoop_CheckpointAttempted_FalseOnErrorExit(t *testing.T) { + ctx := context.Background() + store := newTestStore() + genInputErr := errors.New("gen input failed") + + firstTurnDone := make(chan struct{}) + var callCount int32 + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: "cp-err-exit", + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + n := atomic.AddInt32(&callCount, 1) + if n > 1 { + return nil, genInputErr + } + return &GenInputResult[string, *schema.Message]{Input: &AgentInput{}, Consumed: items}, nil + }, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + close(firstTurnDone) + return nil + }, + }) + loop.Push("msg1") + <-firstTurnDone + loop.Push("msg2") + result := loop.Wait() + + // Loop exited from error, not Stop() — checkpoint should not be saved + assert.ErrorIs(t, result.ExitReason, genInputErr) + assert.False(t, result.CheckpointAttempted) + assert.Nil(t, result.CheckpointErr) +} + +func TestTurnLoop_StopConcurrentWithCallbackError_NoCheckpoint(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "stop-concurrent-err" + + prepareErr := errors.New("prepare agent failed") + firstTurnDone := make(chan struct{}) + stopCalled := make(chan struct{}) + var prepareCount int32 + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + n := atomic.AddInt32(&prepareCount, 1) + if n > 1 { + // Wait until Stop() has been called so stopCtrl.isCommitted() is true. + <-stopCalled + return nil, prepareErr + } + return &turnLoopMockAgent{name: "test"}, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + close(firstTurnDone) + return nil + }, + }) + + loop.Push("msg1") + <-firstTurnDone + loop.Push("msg2") + + // Call Stop() and signal PrepareAgent to proceed with error + go func() { + loop.Stop() + close(stopCalled) + }() + + result := loop.Wait() + + // The loop may exit via Stop (clean) or via PrepareAgent error. + // If it exited via PrepareAgent error with Stop also called: + // checkpoint should NOT be saved. + if result.ExitReason != nil && !errors.As(result.ExitReason, new(*CancelError)) { + assert.ErrorIs(t, result.ExitReason, prepareErr) + assert.False(t, result.CheckpointAttempted, "should not checkpoint when exit is caused by callback error") + } + // If Stop won the race, that's fine — checkpoint may or may not be saved + // depending on idle state. The test is about the error path. +} + +func TestTurnLoop_DeleteWithoutCheckPointDeleter_NoOp(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "no-deleter" + + // First loop: save a checkpoint + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop1.Push("a") + loop1.Stop() + loop1.Run(ctx) + loop1.Wait() + + store.mu.Lock() + _, exists := store.m[cpID] + store.mu.Unlock() + assert.True(t, exists, "checkpoint should be saved") + + // Second loop: exit via context cancel — should try to delete but store + // doesn't implement CheckPointDeleter, so checkpoint persists (no-op) + ctx2, cancel2 := context.WithCancel(ctx) + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + cancel2() + return nil + }, + }) + loop2.Push("b") + loop2.Run(ctx2) + loop2.Wait() + + // Without CheckPointDeleter, the stale checkpoint should NOT be deleted + store.mu.Lock() + v, exists := store.m[cpID] + store.mu.Unlock() + assert.True(t, exists, "checkpoint should still exist without CheckPointDeleter") + assert.NotNil(t, v, "checkpoint should not be set to nil") +} + +func TestTurnLoop_StopWithSkipCheckpoint(t *testing.T) { + ctx := context.Background() + store := newTestStore() + cpID := "skip-cp-session" + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("a") + loop.Push("b") + loop.Stop(WithSkipCheckpoint()) + loop.Run(ctx) + + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + assert.False(t, exit.CheckpointAttempted, "checkpoint should be skipped when WithSkipCheckpoint is used") + + store.mu.Lock() + _, exists := store.m[cpID] + store.mu.Unlock() + assert.False(t, exists, "no checkpoint should be saved when WithSkipCheckpoint is used") +} + +func TestTurnLoop_StopWithSkipCheckpoint_DeletesStaleCheckpoint(t *testing.T) { + ctx := context.Background() + store := &deletableCheckpointStore{ + turnLoopCheckpointStore: turnLoopCheckpointStore{m: make(map[string][]byte)}, + } + cpID := "skip-stale-session" + + loop1 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop1.Push("a") + loop1.Stop() + loop1.Run(ctx) + exit1 := loop1.Wait() + assert.True(t, exit1.CheckpointAttempted) + + store.mu.Lock() + _, exists := store.m[cpID] + store.mu.Unlock() + assert.True(t, exists, "first loop should save checkpoint") + + loop2 := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + loop2.Push("b") + loop2.Stop(WithSkipCheckpoint()) + loop2.Run(ctx) + exit2 := loop2.Wait() + assert.False(t, exit2.CheckpointAttempted, "second loop should skip checkpoint") + + store.mu.Lock() + deleteCalled := store.deleteCalled + store.mu.Unlock() + assert.True(t, deleteCalled, "stale checkpoint should be deleted when SkipCheckpoint is used") +} + +func TestTurnLoop_StopWithStopCause(t *testing.T) { + ctx := context.Background() + cause := "user session timeout" + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Push("a") + loop.Stop(WithStopCause(cause)) + + exit := loop.Wait() + assert.Equal(t, cause, exit.StopCause) +} + +func TestTurnLoop_StopCause_EmptyWhenNoStop(t *testing.T) { + ctx := context.Background() + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAll, + PrepareAgent: prepareTestAgent, + }) + + loop.Stop() + exit := loop.Wait() + assert.Empty(t, exit.StopCause) +} + +func TestTurnLoop_StopCause_InTurnContext(t *testing.T) { + cause := "business shutdown" + gotCause := make(chan string, 1) + agentStarted := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + close(agentStarted) + select { + case <-tc.Stopped: + gotCause <- tc.StopCause() + case <-time.After(5 * time.Second): + t.Error("timed out waiting for Stopped channel") + } + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Stop(WithImmediate(), WithStopCause(cause)) + + select { + case c := <-gotCause: + assert.Equal(t, cause, c) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for StopCause in TurnContext") + } + + exit := loop.Wait() + assert.Equal(t, cause, exit.StopCause) +} + +func TestTurnLoop_StopCause_FirstNonEmptyWins(t *testing.T) { + agentStarted := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + close(agentStarted) + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Stop(WithGraceful(), WithStopCause("first cause")) + loop.Stop(WithStopCause("second cause")) + + exit := loop.Wait() + assert.Equal(t, "first cause", exit.StopCause, "first non-empty StopCause should win") +} + +func TestTurnLoop_StopBeforeRun_PushThenStop(t *testing.T) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + t.Fatal("GenInput should not be called when Stop is called before Run") + return nil, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + t.Fatal("PrepareAgent should not be called when Stop is called before Run") + return nil, nil + }, + }) + + ok, _ := loop.Push("item1") + assert.True(t, ok) + ok, _ = loop.Push("item2") + assert.True(t, ok) + + loop.Stop() + loop.Run(context.Background()) + result := loop.Wait() + + assert.NoError(t, result.ExitReason) + assert.Equal(t, []string{"item1", "item2"}, result.UnhandledItems) + assert.Empty(t, result.InterruptedItems) + assert.Empty(t, result.TakeLateItems()) +} + +func TestTurnLoop_StopBeforeRun_StopThenPush(t *testing.T) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + t.Fatal("GenInput should not be called when Stop is called before Run") + return nil, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + t.Fatal("PrepareAgent should not be called when Stop is called before Run") + return nil, nil + }, + }) + + loop.Stop() + + ok, _ := loop.Push("item1") + assert.False(t, ok) + ok, _ = loop.Push("item2") + assert.False(t, ok) + + loop.Run(context.Background()) + result := loop.Wait() + + assert.NoError(t, result.ExitReason) + assert.Empty(t, result.UnhandledItems) + assert.Empty(t, result.InterruptedItems) + assert.Equal(t, []string{"item1", "item2"}, result.TakeLateItems()) +} + +func TestTurnLoop_SkipCheckpoint_Sticky(t *testing.T) { + agentStarted := make(chan struct{}) + + store := newTestStore() + cpID := "sticky-skip-session" + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: cpID, + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "slow", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + close(agentStarted) + for { + if _, ok := events.Next(); !ok { + break + } + } + return nil + }, + }) + + loop.Push("msg1") + <-agentStarted + loop.Stop(WithGraceful(), WithSkipCheckpoint()) + loop.Stop() + + exit := loop.Wait() + assert.False(t, exit.CheckpointAttempted, "SkipCheckpoint should be sticky across multiple Stop calls") + + store.mu.Lock() + _, exists := store.m[cpID] + store.mu.Unlock() + assert.False(t, exists, "no checkpoint should be saved when SkipCheckpoint was set in any Stop call") +} + +func TestWithGracefulTimeout_NonPositive_Panics(t *testing.T) { + assert.PanicsWithValue(t, "adk: WithGracefulTimeout: gracePeriod must be positive", + func() { WithGracefulTimeout(0) }) + assert.PanicsWithValue(t, "adk: WithGracefulTimeout: gracePeriod must be positive", + func() { WithGracefulTimeout(-1 * time.Second) }) +} + +func TestWithPreempt_ZeroSafePoint_Panics(t *testing.T) { + assert.PanicsWithValue(t, "adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint", + func() { WithPreempt[string, *schema.Message](SafePoint(0)) }) +} + +func TestWithPreemptTimeout_ZeroSafePoint_Panics(t *testing.T) { + assert.PanicsWithValue(t, "adk: SafePoint must not be zero; use AfterToolCalls, AfterChatModel, or AnySafePoint", + func() { WithPreemptTimeout[string, *schema.Message](SafePoint(0), time.Second) }) +} + +func TestSafePoint_ToCancelMode(t *testing.T) { + assert.Equal(t, CancelAfterToolCalls, AfterToolCalls.toCancelMode()) + assert.Equal(t, CancelAfterChatModel, AfterChatModel.toCancelMode()) + assert.Equal(t, CancelAfterToolCalls|CancelAfterChatModel, AnySafePoint.toCancelMode()) +} + +func TestNewTurnLoop_NilGenInput_Panics(t *testing.T) { + assert.PanicsWithValue(t, "adk: NewTurnLoop: GenInput is required", func() { + NewTurnLoop(TurnLoopConfig[string, *schema.Message]{PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return nil, nil + }}) + }) +} + +func TestNewTurnLoop_NilPrepareAgent_Panics(t *testing.T) { + assert.PanicsWithValue(t, "adk: NewTurnLoop: PrepareAgent is required", func() { + NewTurnLoop(TurnLoopConfig[string, *schema.Message]{GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return nil, nil + }}) + }) +} + +func TestDeriveAgentToolCancelContext_NilParent_ReturnsNil(t *testing.T) { + var cc *cancelContext + assert.Nil(t, cc.deriveAgentToolCancelContext(context.Background())) +} + +func TestUntilIdleFor(t *testing.T) { + t.Run("FiresAfterIdleDuration", func(t *testing.T) { + turnDone := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(turnDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-turnDone + + loop.Stop(UntilIdleFor(50 * time.Millisecond)) + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("loop did not exit after idle timeout") + } + }) + + t.Run("ResetsOnPush", func(t *testing.T) { + turnCount := int32(0) + turnDone := make(chan struct{}, 10) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + atomic.AddInt32(&turnCount, 1) + turnDone <- struct{}{} + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-turnDone + + loop.Stop(UntilIdleFor(200 * time.Millisecond)) + + time.Sleep(100 * time.Millisecond) + loop.Push("msg2") + <-turnDone + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("loop did not exit after idle timeout") + } + + assert.Equal(t, int32(2), atomic.LoadInt32(&turnCount)) + }) + + t.Run("EscalatedByStopWithImmediate", func(t *testing.T) { + agentStarted := make(chan *cancelContext, 1) + probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return probe, nil + }, + }) + + loop.Push("msg1") + cc := <-agentStarted + + loop.Stop(UntilIdleFor(10 * time.Minute)) + loop.Stop(WithImmediate()) + + deadline := time.After(2 * time.Second) + for { + if cc.getMode() == CancelImmediate { + break + } + select { + case <-deadline: + t.Fatal("cancel mode did not escalate to CancelImmediate") + default: + } + time.Sleep(1 * time.Millisecond) + } + + exit := loop.Wait() + var ce *CancelError + require.True(t, errors.As(exit.ExitReason, &ce)) + assert.Equal(t, CancelImmediate, ce.Info.Mode) + }) + + t.Run("EscalatedByStopWithGraceful", func(t *testing.T) { + agentStarted := make(chan struct{}) + agentDone := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + <-ctx.Done() + close(agentDone) + return nil, ctx.Err() + }, + }, nil + }, + }) + + loop.Push("msg1") + <-agentStarted + + loop.Stop(UntilIdleFor(10 * time.Minute)) + loop.Stop(WithGracefulTimeout(50 * time.Millisecond)) + + select { + case <-agentDone: + case <-time.After(2 * time.Second): + t.Fatal("agent was not cancelled") + } + + exit := loop.Wait() + assert.Error(t, exit.ExitReason) + }) +} + +// TestUntilIdleFor_DoesNotCancelRunningAgent verifies that Stop(UntilIdleFor) +// records an idle stop policy but does NOT create a pending cancel request for +// the running agent. +func TestUntilIdleFor_DoesNotCancelRunningAgent(t *testing.T) { + t.Run("BeforeRun", func(t *testing.T) { + agentStarted := make(chan struct{}) + agentCtxCanceled := int32(0) + agentDone := make(chan struct{}) + + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + // Block until context is canceled or a short timeout. + select { + case <-ctx.Done(): + atomic.StoreInt32(&agentCtxCanceled, 1) + case <-time.After(200 * time.Millisecond): + } + close(agentDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + // Call Stop(UntilIdleFor) BEFORE Run. + loop.Stop(UntilIdleFor(50 * time.Millisecond)) + loop.Run(context.Background()) + + <-agentStarted + <-agentDone + + exit := loop.Wait() + assert.Nil(t, exit.ExitReason, "UntilIdleFor should not produce a CancelError") + assert.Equal(t, int32(0), atomic.LoadInt32(&agentCtxCanceled), + "agent context should not have been canceled by UntilIdleFor") + }) + + t.Run("DuringRun", func(t *testing.T) { + agentStarted := make(chan struct{}) + agentCtxCanceled := int32(0) + agentDone := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + select { + case <-ctx.Done(): + atomic.StoreInt32(&agentCtxCanceled, 1) + case <-time.After(200 * time.Millisecond): + } + close(agentDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-agentStarted + + // Call Stop(UntilIdleFor) while the agent is running. + loop.Stop(UntilIdleFor(50 * time.Millisecond)) + <-agentDone + + exit := loop.Wait() + assert.Nil(t, exit.ExitReason, "UntilIdleFor should not produce a CancelError") + assert.Equal(t, int32(0), atomic.LoadInt32(&agentCtxCanceled), + "agent context should not have been canceled by UntilIdleFor") + }) + + // Cancel opts paired with UntilIdleFor in the same call are silently + // dropped. The agent must run to completion even when WithImmediate is + // combined with UntilIdleFor. + t.Run("CancelOptsDroppedInSameCall", func(t *testing.T) { + agentStarted := make(chan struct{}) + agentCtxCanceled := int32(0) + agentDone := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage(items[0])}}, + Consumed: items, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + select { + case <-ctx.Done(): + atomic.StoreInt32(&agentCtxCanceled, 1) + case <-time.After(200 * time.Millisecond): + } + close(agentDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-agentStarted + + // WithImmediate in the same call as UntilIdleFor must be ignored. + loop.Stop(UntilIdleFor(50*time.Millisecond), WithImmediate()) + <-agentDone + + exit := loop.Wait() + assert.Nil(t, exit.ExitReason, "cancel opts should be dropped when combined with UntilIdleFor") + assert.Equal(t, int32(0), atomic.LoadInt32(&agentCtxCanceled), + "agent context should not have been canceled") + }) +} + +func TestUntilIdleFor_ContextCancelDuringIdleWait(t *testing.T) { + turnDone := make(chan struct{}) + ctx, cancel := context.WithCancel(context.Background()) + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(turnDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-turnDone + + // Start idle timer, then cancel the parent context while idle. + loop.Stop(UntilIdleFor(10 * time.Minute)) + time.Sleep(20 * time.Millisecond) + cancel() + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + waitOrFail(t, done, "loop should exit when context is canceled during idle wait") + + exit := loop.Wait() + assert.ErrorIs(t, exit.ExitReason, context.Canceled) +} + +func TestCancelRequestState_ImmediateDominatesSafePointModes(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(time.Minute), + }, now) + + state.merge([]AgentCancelOption{WithAgentCancelMode(CancelImmediate)}, now) + + cfg := parseAgentCancelOptions(state.cancelOptions(now)...) + assert.Equal(t, CancelImmediate, cfg.Mode) + assert.Nil(t, cfg.Timeout) +} + +func TestCancelRequestState_NilMergeDoesNotCreateCancelIntent(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{WithAgentCancelMode(CancelAfterChatModel)}, now) + + state.merge(nil, now) + + cfg := parseAgentCancelOptions(state.cancelOptions(now)...) + assert.Equal(t, CancelAfterChatModel, cfg.Mode) +} + +func TestCancelRequestState_EmptyMergeMeansExplicitImmediate(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{WithAgentCancelMode(CancelAfterChatModel)}, now) + + state.merge([]AgentCancelOption{}, now) + + cfg := parseAgentCancelOptions(state.cancelOptions(now)...) + assert.Equal(t, CancelImmediate, cfg.Mode) +} + +func TestCancelRequestState_SafePointModesJoin(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{WithAgentCancelMode(CancelAfterChatModel)}, now) + + state.merge([]AgentCancelOption{WithAgentCancelMode(CancelAfterToolCalls)}, now) + + cfg := parseAgentCancelOptions(state.cancelOptions(now)...) + assert.Equal(t, CancelAfterChatModel|CancelAfterToolCalls, cfg.Mode) +} + +func TestCancelRequestState_RecursiveIsMonotonic(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{WithAgentCancelMode(CancelAfterChatModel)}, now) + + state.merge([]AgentCancelOption{WithAgentCancelMode(CancelAfterToolCalls), WithRecursive()}, now) + state.merge([]AgentCancelOption{WithAgentCancelMode(CancelAfterChatModel)}, now) + + cfg := parseAgentCancelOptions(state.cancelOptions(now)...) + assert.True(t, cfg.Recursive) +} + +func TestCancelRequestState_TimeoutUsesEarliestDeadline(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(10 * time.Second), + }, now) + + state.merge([]AgentCancelOption{ + WithAgentCancelMode(CancelAfterToolCalls), + WithAgentCancelTimeout(time.Second), + }, now.Add(100*time.Millisecond)) + + cfg := parseAgentCancelOptions(state.cancelOptions(now.Add(100 * time.Millisecond))...) + require.NotNil(t, cfg.Timeout) + assert.LessOrEqual(t, *cfg.Timeout, time.Second) +} + +func TestCancelRequestState_ExpiredTimeoutConvertsToImmediate(t *testing.T) { + now := time.Now() + state := newCancelRequestState([]AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(time.Nanosecond), + }, now) + + cfg := parseAgentCancelOptions(state.cancelOptions(now.Add(time.Second))...) + assert.Equal(t, CancelImmediate, cfg.Mode) + assert.Nil(t, cfg.Timeout) +} + +func TestStopController_BareStopCommitsWithoutCancelRequest(t *testing.T) { + c := newStopController() + + decision := c.requestStop(&stopConfig{}) + + assert.True(t, decision.commit) + assert.True(t, c.isCommitted()) + c.beginActiveTurn() + _, ok := c.receiveCancel() + assert.False(t, ok) +} + +func TestStopController_UntilIdleForDoesNotCreateCancelRequest(t *testing.T) { + c := newStopController() + + decision := c.requestStop(&stopConfig{idleFor: time.Second}) + + assert.False(t, decision.commit) + assert.True(t, decision.wakeIdle) + assert.Equal(t, time.Second, c.idleDuration()) + assert.False(t, c.isCommitted()) + c.beginActiveTurn() + _, ok := c.receiveCancel() + assert.False(t, ok) +} + +func TestStopController_CancelOptsDroppedWhenCombinedWithUntilIdleFor(t *testing.T) { + c := newStopController() + + decision := c.requestStop(&stopConfig{ + idleFor: time.Second, + agentCancelOpts: []AgentCancelOption{WithRecursive()}, + }) + + assert.False(t, decision.commit) + c.beginActiveTurn() + _, ok := c.receiveCancel() + assert.False(t, ok) +} + +func TestStopController_ImmediateStopCreatesPendingCancelForActiveTurn(t *testing.T) { + c := newStopController() + c.beginActiveTurn() + + decision := c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithRecursive()}}) + + assert.True(t, decision.commit) + req, ok := c.receiveCancel() + require.True(t, ok) + cfg := parseAgentCancelOptions(req.cancelOptions(time.Now())...) + assert.Equal(t, CancelImmediate, cfg.Mode) + assert.True(t, cfg.Recursive) +} + +func TestStopController_StopBeforeWatcherStartsConsumedAfterBeginActiveTurn(t *testing.T) { + c := newStopController() + + decision := c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithRecursive()}}) + assert.True(t, decision.commit) + + c.beginActiveTurn() + req, ok := c.receiveCancel() + require.True(t, ok) + cfg := parseAgentCancelOptions(req.cancelOptions(time.Now())...) + assert.Equal(t, CancelImmediate, cfg.Mode) +} + +func TestStopController_EndActiveTurnDropsUnconsumedCancel(t *testing.T) { + c := newStopController() + c.beginActiveTurn() + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithRecursive()}}) + + req := c.endActiveTurn() + + require.NotNil(t, req) + _, ok := c.receiveCancel() + assert.False(t, ok) +} + +func TestStopController_RepeatedStopsMergeWithoutDeescalation(t *testing.T) { + c := newStopController() + c.beginActiveTurn() + + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithRecursive()}}) + c.requestStop(&stopConfig{}) + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel | CancelAfterToolCalls), + WithRecursive(), + }}) + + req, ok := c.receiveCancel() + require.True(t, ok) + cfg := parseAgentCancelOptions(req.cancelOptions(time.Now())...) + assert.Equal(t, CancelImmediate, cfg.Mode) + assert.True(t, cfg.Recursive) +} + +func TestStopController_RepeatedStopsUseSharedCancelMergeState(t *testing.T) { + c := newStopController() + c.beginActiveTurn() + + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithAgentCancelMode(CancelAfterChatModel)}}) + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithAgentCancelMode(CancelAfterToolCalls)}}) + + req, ok := c.receiveCancel() + require.True(t, ok) + cfg := parseAgentCancelOptions(req.cancelOptions(time.Now())...) + assert.Equal(t, CancelAfterChatModel|CancelAfterToolCalls, cfg.Mode) +} + +func TestStopController_StopCauseFirstNonEmptyWins(t *testing.T) { + c := newStopController() + + c.requestStop(&stopConfig{}) + c.requestStop(&stopConfig{stopCause: "first"}) + c.requestStop(&stopConfig{stopCause: "second"}) + + assert.Equal(t, "first", c.cause()) +} + +func TestStopController_SkipCheckpointSticky(t *testing.T) { + c := newStopController() + + c.requestStop(&stopConfig{skipCheckpoint: true}) + c.requestStop(&stopConfig{}) + + assert.True(t, c.skipCheckpointEnabled()) +} + +func TestStopController_ConcurrentStopRequestsRaceSafe(t *testing.T) { + c := newStopController() + c.beginActiveTurn() + + var wg sync.WaitGroup + for i := 0; i < 20; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + switch i % 5 { + case 0: + c.requestStop(&stopConfig{}) + case 1: + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithRecursive()}}) + case 2: + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{ + WithAgentCancelMode(CancelAfterChatModel), + WithAgentCancelTimeout(time.Second), + WithRecursive(), + }}) + case 3: + c.requestStop(&stopConfig{idleFor: time.Second}) + case 4: + c.requestStop(&stopConfig{skipCheckpoint: true, stopCause: "cause"}) + } + }(i) + } + wg.Wait() + + assert.True(t, c.isCommitted()) + assert.True(t, c.skipCheckpointEnabled()) +} + +func TestStopController_CloseForLoopExitClearsPendingCancel(t *testing.T) { + c := newStopController() + c.beginActiveTurn() + c.requestStop(&stopConfig{agentCancelOpts: []AgentCancelOption{WithRecursive()}}) + + c.closeForLoopExit() + + _, ok := c.receiveCancel() + assert.False(t, ok) +} + +func TestAttack_UntilIdleFor_ConcurrentPushDuringIdleTimer(t *testing.T) { + turnCount := int32(0) + turnDone := make(chan struct{}, 10) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + atomic.AddInt32(&turnCount, 1) + turnDone <- struct{}{} + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-turnDone + + loop.Stop(UntilIdleFor(200 * time.Millisecond)) + + for i := 0; i < 5; i++ { + time.Sleep(50 * time.Millisecond) + loop.Push("concurrent-" + string(rune('a'+i))) + <-turnDone + } + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + waitOrFail(t, done, "loop did not exit after idle timeout — Push did not reset timer correctly") + + finalCount := atomic.LoadInt32(&turnCount) + assert.Equal(t, int32(6), finalCount, "all 6 pushes should have been processed") +} + +func TestAttack_UntilIdleFor_MultipleStopCallsFirstWins(t *testing.T) { + turnDone := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(turnDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-turnDone + + loop.Stop(UntilIdleFor(100 * time.Millisecond)) + loop.Stop(UntilIdleFor(10 * time.Minute)) + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + waitOrFail(t, done, "second UntilIdleFor should have been ignored; loop should have exited with 100ms timer") +} + +func TestAttack_BareStopOverridesUntilIdleFor(t *testing.T) { + agentStarted := make(chan struct{}) + agentDone := make(chan struct{}) + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + <-agentDone + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-agentStarted + + loop.Stop(UntilIdleFor(10 * time.Minute)) + + loop.Stop() + close(agentDone) + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + waitOrFail(t, done, "bare Stop should override UntilIdleFor and cause immediate shutdown") + + exit := loop.Wait() + assert.NoError(t, exit.ExitReason, "bare Stop should exit cleanly") +} + +func TestAttack_BareStopDoesNotDeescalateExistingCancelIntent(t *testing.T) { + agentStarted := make(chan *cancelContext, 1) + probe := &turnLoopStopModeProbeAgent{ccCh: agentStarted} + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return probe, nil + }, + }) + + loop.Push("msg1") + cc := <-agentStarted + + loop.Stop(WithImmediate()) + + time.Sleep(20 * time.Millisecond) + + loop.Stop() + + time.Sleep(20 * time.Millisecond) + mode := cc.getMode() + assert.Equal(t, CancelImmediate, mode, "bare Stop after WithImmediate must not de-escalate cancel mode") + + exit := loop.Wait() + var ce *CancelError + require.True(t, errors.As(exit.ExitReason, &ce)) + assert.Equal(t, CancelImmediate, ce.Info.Mode) +} + +func TestAttack_InterruptedItems_EmptyWhenAgentFinishesNormally(t *testing.T) { + agentStarted := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-agentStarted + time.Sleep(50 * time.Millisecond) + loop.Stop() + + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + assert.Empty(t, exit.InterruptedItems, "InterruptedItems must be empty when agent finished normally") +} + +func TestAttack_TurnBuffer_WakeupDoesNotLoseItems(t *testing.T) { + tb := newTurnBuffer[string]() + + tb.Send("a") + tb.Send("b") + tb.Wakeup() + tb.Send("c") + + var got []string + for i := 0; i < 3; i++ { + val, ok := tb.Receive() + require.True(t, ok) + got = append(got, val) + } + + assert.Equal(t, []string{"a", "b", "c"}, got, "Wakeup must not cause items to be lost") +} + +func TestAttack_TurnBuffer_ClearWakeupPreventsSpuriousReturn(t *testing.T) { + tb := newTurnBuffer[string]() + + tb.Wakeup() + tb.ClearWakeup() + + received := make(chan string, 1) + go func() { + val, ok := tb.Receive() + if ok { + received <- val + } + }() + + time.Sleep(50 * time.Millisecond) + tb.Send("real") + + select { + case val := <-received: + assert.Equal(t, "real", val, "ClearWakeup should prevent spurious empty return") + case <-time.After(2 * time.Second): + t.Fatal("Receive blocked forever despite Send") + } +} + +func TestAttack_StopBeforeRun_UntilIdleFor_ExitsImmediately(t *testing.T) { + loop := NewTurnLoop(TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: prepareTestAgent, + }) + + loop.Stop(UntilIdleFor(10 * time.Minute)) + loop.Stop() + + loop.Run(context.Background()) + + done := make(chan struct{}) + go func() { + loop.Wait() + close(done) + }() + + waitOrFail(t, done, "loop should exit immediately when Stop() called before Run()") +} + +func TestAttack_PushAfterStop_UntilIdleFor_RoutedToLateItems(t *testing.T) { + turnDone := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(turnDone) + return &AgentOutput{}, nil + }, + }, nil + }, + }) + + loop.Push("msg1") + <-turnDone + + loop.Stop(UntilIdleFor(50 * time.Millisecond)) + exit := loop.Wait() + assert.NoError(t, exit.ExitReason) + + ok, _ := loop.Push("after-stop") + assert.False(t, ok, "Push after loop exited should return false") + + late := exit.TakeLateItems() + assert.Equal(t, []string{"after-stop"}, late) +} + +func TestAttack_ConcurrentStopEscalation_RaceDetector(t *testing.T) { + agentStarted := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + }) + + loop.Push("msg1") + <-agentStarted + + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + switch i % 4 { + case 0: + loop.Stop() + case 1: + loop.Stop(WithImmediate()) + case 2: + loop.Stop(WithGracefulTimeout(100 * time.Millisecond)) + case 3: + loop.Stop(UntilIdleFor(50 * time.Millisecond)) + } + }(i) + } + + wg.Wait() + exit := loop.Wait() + t.Log("ExitReason:", exit.ExitReason) +} + +func TestAttack_SkipCheckpoint_Sticky(t *testing.T) { + agentStarted := make(chan struct{}) + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + close(agentStarted) + <-ctx.Done() + return nil, ctx.Err() + }, + }, nil + }, + Store: newTestStore(), + CheckpointID: "test-sticky", + }) + + loop.Push("msg1") + <-agentStarted + + loop.Stop(WithSkipCheckpoint()) + loop.Stop(WithImmediate()) + + exit := loop.Wait() + assert.False(t, exit.CheckpointAttempted, "SkipCheckpoint is sticky; checkpoint should be skipped") +} + +// turnLoopNestedProbeAgent simulates an agent with a nested sub-agent +// by deriving a child cancelContext. This allows tests to verify that +// TurnLoop's Stop/Push options correctly propagate recursive cancellation. +// +// IMPORTANT: child.markDone() is NOT called by the probe. The test MUST +// call it (e.g. via t.Cleanup) after verifying propagation to avoid a +// race between markDone closing child.doneChan and the deriveAgentToolCancelContext +// goroutines propagating the cancel signal. +type turnLoopNestedProbeAgent struct { + parentCCCh chan *cancelContext + childCCCh chan *cancelContext +} + +func (a *turnLoopNestedProbeAgent) Name(_ context.Context) string { return "nested-probe" } +func (a *turnLoopNestedProbeAgent) Description(_ context.Context) string { return "nested-probe" } +func (a *turnLoopNestedProbeAgent) Run(ctx context.Context, _ *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + o := getCommonOptions(nil, opts...) + cc := o.cancelCtx + + child := cc.deriveAgentToolCancelContext(ctx) + a.parentCCCh <- cc + a.childCCCh <- child + + go func() { + defer gen.Close() + <-cc.cancelChan + for { + if cc.getMode() == CancelImmediate { + gen.Send(&AgentEvent{Err: cc.createCancelError()}) + return + } + time.Sleep(1 * time.Millisecond) + } + }() + return iter +} + +func TestTurnLoop_Stop_WithImmediate_RecursivePropagation(t *testing.T) { + parentCCCh := make(chan *cancelContext, 1) + childCCCh := make(chan *cancelContext, 1) + probe := &turnLoopNestedProbeAgent{parentCCCh: parentCCCh, childCCCh: childCCCh} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return probe, nil + }, + }) + + loop.Push("msg1") + cc := <-parentCCCh + child := <-childCCCh + t.Cleanup(func() { child.markDone() }) + + loop.Stop(WithImmediate()) + + // Child should receive the cancel signal via recursive propagation. + select { + case <-child.cancelChan: + case <-time.After(2 * time.Second): + t.Fatal("child did not receive cancel via recursive propagation") + } + + // Child should also receive the immediate cancel signal. + select { + case <-child.immediateChan: + case <-time.After(2 * time.Second): + t.Fatal("child did not receive immediate cancel via recursive propagation") + } + + assert.True(t, cc.isRecursive(), "WithImmediate should set recursive on parent") + assert.True(t, child.shouldCancel(), "child should be cancelled") + assert.True(t, child.isImmediateCancelled(), "child should have received immediate cancel") + + exit := loop.Wait() + var ce *CancelError + require.True(t, errors.As(exit.ExitReason, &ce)) + assert.Equal(t, CancelImmediate, ce.Info.Mode) +} + +func TestTurnLoop_Push_WithPreemptTimeout_RecursivePropagation(t *testing.T) { + parentCCCh := make(chan *cancelContext, 2) + childCCCh := make(chan *cancelContext, 2) + probe := &turnLoopNestedProbeAgent{parentCCCh: parentCCCh, childCCCh: childCCCh} + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return probe, nil + }, + }) + + loop.Push("first") + cc := <-parentCCCh + child := <-childCCCh + t.Cleanup(func() { child.markDone() }) + + // Preempt with a very short timeout so it escalates to CancelImmediate quickly. + loop.Push("urgent", WithPreemptTimeout[string, *schema.Message](AfterChatModel, 10*time.Millisecond)) + + // After timeout escalation, child should receive the immediate cancel + // via recursive propagation. + select { + case <-child.immediateChan: + case <-time.After(2 * time.Second): + t.Fatal("child did not receive immediate cancel after preempt timeout escalation") + } + + assert.True(t, cc.isRecursive(), "WithPreemptTimeout should set recursive on parent") + assert.True(t, child.isImmediateCancelled(), "child should have received immediate cancel") + + loop.Stop(WithImmediate()) + loop.Wait() +} + +func TestUntilIdleFor_NonPositive_Panics(t *testing.T) { + assert.PanicsWithValue(t, "adk: UntilIdleFor: duration must be positive", + func() { UntilIdleFor(0) }) + assert.PanicsWithValue(t, "adk: UntilIdleFor: duration must be positive", + func() { UntilIdleFor(-1 * time.Second) }) +} + +func TestSaveTurnLoopCheckpoint_NilStore(t *testing.T) { + l := &TurnLoop[string, *schema.Message]{config: TurnLoopConfig[string, *schema.Message]{Store: nil}} + err := l.saveTurnLoopCheckpoint(context.Background(), "cp-1", &turnLoopCheckpoint[string]{}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "checkpoint store is nil") +} + +func TestSetupBridgeStore_NilStore_Resume(t *testing.T) { + l := &TurnLoop[string, *schema.Message]{config: TurnLoopConfig[string, *schema.Message]{Store: nil}} + spec := &turnRunSpec[string, *schema.Message]{isResume: true} + _, _, err := l.setupBridgeStore(spec, nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "checkpoint store is nil") +} + +// TestTurnLoop_Preempt_LoopStalledAfterSecondPreemptPush covers a liveness +// regression where a preempted turn was followed by another preemptive Push and +// the loop stopped making progress before processing the later item. +func TestTurnLoop_Preempt_LoopStalledAfterSecondPreemptPush(t *testing.T) { + // turnCount tracks how many turns have been fully processed. + var turnCount int32 + + // Channels to synchronize the test with each turn's lifecycle. + firstAgentStarted := make(chan struct{}) + secondTurnDone := make(chan struct{}) + thirdTurnDone := make(chan struct{}) + + var firstAgentStartedOnce, secondTurnDoneOnce, thirdTurnDoneOnce sync.Once + + agent := &turnLoopCancellableMockAgent{ + name: "test", + runFunc: func(ctx context.Context, input *AgentInput) (*AgentOutput, error) { + turn := atomic.AddInt32(&turnCount, 1) + switch turn { + case 1: + // First turn: signal started, then block until preempted. + firstAgentStartedOnce.Do(func() { close(firstAgentStarted) }) + <-ctx.Done() + case 2, 3: + // Subsequent turns: complete immediately. + } + return &AgentOutput{}, nil + }, + } + + loop := newAndRunTurnLoop(context.Background(), TurnLoopConfig[string, *schema.Message]{ + PrepareAgent: prepareAgent(agent), + GenInput: genInputConsumeFirst, + OnAgentEvents: func(ctx context.Context, tc *TurnContext[string, *schema.Message], events *AsyncIterator[*AgentEvent]) error { + for { + if _, ok := events.Next(); !ok { + break + } + } + turn := atomic.LoadInt32(&turnCount) + switch turn { + case 2: + secondTurnDoneOnce.Do(func() { close(secondTurnDone) }) + case 3: + thirdTurnDoneOnce.Do(func() { close(thirdTurnDone) }) + } + return nil + }, + }) + + // Step 1: Push item A (no preempt). Wait for agent to start. + loop.Push("A") + waitOrFail(t, firstAgentStarted, "agent did not start for item A") + + // Step 2: Push item B with preempt. This cancels the first turn. + loop.Push("B", WithPreempt[string, *schema.Message](AnySafePoint)) + + // Wait for the second turn (item B) to complete successfully. + waitOrFail(t, secondTurnDone, "second turn (item B) did not complete") + + // Step 3: Push item C with preempt. This is the scenario that triggers + // the bug — the loop should process item C but instead gets stuck. + loop.Push("C", WithPreempt[string, *schema.Message](AnySafePoint)) + + // The loop should process item C. If the bug is present, this will timeout. + waitOrFail(t, thirdTurnDone, "third turn (item C) was never processed — loop is stuck between turns") + + loop.Stop() + result := loop.Wait() + assert.NoError(t, result.ExitReason) + assert.Equal(t, int32(3), atomic.LoadInt32(&turnCount), "expected 3 turns to be processed") +} + +func TestAttack_BusinessInterrupt_NoStore_ExitsWithoutPanic(t *testing.T) { + ctx := context.Background() + interruptAgent := &turnLoopInterruptAgent{interruptInfo: "no_store_test"} + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + GenInput: genInputConsumeAllWithMsg, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return interruptAgent, nil + }, + }) + + loop.Push("msg1") + exit := loop.Wait() + + var intErr *InterruptError + require.True(t, errors.As(exit.ExitReason, &intErr), "expected *InterruptError, got: %v", exit.ExitReason) + assert.Equal(t, []string{"msg1"}, exit.InterruptedItems) + assert.False(t, exit.CheckpointAttempted, "no store → no checkpoint attempt") +} + +func TestAttack_BusinessInterrupt_EmptyConsumed_NoCheckpoint(t *testing.T) { + ctx := context.Background() + store := newTestStore() + interruptAgent := &turnLoopInterruptAgent{interruptInfo: "idle_test"} + + loop := newAndRunTurnLoop(ctx, TurnLoopConfig[string, *schema.Message]{ + Store: store, + CheckpointID: "idle-cp", + GenInput: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], items []string) (*GenInputResult[string, *schema.Message], error) { + return &GenInputResult[string, *schema.Message]{ + Input: &AgentInput{Messages: []Message{schema.UserMessage("x")}}, + Consumed: []string{}, + }, nil + }, + PrepareAgent: func(ctx context.Context, _ *TurnLoop[string, *schema.Message], consumed []string) (Agent, error) { + return interruptAgent, nil + }, + }) + + loop.Push("msg1") + exit := loop.Wait() + + var intErr *InterruptError + require.True(t, errors.As(exit.ExitReason, &intErr), "expected *InterruptError, got: %v", exit.ExitReason) + assert.Empty(t, exit.InterruptedItems, "consumed was empty → InterruptedItems should be empty") +} diff --git a/adk/utils.go b/adk/utils.go new file mode 100644 index 0000000..991abd0 --- /dev/null +++ b/adk/utils.go @@ -0,0 +1,313 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "io" + "strings" + + "github.com/google/uuid" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/schema" +) + +type AsyncIterator[T any] struct { + ch *internal.UnboundedChan[T] +} + +func (ai *AsyncIterator[T]) Next() (T, bool) { + return ai.ch.Receive() +} + +type AsyncGenerator[T any] struct { + ch *internal.UnboundedChan[T] +} + +func (ag *AsyncGenerator[T]) Send(v T) { + ag.ch.Send(v) +} + +func (ag *AsyncGenerator[T]) trySend(v T) bool { + return ag.ch.TrySend(v) +} + +func (ag *AsyncGenerator[T]) Close() { + ag.ch.Close() +} + +// NewAsyncIteratorPair returns a paired async iterator and generator +// that share the same underlying channel. +func NewAsyncIteratorPair[T any]() (*AsyncIterator[T], *AsyncGenerator[T]) { + ch := internal.NewUnboundedChan[T]() + return &AsyncIterator[T]{ch}, &AsyncGenerator[T]{ch} +} + +func copyMap[K comparable, V any](m map[K]V) map[K]V { + res := make(map[K]V, len(m)) + for k, v := range m { + res[k] = v + } + return res +} + +func cloneSlice[T any](s []T) []T { + if s == nil { + return nil + } + res := make([]T, len(s)) + copy(res, s) + return res +} + +func concatInstructions(instructions ...string) string { + var sb strings.Builder + sb.WriteString(instructions[0]) + for i := 1; i < len(instructions); i++ { + sb.WriteString("\n\n") + sb.WriteString(instructions[i]) + } + + return sb.String() +} + +// GenTransferMessages generates assistant and tool messages to instruct a +// transfer-to-agent tool call targeting the destination agent. +// +// NOT RECOMMENDED: Agent transfer with full context sharing between agents has not proven +// to be more effective empirically. Consider using ChatModelAgent with AgentTool +// or DeepAgent instead for most multi-agent scenarios. +func GenTransferMessages(_ context.Context, destAgentName string) (Message, Message) { + toolCallID := uuid.NewString() + tooCall := schema.ToolCall{ID: toolCallID, Function: schema.FunctionCall{Name: TransferToAgentToolName, Arguments: destAgentName}} + assistantMessage := schema.AssistantMessage("", []schema.ToolCall{tooCall}) + msg := transferToAgentToolOutput(destAgentName) + toolMessage := schema.ToolMessage(msg, toolCallID, schema.WithToolName(TransferToAgentToolName)) + return assistantMessage, toolMessage +} + +func typedSetAutomaticClose[M MessageType](e *TypedAgentEvent[M]) { + if e.Output == nil || e.Output.MessageOutput == nil || !e.Output.MessageOutput.IsStreaming { + return + } + + e.Output.MessageOutput.MessageStream.SetAutomaticClose() +} + +// set automatic close for event's message stream +func setAutomaticClose(e *AgentEvent) { + typedSetAutomaticClose(e) +} + +// getMessageFromWrappedEvent extracts the message from an AgentEvent. +// If the stream contains an error chunk, this function returns (nil, err) and +// sets StreamErr to prevent re-consumption. The nil message ensures that +// failed stream responses are not included in subsequent agents' context windows. +func getMessageFromTypedWrappedEvent[M MessageType](e *typedAgentEventWrapper[M]) (M, error) { + var zero M + if e.event.Output == nil || e.event.Output.MessageOutput == nil { + return zero, nil + } + + if !e.event.Output.MessageOutput.IsStreaming { + return e.event.Output.MessageOutput.Message, nil + } + + if e.StreamErr != nil { + return zero, e.StreamErr + } + + if !isNilMessage(e.concatenatedMessage) { + return e.concatenatedMessage, nil + } + + e.consumeStream() + + if e.StreamErr != nil { + return zero, e.StreamErr + } + return e.concatenatedMessage, nil +} + +func getMessageFromWrappedEvent(e *agentEventWrapper) (Message, error) { + if e.AgentEvent.Output == nil || e.AgentEvent.Output.MessageOutput == nil { + return nil, nil + } + + if !e.AgentEvent.Output.MessageOutput.IsStreaming { + return e.AgentEvent.Output.MessageOutput.Message, nil + } + + if e.StreamErr != nil { + return nil, e.StreamErr + } + + if e.concatenatedMessage != nil { + return e.concatenatedMessage, nil + } + + e.consumeStream() + + if e.StreamErr != nil { + return nil, e.StreamErr + } + return e.concatenatedMessage, nil +} + +// consumeStream drains the message stream, setting concatenatedMessage on +// success or StreamErr on failure. The stream is always replaced with an +// error-free, materialized version safe for gob encoding. +// Must be called at most once (guarded by callers checking concatenatedMessage/StreamErr). +func (e *agentEventWrapper) consumeStream() { + e.mu.Lock() + defer e.mu.Unlock() + + if e.concatenatedMessage != nil { + return + } + + s := e.AgentEvent.Output.MessageOutput.MessageStream + var msgs []Message + + defer s.Close() + for { + msg, err := s.Recv() + if err != nil { + if err == io.EOF { + break + } + e.StreamErr = err + e.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs) + return + } + msgs = append(msgs, msg) + } + + if len(msgs) == 0 { + e.StreamErr = errors.New("no messages in MessageVariant.MessageStream") + // Defensively replace the stream. The defer s.Close() above already + // ensures subsequent Recv() returns io.EOF, but we replace it anyway + // to make the invariant explicit: after consumeStream, MessageStream + // is always safe for MessageVariant.GobEncode to consume. + e.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs) + return + } + + if len(msgs) == 1 { + e.concatenatedMessage = msgs[0] + } else { + var err error + e.concatenatedMessage, err = schema.ConcatMessages(msgs) + if err != nil { + e.StreamErr = err + e.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray(msgs) + return + } + } + + e.AgentEvent.Output.MessageOutput.MessageStream = schema.StreamReaderFromArray([]Message{e.concatenatedMessage}) +} + +// copyTypedAgentEvent copies a TypedAgentEvent. +// If the MessageVariant is streaming, the MessageStream will be copied. +// RunPath will be deep copied. +// The result of Copy will be a new TypedAgentEvent that is: +// - safe to set fields of TypedAgentEvent +// - safe to extend RunPath +// - safe to receive from MessageStream +// NOTE: even if the event is copied, it's still not recommended to modify +// the Message itself or Chunks of the MessageStream, as they are not copied. +// NOTE: if you have CustomizedOutput or CustomizedAction, they are NOT copied. +func copyTypedAgentEvent[M MessageType](ae *TypedAgentEvent[M]) *TypedAgentEvent[M] { + rp := make([]RunStep, len(ae.RunPath)) + copy(rp, ae.RunPath) + + copied := &TypedAgentEvent[M]{ + AgentName: ae.AgentName, + RunPath: rp, + Action: ae.Action, + Err: ae.Err, + } + + if ae.Output == nil { + return copied + } + + copied.Output = &TypedAgentOutput[M]{ + CustomizedOutput: ae.Output.CustomizedOutput, + } + + mv := ae.Output.MessageOutput + if mv == nil { + return copied + } + + copied.Output.MessageOutput = &TypedMessageVariant[M]{ + IsStreaming: mv.IsStreaming, + Role: mv.Role, + AgenticRole: mv.AgenticRole, + ToolName: mv.ToolName, + } + if mv.IsStreaming { + sts := ae.Output.MessageOutput.MessageStream.Copy(2) + mv.MessageStream = sts[0] + copied.Output.MessageOutput.MessageStream = sts[1] + } else { + copied.Output.MessageOutput.Message = mv.Message + } + + return copied +} + +// TypedGetMessage extracts the message from a TypedAgentEvent, concatenating a stream if present. +func TypedGetMessage[M MessageType](e *TypedAgentEvent[M]) (M, *TypedAgentEvent[M], error) { + var zero M + if e.Output == nil || e.Output.MessageOutput == nil { + return zero, e, nil + } + + msgOutput := e.Output.MessageOutput + if msgOutput.IsStreaming { + ss := msgOutput.MessageStream.Copy(2) + e.Output.MessageOutput.MessageStream = ss[0] + + msg, err := concatMessageStream(ss[1]) + + return msg, e, err + } + + return msgOutput.Message, e, nil +} + +// GetMessage extracts the Message from an AgentEvent. For streaming output, +// it duplicates the stream and concatenates it into a single Message. +func GetMessage(e *AgentEvent) (Message, *AgentEvent, error) { + return TypedGetMessage(e) +} + +func typedErrorIter[M MessageType](err error) *AsyncIterator[*TypedAgentEvent[M]] { + iterator, generator := NewAsyncIteratorPair[*TypedAgentEvent[M]]() + generator.Send(&TypedAgentEvent[M]{Err: err}) + generator.Close() + return iterator +} + +func genErrorIter(err error) *AsyncIterator[*AgentEvent] { + return typedErrorIter[*schema.Message](err) +} diff --git a/adk/utils_test.go b/adk/utils_test.go new file mode 100644 index 0000000..b4c2cbc --- /dev/null +++ b/adk/utils_test.go @@ -0,0 +1,455 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "bytes" + "encoding/gob" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestAsyncIteratorPair_Basic(t *testing.T) { + // Create a new iterator-generator pair + iterator, generator := NewAsyncIteratorPair[string]() + + // Test sending and receiving a value + generator.Send("test1") + val, ok := iterator.Next() + if !ok { + t.Error("receive should succeed") + } + if val != "test1" { + t.Errorf("expected 'test1', got '%s'", val) + } + + // Test sending and receiving multiple values + generator.Send("test2") + generator.Send("test3") + + val, ok = iterator.Next() + if !ok { + t.Error("receive should succeed") + } + if val != "test2" { + t.Errorf("expected 'test2', got '%s'", val) + } + + val, ok = iterator.Next() + if !ok { + t.Error("receive should succeed") + } + if val != "test3" { + t.Errorf("expected 'test3', got '%s'", val) + } +} + +func TestAsyncIteratorPair_Close(t *testing.T) { + iterator, generator := NewAsyncIteratorPair[int]() + + // Send some values + generator.Send(1) + generator.Send(2) + + // Close the generator + generator.Close() + + // Should still be able to read existing values + val, ok := iterator.Next() + if !ok { + t.Error("receive should succeed") + } + if val != 1 { + t.Errorf("expected 1, got %d", val) + } + + val, ok = iterator.Next() + if !ok { + t.Error("receive should succeed") + } + if val != 2 { + t.Errorf("expected 2, got %d", val) + } + + // After consuming all values, Next should return false + _, ok = iterator.Next() + if ok { + t.Error("receive from closed, empty channel should return ok=false") + } +} + +func TestAsyncIteratorPair_Concurrency(t *testing.T) { + iterator, generator := NewAsyncIteratorPair[int]() + const numSenders = 5 + const numReceivers = 3 + const messagesPerSender = 100 + + var rwg, swg sync.WaitGroup + rwg.Add(numReceivers) + swg.Add(numSenders) + + // Start senders + for i := 0; i < numSenders; i++ { + go func(id int) { + defer swg.Done() + for j := 0; j < messagesPerSender; j++ { + generator.Send(id*messagesPerSender + j) + time.Sleep(time.Microsecond) // Small delay to increase concurrency chance + } + }(i) + } + + // Start receivers + received := make([]int, 0, numSenders*messagesPerSender) + var mu sync.Mutex + + for i := 0; i < numReceivers; i++ { + go func() { + defer rwg.Done() + for { + val, ok := iterator.Next() + if !ok { + return + } + mu.Lock() + received = append(received, val) + mu.Unlock() + } + }() + } + + // Wait for senders to finish + swg.Wait() + generator.Close() + + // Wait for all goroutines to finish + rwg.Wait() + + // Verify we received all messages + if len(received) != numSenders*messagesPerSender { + t.Errorf("expected %d messages, got %d", numSenders*messagesPerSender, len(received)) + } + + // Create a map to check for duplicates and missing values + receivedMap := make(map[int]bool) + for _, val := range received { + receivedMap[val] = true + } + + if len(receivedMap) != numSenders*messagesPerSender { + t.Error("duplicate or missing messages detected") + } +} + +func TestGenErrorIter(t *testing.T) { + iter := genErrorIter(fmt.Errorf("test")) + e, ok := iter.Next() + assert.True(t, ok) + assert.Equal(t, "test", e.Err.Error()) + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestGetMessageFromWrappedEvent_StreamError_MultipleCallsGuard(t *testing.T) { + streamErr := errors.New("stream error") + + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("chunk1", nil), nil) + sw.Send(schema.AssistantMessage("chunk2", nil), nil) + sw.Send(nil, streamErr) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + } + + msg1, err1 := getMessageFromWrappedEvent(wrapper) + assert.Nil(t, msg1) + assert.NotNil(t, err1) + assert.Equal(t, "stream error", err1.Error()) + + assert.NotEmpty(t, wrapper.StreamErr) + assert.Equal(t, err1, wrapper.StreamErr) + + msg2, err2 := getMessageFromWrappedEvent(wrapper) + assert.Nil(t, msg2) + assert.Equal(t, err1, err2) +} + +func TestGetMessageFromWrappedEvent_StreamSuccess_MultipleCallsCached(t *testing.T) { + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("chunk1", nil), nil) + sw.Send(schema.AssistantMessage("chunk2", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + } + + msg1, err1 := getMessageFromWrappedEvent(wrapper) + assert.NotNil(t, msg1) + assert.Nil(t, err1) + assert.Equal(t, "chunk1chunk2", msg1.Content) + + assert.NotNil(t, wrapper.concatenatedMessage) + + msg2, err2 := getMessageFromWrappedEvent(wrapper) + assert.NotNil(t, msg2) + assert.Nil(t, err2) + assert.Equal(t, "chunk1chunk2", msg2.Content) + assert.Same(t, msg1, msg2) +} + +func TestGetMessageFromWrappedEvent_StreamError_PartialMessagesPreserved(t *testing.T) { + streamErr := errors.New("stream error at chunk3") + + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("chunk1", nil), nil) + sw.Send(schema.AssistantMessage("chunk2", nil), nil) + sw.Send(nil, streamErr) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + } + + _, err := getMessageFromWrappedEvent(wrapper) + assert.NotNil(t, err) + assert.Equal(t, streamErr, wrapper.StreamErr) + + newStream := wrapper.AgentEvent.Output.MessageOutput.MessageStream + assert.NotNil(t, newStream) + + var msgs []Message + for { + msg, err := newStream.Recv() + if err != nil { + break + } + msgs = append(msgs, msg) + } + + assert.Equal(t, 2, len(msgs)) + assert.Equal(t, "chunk1", msgs[0].Content) + assert.Equal(t, "chunk2", msgs[1].Content) +} + +func TestAgentEventWrapper_GobEncoding_WithWillRetryError(t *testing.T) { + streamErr := &WillRetryError{ErrStr: "stream error", RetryAttempt: 2} + + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("partial1", nil), nil) + sw.Send(schema.AssistantMessage("partial2", nil), nil) + sw.Send(nil, streamErr) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + TS: 12345, + } + + _, err := getMessageFromWrappedEvent(wrapper) + assert.NotNil(t, err) + var wrapperErr *WillRetryError + assert.True(t, errors.As(wrapper.StreamErr, &wrapperErr)) + assert.Equal(t, streamErr.ErrStr, wrapperErr.ErrStr) + assert.Equal(t, streamErr.RetryAttempt, wrapperErr.RetryAttempt) + + var buf bytes.Buffer + enc := gob.NewEncoder(&buf) + err = enc.Encode(wrapper) + assert.NoError(t, err) + + var decoded agentEventWrapper + dec := gob.NewDecoder(&buf) + err = dec.Decode(&decoded) + assert.NoError(t, err) + + assert.Equal(t, "TestAgent", decoded.AgentName) + assert.Equal(t, int64(12345), decoded.TS) + var decodedErr *WillRetryError + assert.True(t, errors.As(decoded.StreamErr, &decodedErr)) + assert.Equal(t, streamErr.ErrStr, decodedErr.ErrStr) + assert.Equal(t, streamErr.RetryAttempt, decodedErr.RetryAttempt) +} + +func TestAgentEventWrapper_GobEncoding_WithUnregisteredError(t *testing.T) { + streamErr := errors.New("unregistered error type") + + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("partial1", nil), nil) + sw.Send(nil, streamErr) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + TS: 22222, + } + + _, err := getMessageFromWrappedEvent(wrapper) + assert.NotNil(t, err) + assert.Equal(t, streamErr, wrapper.StreamErr) + + var buf bytes.Buffer + enc := gob.NewEncoder(&buf) + err = enc.Encode(wrapper) + assert.Error(t, err, "gob encoding should fail for unregistered error types") +} + +func TestAgentEventWrapper_GobEncoding_WithStreamSuccess(t *testing.T) { + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("success1", nil), nil) + sw.Send(schema.AssistantMessage("success2", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + AgentName: "TestAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + TS: 67890, + } + + msg, err := getMessageFromWrappedEvent(wrapper) + assert.NoError(t, err) + assert.Equal(t, "success1success2", msg.Content) + + var buf bytes.Buffer + enc := gob.NewEncoder(&buf) + err = enc.Encode(wrapper) + assert.NoError(t, err) + + var decoded agentEventWrapper + dec := gob.NewDecoder(&buf) + err = dec.Decode(&decoded) + assert.NoError(t, err) + + assert.Equal(t, "TestAgent", decoded.AgentName) + assert.Equal(t, int64(67890), decoded.TS) + assert.Empty(t, decoded.StreamErr) +} + +func TestConsumeStream_EdgeCases(t *testing.T) { + t.Run("double call is no-op", func(t *testing.T) { + sr, sw := schema.Pipe[Message](10) + go func() { + defer sw.Close() + sw.Send(schema.AssistantMessage("msg1", nil), nil) + }() + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + } + + wrapper.consumeStream() + assert.NotNil(t, wrapper.concatenatedMessage) + assert.Equal(t, "msg1", wrapper.concatenatedMessage.Content) + + // Second call should be a no-op + wrapper.consumeStream() + assert.Equal(t, "msg1", wrapper.concatenatedMessage.Content) + }) + + t.Run("empty stream sets StreamErr", func(t *testing.T) { + sr, sw := schema.Pipe[Message](10) + sw.Close() // immediately close => EOF with 0 messages + + wrapper := &agentEventWrapper{ + AgentEvent: &AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: true, + MessageStream: sr, + }, + }, + }, + } + + wrapper.consumeStream() + assert.Nil(t, wrapper.concatenatedMessage) + assert.Error(t, wrapper.StreamErr) + assert.Contains(t, wrapper.StreamErr.Error(), "no messages") + }) +} diff --git a/adk/workflow.go b/adk/workflow.go new file mode 100644 index 0000000..161c434 --- /dev/null +++ b/adk/workflow.go @@ -0,0 +1,706 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "fmt" + "runtime/debug" + "sync" + + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +type workflowAgentMode int + +const ( + workflowAgentModeUnknown workflowAgentMode = iota + workflowAgentModeSequential + workflowAgentModeLoop + workflowAgentModeParallel +) + +type workflowAgent struct { + name string + description string + subAgents []*flowAgent + + mode workflowAgentMode + + maxIterations int +} + +func (a *workflowAgent) Name(_ context.Context) string { + return a.name +} + +func (a *workflowAgent) Description(_ context.Context) string { + return a.description +} + +func (a *workflowAgent) GetType() string { + switch a.mode { + case workflowAgentModeSequential: + return "Sequential" + case workflowAgentModeParallel: + return "Parallel" + case workflowAgentModeLoop: + return "Loop" + default: + return "WorkflowAgent" + } +} + +func (a *workflowAgent) Run(ctx context.Context, _ *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + + var err error + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&AgentEvent{Err: e}) + } else if err != nil { + generator.Send(&AgentEvent{Err: err}) + } + + generator.Close() + }() + + // Different workflow execution based on mode + switch a.mode { + case workflowAgentModeSequential: + err = a.runSequential(ctx, generator, nil, nil, opts...) + case workflowAgentModeLoop: + err = a.runLoop(ctx, generator, nil, nil, opts...) + case workflowAgentModeParallel: + err = a.runParallel(ctx, generator, nil, nil, opts...) + default: + err = fmt.Errorf("unsupported workflow agent mode: %d", a.mode) + } + }() + + return iterator +} + +type sequentialWorkflowState struct { + InterruptIndex int +} + +type parallelWorkflowState struct { + SubAgentEvents map[int][]*agentEventWrapper +} + +type loopWorkflowState struct { + LoopIterations int + SubAgentIndex int +} + +func init() { + schema.RegisterName[*sequentialWorkflowState]("eino_adk_sequential_workflow_state") + schema.RegisterName[*parallelWorkflowState]("eino_adk_parallel_workflow_state") + schema.RegisterName[*loopWorkflowState]("eino_adk_loop_workflow_state") +} + +func (a *workflowAgent) Resume(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + var err error + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&AgentEvent{Err: e}) + } else if err != nil { + generator.Send(&AgentEvent{Err: err}) + } + + generator.Close() + }() + + state := info.InterruptState + if state == nil { + panic(fmt.Sprintf("workflowAgent.Resume: agent '%s' was asked to resume but has no state", a.Name(ctx))) + } + + // Different workflow execution based on the type of our restored state. + switch s := state.(type) { + case *sequentialWorkflowState: + err = a.runSequential(ctx, generator, s, info, opts...) + case *parallelWorkflowState: + err = a.runParallel(ctx, generator, s, info, opts...) + case *loopWorkflowState: + err = a.runLoop(ctx, generator, s, info, opts...) + default: + err = fmt.Errorf("unsupported workflow agent state type: %T", s) + } + }() + return iterator +} + +// WorkflowInterruptInfo stores interrupt information for workflow agents. +// CheckpointSchema: persisted via InterruptInfo.Data (gob). +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +type WorkflowInterruptInfo struct { + OrigInput *AgentInput + + SequentialInterruptIndex int + SequentialInterruptInfo *InterruptInfo + + LoopIterations int + + ParallelInterruptInfo map[int] /*index*/ *InterruptInfo +} + +func (a *workflowAgent) runSequential(ctx context.Context, + generator *AsyncGenerator[*AgentEvent], seqState *sequentialWorkflowState, info *ResumeInfo, + opts ...AgentRunOption) (err error) { + + startIdx := 0 + + seqCtx := ctx + + // If we are resuming, find which sub-agent to start from and prepare its context. + if seqState != nil { + startIdx = seqState.InterruptIndex + + var steps []string + for i := 0; i < startIdx; i++ { + steps = append(steps, a.subAgents[i].Name(seqCtx)) + } + + seqCtx = updateRunPathOnly(seqCtx, steps...) + } + + for i := startIdx; i < len(a.subAgents); i++ { + subAgent := a.subAgents[i] + + // Cancel check at transition boundary between sub-agents. + // Transition boundaries are always safe to cancel at — no sub-agent + // work is in progress, so any cancel mode is honoured. + if cancelCtx := getCancelContext(ctx); cancelCtx != nil && cancelCtx.shouldCancel() { + state := &sequentialWorkflowState{InterruptIndex: i} + event := cancelAtTransition(ctx, "Sequential workflow cancel at transition", state) + generator.Send(event) + return nil + } + + var subIterator *AsyncIterator[*AgentEvent] + if seqState != nil { + wfInfo, _ := info.Data.(*WorkflowInterruptInfo) + if wfInfo != nil && wfInfo.SequentialInterruptInfo != nil { + // Sub-agent was interrupted — resume it. + subIterator = subAgent.Resume(seqCtx, &ResumeInfo{ + EnableStreaming: info.EnableStreaming, + InterruptInfo: wfInfo.SequentialInterruptInfo, + }, opts...) + } else { + subIterator = subAgent.Run(seqCtx, nil, opts...) + } + seqState = nil + } else { + subIterator = subAgent.Run(seqCtx, nil, opts...) + } + + seqCtx = updateRunPathOnly(seqCtx, subAgent.Name(seqCtx)) + + var lastActionEvent *AgentEvent + for { + event, ok := subIterator.Next() + if !ok { + break + } + + if event.Err != nil { + // exit if report error + generator.Send(event) + return nil + } + + if lastActionEvent != nil { + generator.Send(lastActionEvent) + lastActionEvent = nil + } + + if event.Action != nil { + lastActionEvent = event + continue + } + generator.Send(event) + } + + if lastActionEvent != nil { + if lastActionEvent.Action.internalInterrupted != nil { + // A sub-agent interrupted. Wrap it with our own state, including the index. + state := &sequentialWorkflowState{ + InterruptIndex: i, + } + // Use CompositeInterrupt to funnel the sub-interrupt and add our own state. + // The context for the composite interrupt must be the one from *before* the sub-agent ran. + event := CompositeInterrupt(ctx, "Sequential workflow interrupted", state, + lastActionEvent.Action.internalInterrupted) + + // For backward compatibility, populate the deprecated Data field. + event.Action.Interrupted.Data = &WorkflowInterruptInfo{ + OrigInput: getRunCtx(ctx).RootInput, + SequentialInterruptIndex: i, + SequentialInterruptInfo: lastActionEvent.Action.Interrupted, + } + event.AgentName = lastActionEvent.AgentName + event.RunPath = lastActionEvent.RunPath + + generator.Send(event) + return nil + } + + if lastActionEvent.Action.Exit { + // Forward the event + generator.Send(lastActionEvent) + return nil + } + + generator.Send(lastActionEvent) + } + } + + return nil +} + +// BreakLoopAction is a programmatic-only agent action used to prematurely +// terminate the execution of a loop workflow agent. +// When a loop workflow agent receives this action from a sub-agent, it will stop its +// current iteration and will not proceed to the next one. +// It will mark the BreakLoopAction as Done, signalling to any 'upper level' loop agent +// that this action has been processed and should be ignored further up. +// This action is not intended to be used by LLMs. +type BreakLoopAction struct { + // From records the name of the agent that initiated the break loop action. + From string + // Done is a state flag that can be used by the framework to mark when the + // action has been handled. + Done bool + // CurrentIterations is populated by the framework to record at which + // iteration the loop was broken. + CurrentIterations int +} + +// NewBreakLoopAction creates a new BreakLoopAction, signaling a request +// to terminate the current loop. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +func NewBreakLoopAction(agentName string) *AgentAction { + return &AgentAction{BreakLoop: &BreakLoopAction{ + From: agentName, + }} +} + +func (a *workflowAgent) runLoop(ctx context.Context, generator *AsyncGenerator[*AgentEvent], + loopState *loopWorkflowState, resumeInfo *ResumeInfo, opts ...AgentRunOption) (err error) { + + if len(a.subAgents) == 0 { + return nil + } + + startIter := 0 + startIdx := 0 + + loopCtx := ctx + + if loopState != nil { + // We are resuming. + startIter = loopState.LoopIterations + startIdx = loopState.SubAgentIndex + + // Rebuild the loopCtx to have the correct RunPath up to the point of resumption. + var steps []string + for i := 0; i < startIter; i++ { + for _, subAgent := range a.subAgents { + steps = append(steps, subAgent.Name(loopCtx)) + } + } + for i := 0; i < startIdx; i++ { + steps = append(steps, a.subAgents[i].Name(loopCtx)) + } + loopCtx = updateRunPathOnly(loopCtx, steps...) + } + + for i := startIter; i < a.maxIterations || a.maxIterations == 0; i++ { + for j := startIdx; j < len(a.subAgents); j++ { + subAgent := a.subAgents[j] + + if cancelCtx := getCancelContext(ctx); cancelCtx != nil && cancelCtx.shouldCancel() { + state := &loopWorkflowState{LoopIterations: i, SubAgentIndex: j} + event := cancelAtTransition(ctx, "Loop workflow cancel at transition", state) + generator.Send(event) + return nil + } + + var subIterator *AsyncIterator[*AgentEvent] + if loopState != nil { + wfInfo, _ := resumeInfo.Data.(*WorkflowInterruptInfo) + if wfInfo != nil && wfInfo.SequentialInterruptInfo != nil { + // Sub-agent was interrupted — resume it. + subIterator = subAgent.Resume(loopCtx, &ResumeInfo{ + EnableStreaming: resumeInfo.EnableStreaming, + InterruptInfo: wfInfo.SequentialInterruptInfo, + }, opts...) + } else { + subIterator = subAgent.Run(loopCtx, nil, opts...) + } + loopState = nil // Only resume the first time. + } else { + subIterator = subAgent.Run(loopCtx, nil, opts...) + } + + loopCtx = updateRunPathOnly(loopCtx, subAgent.Name(loopCtx)) + + var lastActionEvent *AgentEvent + var breakLoopEvent *AgentEvent + for { + event, ok := subIterator.Next() + if !ok { + break + } + + if event.Err != nil { + generator.Send(event) + return nil + } + + if lastActionEvent != nil { + if lastActionEvent.Action.BreakLoop != nil && !lastActionEvent.Action.BreakLoop.Done { + lastActionEvent.Action.BreakLoop.Done = true + lastActionEvent.Action.BreakLoop.CurrentIterations = i + breakLoopEvent = lastActionEvent + } + generator.Send(lastActionEvent) + lastActionEvent = nil + } + + if event.Action != nil { + lastActionEvent = event + continue + } + generator.Send(event) + } + + if lastActionEvent != nil { + if lastActionEvent.Action.BreakLoop != nil && !lastActionEvent.Action.BreakLoop.Done { + lastActionEvent.Action.BreakLoop.Done = true + lastActionEvent.Action.BreakLoop.CurrentIterations = i + breakLoopEvent = lastActionEvent + } + + if lastActionEvent.Action.internalInterrupted != nil { + // A sub-agent interrupted. Wrap it with our own loop state. + state := &loopWorkflowState{ + LoopIterations: i, + SubAgentIndex: j, + } + // Use CompositeInterrupt to funnel the sub-interrupt and add our own state. + event := CompositeInterrupt(ctx, "Loop workflow interrupted", state, + lastActionEvent.Action.internalInterrupted) + + // For backward compatibility, populate the deprecated Data field. + event.Action.Interrupted.Data = &WorkflowInterruptInfo{ + OrigInput: getRunCtx(ctx).RootInput, + LoopIterations: i, + SequentialInterruptIndex: j, + SequentialInterruptInfo: lastActionEvent.Action.Interrupted, + } + event.AgentName = lastActionEvent.AgentName + event.RunPath = lastActionEvent.RunPath + + generator.Send(event) + return + } + + if lastActionEvent.Action.Exit { + generator.Send(lastActionEvent) + return + } + + generator.Send(lastActionEvent) + } + + if breakLoopEvent != nil { + return + } + } + + // Reset the sub-agent index for the next iteration of the outer loop. + startIdx = 0 + } + + return nil +} + +func (a *workflowAgent) runParallel(ctx context.Context, generator *AsyncGenerator[*AgentEvent], + parState *parallelWorkflowState, resumeInfo *ResumeInfo, opts ...AgentRunOption) error { + + if len(a.subAgents) == 0 { + return nil + } + + var ( + wg sync.WaitGroup + subInterruptSignals []*core.InterruptSignal + dataMap = make(map[int]*InterruptInfo) + mu sync.Mutex + agentNames map[string]bool + err error + childContexts = make([]context.Context, len(a.subAgents)) + ) + + // If resuming, get the scoped ResumeInfo for each child that needs to be resumed. + if parState != nil { + agentNames, err = getNextResumeAgents(ctx, resumeInfo) + if err != nil { + return err + } + } + + // Fork contexts for each sub-agent + for i := range a.subAgents { + childContexts[i] = forkRunCtx(ctx) + + // If we're resuming and this agent has existing events, add them to the child context + if parState != nil && parState.SubAgentEvents != nil { + if existingEvents, ok := parState.SubAgentEvents[i]; ok { + // Add existing events to the child's lane events + childRunCtx := getRunCtx(childContexts[i]) + if childRunCtx != nil && childRunCtx.Session != nil { + if childRunCtx.Session.LaneEvents == nil { + childRunCtx.Session.LaneEvents = &laneEvents{} + } + childRunCtx.Session.LaneEvents.Events = append(childRunCtx.Session.LaneEvents.Events, existingEvents...) + } + } + } + } + + // Cancel check before spawning parallel goroutines. No sub-agent work + // is in progress, so any cancel mode is honoured at this boundary. + if cancelCtx := getCancelContext(ctx); cancelCtx != nil && cancelCtx.shouldCancel() { + state := ¶llelWorkflowState{} + event := cancelAtTransition(ctx, "Parallel workflow cancel before spawn", state) + generator.Send(event) + return nil + } + + for i := range a.subAgents { + wg.Add(1) + go func(idx int, agent *flowAgent) { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + generator.Send(&AgentEvent{Err: e}) + } + wg.Done() + }() + + var iterator *AsyncIterator[*AgentEvent] + + if _, ok := agentNames[agent.Name(ctx)]; ok { + childResumeInfo := &ResumeInfo{ + EnableStreaming: resumeInfo.EnableStreaming, + } + if wfInfo, ok := resumeInfo.Data.(*WorkflowInterruptInfo); ok && wfInfo != nil { + childResumeInfo.InterruptInfo = wfInfo.ParallelInterruptInfo[idx] + } + iterator = agent.Resume(childContexts[idx], childResumeInfo, opts...) + } else if parState != nil { + // We are resuming, but this child is not in the next points map. + // This means it finished successfully, so we don't run it. + return + } else { + iterator = agent.Run(childContexts[idx], nil, opts...) + } + + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Action != nil && event.Action.internalInterrupted != nil { + mu.Lock() + subInterruptSignals = append(subInterruptSignals, event.Action.internalInterrupted) + dataMap[idx] = event.Action.Interrupted + mu.Unlock() + break + } + generator.Send(event) + } + }(i, a.subAgents[i]) + } + + wg.Wait() + + if len(subInterruptSignals) == 0 { + // Join all child contexts back to the parent + joinRunCtxs(ctx, childContexts...) + return nil + } + + if len(subInterruptSignals) > 0 { + // Before interrupting, collect the current events from each child context + subAgentEvents := make(map[int][]*agentEventWrapper) + for i, childCtx := range childContexts { + childRunCtx := getRunCtx(childCtx) + if childRunCtx != nil && childRunCtx.Session != nil && childRunCtx.Session.LaneEvents != nil { + subAgentEvents[i] = childRunCtx.Session.LaneEvents.Events + } + } + + state := ¶llelWorkflowState{ + SubAgentEvents: subAgentEvents, + } + event := CompositeInterrupt(ctx, "Parallel workflow interrupted", state, subInterruptSignals...) + + // For backward compatibility, populate the deprecated Data field. + event.Action.Interrupted.Data = &WorkflowInterruptInfo{ + OrigInput: getRunCtx(ctx).RootInput, + ParallelInterruptInfo: dataMap, + } + event.AgentName = a.Name(ctx) + event.RunPath = getRunCtx(ctx).RunPath + + generator.Send(event) + } + + return nil +} + +func cancelAtTransition(ctx context.Context, info string, state any) *AgentEvent { + // state is the workflow checkpoint state (e.g. sequentialWorkflowState); + // nil for subContexts because this is a leaf interrupt with no child signals. + is, err := core.Interrupt(ctx, info, state, nil, + core.WithLayerPayload(getRunCtx(ctx).RunPath)) + if err != nil { + return &AgentEvent{Err: err} + } + + contexts := core.ToInterruptContexts(is, allowedAddressSegmentTypes) + + return &AgentEvent{ + Action: &AgentAction{ + Interrupted: &InterruptInfo{ + InterruptContexts: contexts, + }, + internalInterrupted: is, + }, + } +} + +// SequentialAgentConfig is the configuration for NewSequentialAgent. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +type SequentialAgentConfig struct { + Name string + Description string + SubAgents []Agent +} + +// ParallelAgentConfig is the configuration for NewParallelAgent. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +type ParallelAgentConfig struct { + Name string + Description string + SubAgents []Agent +} + +// LoopAgentConfig is the configuration for NewLoopAgent. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +type LoopAgentConfig struct { + Name string + Description string + SubAgents []Agent + + MaxIterations int +} + +func newWorkflowAgent(ctx context.Context, name, desc string, + subAgents []Agent, mode workflowAgentMode, maxIterations int) (*flowAgent, error) { + + wa := &workflowAgent{ + name: name, + description: desc, + mode: mode, + + maxIterations: maxIterations, + } + + fas := make([]Agent, len(subAgents)) + for i, subAgent := range subAgents { + fas[i] = toFlowAgent(ctx, subAgent, WithDisallowTransferToParent()) + } + + fa, err := setSubAgents(ctx, wa, fas) + if err != nil { + return nil, err + } + + wa.subAgents = fa.subAgents + + return fa, nil +} + +// NewSequentialAgent creates an agent that runs sub-agents sequentially. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +func NewSequentialAgent(ctx context.Context, config *SequentialAgentConfig) (ResumableAgent, error) { + return newWorkflowAgent(ctx, config.Name, config.Description, config.SubAgents, workflowAgentModeSequential, 0) +} + +// NewParallelAgent creates an agent that runs sub-agents in parallel. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +func NewParallelAgent(ctx context.Context, config *ParallelAgentConfig) (ResumableAgent, error) { + return newWorkflowAgent(ctx, config.Name, config.Description, config.SubAgents, workflowAgentModeParallel, 0) +} + +// NewLoopAgent creates an agent that loops over sub-agents with a max iteration limit. +// +// NOT RECOMMENDED: Workflow agents are built on agent transfer with full context sharing, +// which has not proven to be more effective empirically. Consider using +// ChatModelAgent with AgentTool or DeepAgent instead for most multi-agent scenarios. +func NewLoopAgent(ctx context.Context, config *LoopAgentConfig) (ResumableAgent, error) { + return newWorkflowAgent(ctx, config.Name, config.Description, config.SubAgents, workflowAgentModeLoop, config.MaxIterations) +} diff --git a/adk/workflow_test.go b/adk/workflow_test.go new file mode 100644 index 0000000..6e68d21 --- /dev/null +++ b/adk/workflow_test.go @@ -0,0 +1,1361 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "fmt" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// mockAgent is a simple implementation of the Agent interface for testing +type mockAgent struct { + name string + description string + responses []*AgentEvent +} + +func (a *mockAgent) Name(_ context.Context) string { + return a.name +} + +func (a *mockAgent) Description(_ context.Context) string { + return a.description +} + +func (a *mockAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iterator, generator := NewAsyncIteratorPair[*AgentEvent]() + + go func() { + defer generator.Close() + + for _, event := range a.responses { + generator.Send(event) + + // If the event has an Exit action, stop sending events + if event.Action != nil && event.Action.Exit { + break + } + } + }() + + return iterator +} + +// newMockAgent creates a new mock agent with the given name, description, and responses +func newMockAgent(name, description string, responses []*AgentEvent) *mockAgent { + return &mockAgent{ + name: name, + description: description, + responses: responses, + } +} + +// TestSequentialAgent tests the sequential workflow agent +func TestSequentialAgent(t *testing.T) { + ctx := context.Background() + + // Create mock agents with predefined responses + agent1 := newMockAgent("Agent1", "First agent", []*AgentEvent{ + { + AgentName: "Agent1", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from Agent1", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + agent2 := newMockAgent("Agent2", "Second agent", []*AgentEvent{ + { + AgentName: "Agent2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from Agent2", nil), + Role: schema.Assistant, + }, + }}, + }) + + // Create a sequential agent with the mock agents + config := &SequentialAgentConfig{ + Name: "SequentialTestAgent", + Description: "Test sequential agent", + SubAgents: []Agent{agent1, agent2}, + } + + sequentialAgent, err := NewSequentialAgent(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, sequentialAgent) + + assert.Equal(t, "Test sequential agent", sequentialAgent.Description(ctx)) + + // Run the sequential agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + + // Initialize the run context + ctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input) + + iterator := sequentialAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // First event should be from agent1 + event1, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event1) + assert.Nil(t, event1.Err) + assert.NotNil(t, event1.Output) + assert.NotNil(t, event1.Output.MessageOutput) + + // Get the message content from agent1 + msg1 := event1.Output.MessageOutput.Message + assert.NotNil(t, msg1) + assert.Equal(t, "Response from Agent1", msg1.Content) + + // Second event should be from agent2 + event2, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event2) + assert.Nil(t, event2.Err) + assert.NotNil(t, event2.Output) + assert.NotNil(t, event2.Output.MessageOutput) + + // Get the message content from agent2 + msg2 := event2.Output.MessageOutput.Message + assert.NotNil(t, msg2) + assert.Equal(t, "Response from Agent2", msg2.Content) + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) +} + +// TestSequentialAgentWithExit tests the sequential workflow agent with an exit action +func TestSequentialAgentWithExit(t *testing.T) { + ctx := context.Background() + + // Create mock agents with predefined responses + agent1 := newMockAgent("Agent1", "First agent", []*AgentEvent{ + { + AgentName: "Agent1", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from Agent1", nil), + Role: schema.Assistant, + }, + }, + Action: &AgentAction{ + Exit: true, + }, + }, + }) + + agent2 := newMockAgent("Agent2", "Second agent", []*AgentEvent{ + { + AgentName: "Agent2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from Agent2", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // Create a sequential agent with the mock agents + config := &SequentialAgentConfig{ + Name: "SequentialTestAgent", + Description: "Test sequential agent", + SubAgents: []Agent{agent1, agent2}, + } + + sequentialAgent, err := NewSequentialAgent(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, sequentialAgent) + + // Run the sequential agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + + ctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input) + + iterator := sequentialAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // First event should be from agent1 with exit action + event1, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event1) + assert.Nil(t, event1.Err) + assert.NotNil(t, event1.Output) + assert.NotNil(t, event1.Output.MessageOutput) + assert.NotNil(t, event1.Action) + assert.True(t, event1.Action.Exit) + + // No more events due to exit action + _, ok = iterator.Next() + assert.False(t, ok) +} + +// TestParallelAgent tests the parallel workflow agent +func TestParallelAgent(t *testing.T) { + ctx := context.Background() + + // Create mock agents with predefined responses + agent1 := newMockAgent("Agent1", "First agent", []*AgentEvent{ + { + AgentName: "Agent1", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from Agent1", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + agent2 := newMockAgent("Agent2", "Second agent", []*AgentEvent{ + { + AgentName: "Agent2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Response from Agent2", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // Create a parallel agent with the mock agents + config := &ParallelAgentConfig{ + Name: "ParallelTestAgent", + Description: "Test parallel agent", + SubAgents: []Agent{agent1, agent2}, + } + + parallelAgent, err := NewParallelAgent(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, parallelAgent) + + // Run the parallel agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + + ctx, _ = initRunCtx(ctx, parallelAgent.Name(ctx), input) + + iterator := parallelAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Collect all events + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + // Should have two events, one from each agent + assert.Equal(t, 2, len(events)) + + // Verify the events + for _, event := range events { + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.NoError(t, err) + + // Check the source agent name and message content + switch event.AgentName { + case "Agent1": + assert.Equal(t, "Response from Agent1", msg.Content) + case "Agent2": + assert.Equal(t, "Response from Agent2", msg.Content) + default: + t.Fatalf("Unexpected source agent name: %s", event.AgentName) + } + } +} + +// TestLoopAgent tests the loop workflow agent +func TestLoopAgent(t *testing.T) { + ctx := context.Background() + + // Create a mock agent that will be called multiple times + agent := newMockAgent("LoopAgent", "Loop agent", []*AgentEvent{ + { + AgentName: "LoopAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Loop iteration", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + // Create a loop agent with the mock agent and max iterations set to 3 + config := &LoopAgentConfig{ + Name: "LoopTestAgent", + Description: "Test loop agent", + SubAgents: []Agent{agent}, + + MaxIterations: 3, + } + + loopAgent, err := NewLoopAgent(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, loopAgent) + + // Run the loop agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + + ctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input) + + iterator := loopAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Collect all events + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + // Should have 3 events (one for each iteration) + assert.Equal(t, 3, len(events)) + + // Verify all events + for _, event := range events { + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.Equal(t, "Loop iteration", msg.Content) + } +} + +// TestLoopAgentWithBreakLoop tests the loop workflow agent with an break loop action +func TestLoopAgentWithBreakLoop(t *testing.T) { + ctx := context.Background() + + // Create a mock agent that will break the loop after the first iteration + agent := newMockAgent("LoopAgent", "Loop agent", []*AgentEvent{ + { + AgentName: "LoopAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Loop iteration with break loop", nil), + Role: schema.Assistant, + }, + }, + Action: NewBreakLoopAction("LoopAgent"), + }, + }) + + // Create a loop agent with the mock agent and max iterations set to 3 + config := &LoopAgentConfig{ + Name: "LoopTestAgent", + Description: "Test loop agent", + SubAgents: []Agent{agent}, + MaxIterations: 3, + } + + loopAgent, err := NewLoopAgent(ctx, config) + assert.NoError(t, err) + assert.NotNil(t, loopAgent) + + // Run the loop agent + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + ctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input) + + iterator := loopAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Collect all events + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + // Should have only 1 event due to break loop action + assert.Equal(t, 1, len(events)) + + // Verify the event + event := events[0] + assert.Nil(t, event.Err) + assert.NotNil(t, event.Output) + assert.NotNil(t, event.Output.MessageOutput) + assert.NotNil(t, event.Action) + assert.NotNil(t, event.Action.BreakLoop) + assert.True(t, event.Action.BreakLoop.Done) + assert.Equal(t, "LoopAgent", event.Action.BreakLoop.From) + assert.Equal(t, 0, event.Action.BreakLoop.CurrentIterations) + + msg := event.Output.MessageOutput.Message + assert.NotNil(t, msg) + assert.Equal(t, "Loop iteration with break loop", msg.Content) +} + +// Add these test functions to the existing workflow_test.go file + +// Replace the existing TestWorkflowAgentPanicRecovery function +func TestWorkflowAgentPanicRecovery(t *testing.T) { + ctx := context.Background() + + // Create a panic agent that panics in Run method + panicAgent := &panicMockAgent{ + mockAgent: mockAgent{ + name: "PanicAgent", + description: "Agent that panics", + responses: []*AgentEvent{}, + }, + } + + // Create a sequential agent with the panic agent + config := &SequentialAgentConfig{ + Name: "PanicTestAgent", + Description: "Test agent with panic", + SubAgents: []Agent{panicAgent}, + } + + sequentialAgent, err := NewSequentialAgent(ctx, config) + assert.NoError(t, err) + + // Run the agent and expect panic recovery + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + + ctx, _ = initRunCtx(ctx, sequentialAgent.Name(ctx), input) + iterator := sequentialAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Should receive an error event due to panic recovery + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "panic") + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) +} + +// Add these new mock agent types that properly panic +type panicMockAgent struct { + mockAgent +} + +func (a *panicMockAgent) Run(_ context.Context, _ *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] { + panic("test panic in agent") +} + +func TestParallelWorkflowResumeWithEvents(t *testing.T) { + ctx := context.Background() + + // Create interruptible agents + sa1 := &myAgent{ + name: "sa1", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + // Send a normal message event first, called event1 + generator.Send(&AgentEvent{ + AgentName: "sa1", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa1 normal message"), + }, + }, + }) + intEvent := Interrupt(ctx, "sa1 interrupt data") + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.WasInterrupted) + assert.Nil(t, info.InterruptState) + assert.True(t, info.IsResumeTarget) + assert.Equal(t, "resume sa1", info.ResumeData) + + // Get the events from session and verify visibility + runCtx := getRunCtx(ctx) + assert.NotNil(t, runCtx.Session, "sa1 resumer should have session") + allEvents := runCtx.Session.getEvents() + + // Assert that allEvents only have 1 event, that is event1 + assert.Equal(t, 1, len(allEvents), "sa1 should only see its own event in session") + assert.Equal(t, "sa1", allEvents[0].AgentEvent.AgentName, "sa1 should see its own event") + assert.Equal(t, "sa1 normal message", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, "sa1 should see its own message content") + + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } + + sa2 := &myAgent{ + name: "sa2", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + // Send a normal message event first, called event2 + generator.Send(&AgentEvent{ + AgentName: "sa2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa2 normal message"), + }, + }, + }) + intEvent := StatefulInterrupt(ctx, "sa2 interrupt data", "sa2 interrupt") + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.WasInterrupted) + assert.NotNil(t, info.InterruptState) + assert.Equal(t, "sa2 interrupt", info.InterruptState) + assert.True(t, info.IsResumeTarget) + assert.Equal(t, "resume sa2", info.ResumeData) + + // Get the events from session and verify visibility + runCtx := getRunCtx(ctx) + assert.NotNil(t, runCtx.Session, "sa2 resumer should have session") + allEvents := runCtx.Session.getEvents() + + // Assert that allEvents only have 1 event, that is event2 + assert.Equal(t, 1, len(allEvents), "sa2 should only see its own event in session") + assert.Equal(t, "sa2", allEvents[0].AgentEvent.AgentName, "sa2 should see its own event") + assert.Equal(t, "sa2 normal message", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, "sa2 should see its own message content") + + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } + + sa3 := &myAgent{ + name: "sa3", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "sa3", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa3 completed"), + }, + }, + }) + generator.Close() + return iter + }, + } + + sa4 := &myAgent{ + name: "sa4", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "sa4", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("sa4 completed"), + }, + }, + }) + generator.Close() + return iter + }, + } + + t.Run("test parallel workflow agent", func(t *testing.T) { + // parallel + a, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "parallel agent", + SubAgents: []Agent{sa1, sa2, sa3, sa4}, + }) + assert.NoError(t, err) + runner := NewRunner(ctx, RunnerConfig{ + Agent: a, + CheckPointStore: newMyStore(), + }) + iter := runner.Query(ctx, "hello world", WithCheckPointID("1")) + var ( + events []*AgentEvent + interruptEvent *AgentEvent + ) + + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + continue + } + events = append(events, event) + } + assert.Equal(t, 4, len(events), "should have 4 events (2 normal messages + 2 completed agents)") + + // Verify specific properties of each event + var sa3Event, sa4Event *AgentEvent + for _, event := range events { + switch event.AgentName { + case "sa3": + sa3Event = event + case "sa4": + sa4Event = event + } + } + + // Verify sa3 event properties + assert.NotNil(t, sa3Event, "should have event from sa3") + assert.Equal(t, "sa3", sa3Event.AgentName, "sa3 event should have correct agent name") + assert.Equal(t, []RunStep{{"parallel agent"}, {"sa3"}}, sa3Event.RunPath, "sa3 event should have correct run path") + assert.NotNil(t, sa3Event.Output, "sa3 event should have output") + assert.NotNil(t, sa3Event.Output.MessageOutput, "sa3 event should have message output") + assert.Equal(t, "sa3 completed", sa3Event.Output.MessageOutput.Message.Content, "sa3 event should have correct message content") + + // Verify sa4 event properties + assert.NotNil(t, sa4Event, "should have event from sa4") + assert.Equal(t, "sa4", sa4Event.AgentName, "sa4 event should have correct agent name") + assert.Equal(t, []RunStep{{"parallel agent"}, {"sa4"}}, sa4Event.RunPath, "sa4 event should have correct run path") + assert.NotNil(t, sa4Event.Output, "sa4 event should have output") + assert.NotNil(t, sa4Event.Output.MessageOutput, "sa4 event should have message output") + assert.Equal(t, "sa4 completed", sa4Event.Output.MessageOutput.Message.Content, "sa4 event should have correct message content") + + assert.NotNil(t, interruptEvent) + assert.Equal(t, "parallel agent", interruptEvent.AgentName) + assert.Equal(t, []RunStep{{"parallel agent"}}, interruptEvent.RunPath) + assert.NotNil(t, interruptEvent.Action.Interrupted) + + var sa1InfoFound, sa2InfoFound bool + for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + switch ctx.Info { + case "sa1 interrupt data": + sa1InfoFound = true + case "sa2 interrupt data": + sa2InfoFound = true + } + } + + assert.Equal(t, 2, len(interruptEvent.Action.Interrupted.InterruptContexts)) + assert.True(t, sa1InfoFound) + assert.True(t, sa2InfoFound) + + var parallelInterruptID1, parallelInterruptID2 string + for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + switch ctx.Info { + case "sa1 interrupt data": + parallelInterruptID1 = ctx.ID + case "sa2 interrupt data": + parallelInterruptID2 = ctx.ID + } + } + assert.NotEmpty(t, parallelInterruptID1) + assert.NotEmpty(t, parallelInterruptID2) + + iter, err = runner.ResumeWithParams(ctx, "1", &ResumeParams{ + Targets: map[string]any{ + parallelInterruptID1: "resume sa1", + parallelInterruptID2: "resume sa2", + }, + }) + assert.NoError(t, err) + _, ok := iter.Next() + assert.False(t, ok) + }) +} + +func TestNestedParallelWorkflow(t *testing.T) { + ctx := context.Background() + + // Create predecessor agent that runs before the parallel structure + predecessorAgent := &myAgent{ + name: "predecessor", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "predecessor", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("predecessor completed"), + }, + }, + }) + generator.Close() + return iter + }, + } + + // Create interruptible inner agents + innerAgent1 := &myAgent{ + name: "inner1", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + + // Verify inner1 can see predecessor's event + runCtx := getRunCtx(ctx) + allEvents := runCtx.Session.getEvents() + assert.Equal(t, 1, len(allEvents), "inner1 should see exactly 1 event (predecessor)") + + assert.Equal(t, "predecessor", allEvents[0].AgentEvent.AgentName, "inner1 should see predecessor event") + assert.Equal(t, "predecessor completed", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, "inner1 should see predecessor message content") + + generator.Send(&AgentEvent{ + AgentName: "inner1", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("inner1 normal"), + }, + }, + }) + intEvent := Interrupt(ctx, "inner1 interrupt") + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.WasInterrupted) + assert.Equal(t, "resume inner1", info.ResumeData) + + // Verify inner1 can see predecessor's event during resume + runCtx := getRunCtx(ctx) + allEvents := runCtx.Session.getEvents() + assert.Equal(t, 2, len(allEvents), "inner1 should see exactly 2 events (predecessor + own normal message) during resume") + + // Find and verify predecessor event + var foundPredecessor bool + for _, event := range allEvents { + if event.AgentEvent != nil && event.AgentEvent.AgentName == "predecessor" { + foundPredecessor = true + assert.Equal(t, "predecessor completed", event.AgentEvent.Output.MessageOutput.Message.Content) + } + } + assert.True(t, foundPredecessor, "inner1 should see predecessor event during resume") + + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } + + innerAgent2 := &myAgent{ + name: "inner2", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + + // Verify inner2 can see predecessor's event + runCtx := getRunCtx(ctx) + allEvents := runCtx.Session.getEvents() + assert.Equal(t, 1, len(allEvents), "inner2 should see exactly 1 event (predecessor)") + + assert.Equal(t, "predecessor", allEvents[0].AgentEvent.AgentName, "inner2 should see predecessor event") + assert.Equal(t, "predecessor completed", allEvents[0].AgentEvent.Output.MessageOutput.Message.Content, "inner2 should see predecessor message content") + + generator.Send(&AgentEvent{ + AgentName: "inner2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("inner2 normal"), + }, + }, + }) + intEvent := StatefulInterrupt(ctx, "inner2 interrupt", "inner2 state") + generator.Send(intEvent) + generator.Close() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + assert.True(t, info.WasInterrupted) + assert.Equal(t, "inner2 state", info.InterruptState) + assert.Equal(t, "resume inner2", info.ResumeData) + + // Verify inner2 can see predecessor's event during resume + runCtx := getRunCtx(ctx) + allEvents := runCtx.Session.getEvents() + assert.Equal(t, 2, len(allEvents), "inner2 should see exactly 2 events (predecessor + own normal message) during resume") + + // Find and verify predecessor event + var foundPredecessor bool + for _, event := range allEvents { + if event.AgentEvent != nil && event.AgentEvent.AgentName == "predecessor" { + foundPredecessor = true + assert.Equal(t, "predecessor completed", event.AgentEvent.Output.MessageOutput.Message.Content) + } + } + assert.True(t, foundPredecessor, "inner2 should see predecessor event during resume") + + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Close() + return iter + }, + } + + // Create inner parallel workflow + innerParallel, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "inner parallel", + SubAgents: []Agent{innerAgent1, innerAgent2}, + }) + assert.NoError(t, err) + + // Create simple outer agents + outerAgent1 := &myAgent{ + name: "outer1", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "outer1", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("outer1 completed"), + }, + }, + }) + generator.Close() + return iter + }, + } + + outerAgent2 := &myAgent{ + name: "outer2", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + generator.Send(&AgentEvent{ + AgentName: "outer2", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("outer2 completed"), + }, + }, + }) + generator.Close() + return iter + }, + } + + // Create outer parallel workflow with nested parallel agent + outerParallel, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + Name: "outer parallel", + SubAgents: []Agent{outerAgent1, innerParallel, outerAgent2}, + }) + assert.NoError(t, err) + + // Create successor agent that runs after the parallel structure + successorAgent := &myAgent{ + name: "successor", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + + // Verify successor can see all events from predecessor and parallel agents + runCtx := getRunCtx(ctx) + allEvents := runCtx.Session.getEvents() + assert.GreaterOrEqual(t, len(allEvents), 5, "successor should see all events") + + var foundPredecessor, foundOuter1, foundOuter2, foundInner1, foundInner2 bool + for _, event := range allEvents { + if event.AgentEvent != nil { + switch event.AgentEvent.AgentName { + case "predecessor": + foundPredecessor = true + assert.Equal(t, "predecessor completed", event.AgentEvent.Output.MessageOutput.Message.Content) + case "outer1": + foundOuter1 = true + case "outer2": + foundOuter2 = true + case "inner1": + foundInner1 = true + case "inner2": + foundInner2 = true + } + } + } + + assert.True(t, foundPredecessor, "successor should see predecessor event") + assert.True(t, foundOuter1, "successor should see outer1 event") + assert.True(t, foundOuter2, "successor should see outer2 event") + assert.True(t, foundInner1, "successor should see inner1 event") + assert.True(t, foundInner2, "successor should see inner2 event") + + generator.Send(&AgentEvent{ + AgentName: "successor", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.UserMessage("successor completed"), + }, + }, + }) + generator.Close() + return iter + }, + } + + // Create sequential workflow: predecessor -> parallel -> successor + sequentialWorkflow, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + Name: "sequential workflow", + SubAgents: []Agent{predecessorAgent, outerParallel, successorAgent}, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: sequentialWorkflow, + CheckPointStore: newMyStore(), + }) + + iter := runner.Query(ctx, "test nested parallel with predecessor and successor", WithCheckPointID("nested-parallel-test")) + + var events []*AgentEvent + var interruptEvent *AgentEvent + for event, ok := iter.Next(); ok; event, ok = iter.Next() { + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + continue + } + events = append(events, event) + } + + // Should get events from predecessor, outer agents, and inner normal messages (successor doesn't run due to interruption) + assert.Equal(t, 5, len(events), "should have 5 events (predecessor + 2 outer + 2 inner)") + if interruptEvent == nil { + t.Fatal("should have interrupt event") + } + + // Resume the inner parallel workflow + var innerInterruptID1, innerInterruptID2 string + for _, ctx := range interruptEvent.Action.Interrupted.InterruptContexts { + switch ctx.Info { + case "inner1 interrupt": + innerInterruptID1 = ctx.ID + case "inner2 interrupt": + innerInterruptID2 = ctx.ID + } + } + + iter, err = runner.ResumeWithParams(ctx, "nested-parallel-test", &ResumeParams{ + Targets: map[string]any{ + innerInterruptID1: "resume inner1", + innerInterruptID2: "resume inner2", + }, + }) + assert.NoError(t, err) + + // Verify resume completes successfully and successor runs + var resumeEvents []*AgentEvent + for event, ok := iter.Next(); ok; event, ok = iter.Next() { + resumeEvents = append(resumeEvents, event) + } + + // Should get successor event after resume + assert.Equal(t, 1, len(resumeEvents), "should have successor event after resume") + assert.Equal(t, "successor", resumeEvents[0].AgentName) +} + +// TestWorkflowAgentUnsupportedMode tests unsupported workflow mode error (lines 65-71) +func TestWorkflowAgentUnsupportedMode(t *testing.T) { + ctx := context.Background() + + // Create a workflow agent with unsupported mode + agent := &workflowAgent{ + name: "UnsupportedModeAgent", + description: "Agent with unsupported mode", + subAgents: []*flowAgent{}, + mode: workflowAgentMode(999), + } + + // Run the agent and expect error + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + + ctx, _ = initRunCtx(ctx, agent.Name(ctx), input) + iterator := agent.Run(ctx, input) + assert.NotNil(t, iterator) + + // Should receive an error event due to unsupported mode + event, ok := iterator.Next() + assert.True(t, ok) + assert.NotNil(t, event) + assert.NotNil(t, event.Err) + assert.Contains(t, event.Err.Error(), "unsupported workflow agent mode") + + // No more events + _, ok = iterator.Next() + assert.False(t, ok) +} + +func TestFilterOptions(t *testing.T) { + a1 := &myAgent{ + name: "Agent1", + runFn: func(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + o := GetImplSpecificOptions[myAgentOptions](nil, opts...) + assert.Equal(t, "Agent1", o.value) + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Close() + return iter + }, + } + a2 := &myAgent{ + name: "Agent2", + runFn: func(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + o := GetImplSpecificOptions[myAgentOptions](nil, opts...) + assert.Equal(t, "Agent2", o.value) + iter, gen := NewAsyncIteratorPair[*AgentEvent]() + gen.Close() + return iter + }, + } + ctx := context.Background() + // sequential + seqAgent, err := NewSequentialAgent(ctx, &SequentialAgentConfig{ + SubAgents: []Agent{a1, a2}, + }) + assert.NoError(t, err) + iter := seqAgent.Run(ctx, &AgentInput{}, withValue("Agent1").DesignateAgent("Agent1"), withValue("Agent2").DesignateAgent("Agent2")) + _, ok := iter.Next() + assert.False(t, ok) + + // parallel + parAgent, err := NewParallelAgent(ctx, &ParallelAgentConfig{ + SubAgents: []Agent{a1, a2}, + }) + assert.NoError(t, err) + iter = parAgent.Run(ctx, &AgentInput{}, withValue("Agent1").DesignateAgent("Agent1"), withValue("Agent2").DesignateAgent("Agent2")) + _, ok = iter.Next() + assert.False(t, ok) +} + +func TestLoopAgentWithError(t *testing.T) { + ctx := context.Background() + + iterationCount := 0 + agent := &myAgent{ + name: "ErrorAgent", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer generator.Close() + iterationCount++ + if iterationCount == 3 { + generator.Send(&AgentEvent{Err: fmt.Errorf("error on iteration %d", iterationCount)}) + return + } + generator.Send(&AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage(fmt.Sprintf("iteration %d", iterationCount), nil), + Role: schema.Assistant, + }, + }, + }) + }() + return iter + }, + } + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "LoopErrorTestAgent", + SubAgents: []Agent{agent}, + MaxIterations: 10, + }) + assert.NoError(t, err) + + input := &AgentInput{Messages: []Message{schema.UserMessage("test")}} + ctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input) + iterator := loopAgent.Run(ctx, input) + + var events []*AgentEvent + var errorEvent *AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Err != nil { + errorEvent = event + } else { + events = append(events, event) + } + } + + assert.Equal(t, 2, len(events), "should have 2 successful iterations before error") + assert.NotNil(t, errorEvent, "should have received error event") + assert.Contains(t, errorEvent.Err.Error(), "error on iteration 3") + assert.Equal(t, 3, iterationCount, "loop should stop at iteration 3") +} + +func TestWorkflowCallbackHandlerNotDoubled(t *testing.T) { + ctx := context.Background() + store := newMyStore() + + var globalCallbackCount int + var designatedCallbackCount int + var mu sync.Mutex + + globalHandler := callbacks.NewHandlerBuilder().OnStartFn( + func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component == ComponentOfAgent && info.Name == "SubSubAgent" { + mu.Lock() + globalCallbackCount++ + mu.Unlock() + } + return ctx + }).Build() + + designatedHandler := callbacks.NewHandlerBuilder().OnStartFn( + func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Component == ComponentOfAgent && info.Name == "SubSubAgent" { + mu.Lock() + designatedCallbackCount++ + mu.Unlock() + } + return ctx + }).Build() + + iterationCount := 0 + shouldInterrupt := true + subSubAgent := &myAgent{ + name: "SubSubAgent", + runFn: func(ctx context.Context, input *AgentInput, options ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer generator.Close() + iterationCount++ + if shouldInterrupt && iterationCount == 2 { + generator.Send(Interrupt(ctx, "test_interrupt")) + return + } + generator.Send(&AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage(fmt.Sprintf("iteration %d", iterationCount), nil), + Role: schema.Assistant, + }, + }, + }) + }() + return iter + }, + resumeFn: func(ctx context.Context, info *ResumeInfo, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] { + iter, generator := NewAsyncIteratorPair[*AgentEvent]() + go func() { + defer generator.Close() + iterationCount++ + generator.Send(&AgentEvent{ + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + Message: schema.AssistantMessage(fmt.Sprintf("resumed iteration %d", iterationCount), nil), + Role: schema.Assistant, + }, + }, + }) + }() + return iter + }, + } + + subWorkflow, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "SubWorkflow", + SubAgents: []Agent{subSubAgent}, + MaxIterations: 2, + }) + assert.NoError(t, err) + + parentWorkflow, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "ParentWorkflow", + SubAgents: []Agent{subWorkflow}, + MaxIterations: 2, + }) + assert.NoError(t, err) + + runner := NewRunner(ctx, RunnerConfig{ + Agent: parentWorkflow, + CheckPointStore: store, + }) + + opts := []AgentRunOption{ + WithCallbacks(globalHandler), + WithCallbacks(designatedHandler).DesignateAgent("ParentWorkflow", "SubSubAgent"), + WithCheckPointID("cp1"), + } + + iterator := runner.Run(ctx, []Message{schema.UserMessage("test")}, opts...) + + var interruptEvent *AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + if event.Action != nil && event.Action.Interrupted != nil { + interruptEvent = event + } + } + + assert.NotNil(t, interruptEvent) + assert.Equal(t, 2, iterationCount) + assert.Equal(t, 2, globalCallbackCount) + assert.Equal(t, 2, designatedCallbackCount) + + shouldInterrupt = false + var rootCauseID string + for _, intCtx := range interruptEvent.Action.Interrupted.InterruptContexts { + if intCtx.IsRootCause { + rootCauseID = intCtx.ID + break + } + } + + resumeIter, err := runner.ResumeWithParams(ctx, "cp1", &ResumeParams{ + Targets: map[string]any{rootCauseID: nil}, + }, opts...) + assert.NoError(t, err) + + for { + _, ok := resumeIter.Next() + if !ok { + break + } + } + + assert.Equal(t, 5, iterationCount) + assert.Equal(t, 5, globalCallbackCount) + assert.Equal(t, 5, designatedCallbackCount) +} + +func TestLoopAgentWithBreakLoopFollowedByMoreEvents(t *testing.T) { + ctx := context.Background() + + agent := newMockAgent("SubAgent", "Sub agent", []*AgentEvent{ + { + AgentName: "SubAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.ToolMessage("tool result", "call_123"), + Role: schema.Tool, + }, + }, + Action: NewBreakLoopAction("SubAgent"), + }, + { + AgentName: "SubAgent", + Output: &AgentOutput{ + MessageOutput: &MessageVariant{ + IsStreaming: false, + Message: schema.AssistantMessage("Final response after tool", nil), + Role: schema.Assistant, + }, + }, + }, + }) + + loopAgent, err := NewLoopAgent(ctx, &LoopAgentConfig{ + Name: "LoopTestAgent", + Description: "Test loop agent", + SubAgents: []Agent{agent}, + MaxIterations: 3, + }) + assert.NoError(t, err) + assert.NotNil(t, loopAgent) + + input := &AgentInput{ + Messages: []Message{ + schema.UserMessage("Test input"), + }, + } + ctx, _ = initRunCtx(ctx, loopAgent.Name(ctx), input) + + iterator := loopAgent.Run(ctx, input) + assert.NotNil(t, iterator) + + var events []*AgentEvent + for { + event, ok := iterator.Next() + if !ok { + break + } + events = append(events, event) + } + + assert.Equal(t, 2, len(events), "should have 2 events (tool event with BreakLoop + final response) and loop should break") + + assert.NotNil(t, events[0].Action, "first event should have an action") + assert.NotNil(t, events[0].Action.BreakLoop, "first event should have BreakLoop action") + assert.True(t, events[0].Action.BreakLoop.Done, "BreakLoop should be marked as Done") + assert.Equal(t, "SubAgent", events[0].Action.BreakLoop.From) + assert.Equal(t, 0, events[0].Action.BreakLoop.CurrentIterations) + assert.Equal(t, schema.Tool, events[0].Output.MessageOutput.Role, "first event should be tool message") + + assert.Nil(t, events[1].Action, "second event should not have an action") + assert.Equal(t, schema.Assistant, events[1].Output.MessageOutput.Role, "second event should be assistant message") + assert.Equal(t, "Final response after tool", events[1].Output.MessageOutput.Message.Content) +} diff --git a/adk/wrappers.go b/adk/wrappers.go new file mode 100644 index 0000000..ef5c035 --- /dev/null +++ b/adk/wrappers.go @@ -0,0 +1,1448 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "io" + "reflect" + "sync" + + "github.com/google/uuid" + + "github.com/cloudwego/eino/adk/internal" + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +type typedGenerateEndpoint[M MessageType] func(ctx context.Context, input []M, opts ...model.Option) (M, error) +type typedStreamEndpoint[M MessageType] func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) + +type typedModelWrapperConfig[M MessageType] struct { + handlers []TypedChatModelAgentMiddleware[M] + middlewares []AgentMiddleware + retryConfig *TypedModelRetryConfig[M] + failoverConfig *ModelFailoverConfig[M] + toolInfos []*schema.ToolInfo + cancelContext *cancelContext +} + +type modelWrapperConfig = typedModelWrapperConfig[*schema.Message] + +func buildModelWrappers[M MessageType](m model.BaseModel[M], config *typedModelWrapperConfig[M]) model.BaseModel[M] { + return buildModelWrappersImpl(m, config) +} + +func buildModelWrappersImpl[M MessageType](m model.BaseModel[M], config *typedModelWrapperConfig[M]) model.BaseModel[M] { + var wrapped = m + + if config.failoverConfig != nil { + wrapped = &typedFailoverProxyModel[M]{} + } + + if !components.IsCallbacksEnabled(wrapped) { + wrapped = typedCallbackInjectionModelWrapper[M]{}.wrapModel(wrapped) + } + + wrapped = &typedStateModelWrapper[M]{ + inner: wrapped, + original: m, + handlers: config.handlers, + middlewares: config.middlewares, + toolInfos: config.toolInfos, + modelRetryConfig: config.retryConfig, + modelFailoverConfig: config.failoverConfig, + cancelContext: config.cancelContext, + } + + return wrapped +} + +type typedCallbackInjectionModelWrapper[M MessageType] struct{} + +func (w typedCallbackInjectionModelWrapper[M]) wrapModel(m model.BaseModel[M]) model.BaseModel[M] { + return &typedCallbackInjectedModel[M]{inner: m} +} + +type typedCallbackInjectedModel[M MessageType] struct { + inner model.BaseModel[M] +} + +func (m *typedCallbackInjectedModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + ctx = callbacks.OnStart(ctx, input) + result, err := m.inner.Generate(ctx, input, opts...) + if err != nil { + callbacks.OnError(ctx, err) + var zero M + return zero, err + } + callbacks.OnEnd(ctx, result) + return result, nil +} + +func (m *typedCallbackInjectedModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + ctx = callbacks.OnStart(ctx, input) + result, err := m.inner.Stream(ctx, input, opts...) + if err != nil { + callbacks.OnError(ctx, err) + return nil, err + } + _, wrappedStream := callbacks.OnEndWithStreamOutput(ctx, result) + return wrappedStream, nil +} + +func handlersToToolMiddlewares[M MessageType](handlers []TypedChatModelAgentMiddleware[M]) []compose.ToolMiddleware { + var middlewares []compose.ToolMiddleware + // Forward iteration: compose.wrapToolCall applies middlewares in reverse order + // (len-1 down to 0), so keeping the original handler order here means + // handlers[0] ends up outermost — matching the model wrapping convention. + for _, handler := range handlers { + + m := compose.ToolMiddleware{} + + h := handler + m.Invokable = func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + tCtx := &ToolContext{ + Name: input.Name, + CallID: input.CallID, + } + wrappedEndpoint, err := h.WrapInvokableToolCall( + ctx, + func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + output, err := next(ctx, &compose.ToolInput{ + Name: input.Name, + CallID: input.CallID, + Arguments: argumentsInJSON, + CallOptions: opts, + }) + if err != nil { + return "", err + } + return output.Result, nil + }, + tCtx, + ) + if err != nil { + return nil, err + } + result, err := wrappedEndpoint(ctx, input.Arguments, input.CallOptions...) + if err != nil { + return nil, err + } + return &compose.ToolOutput{Result: result}, nil + } + } + + m.Streamable = func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + tCtx := &ToolContext{ + Name: input.Name, + CallID: input.CallID, + } + wrappedEndpoint, err := h.WrapStreamableToolCall( + ctx, + func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + output, err := next(ctx, &compose.ToolInput{ + Name: input.Name, + CallID: input.CallID, + Arguments: argumentsInJSON, + CallOptions: opts, + }) + if err != nil { + return nil, err + } + return output.Result, nil + }, + tCtx, + ) + if err != nil { + return nil, err + } + result, err := wrappedEndpoint(ctx, input.Arguments, input.CallOptions...) + if err != nil { + return nil, err + } + return &compose.StreamToolOutput{Result: result}, nil + } + } + + m.EnhancedInvokable = func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + tCtx := &ToolContext{ + Name: input.Name, + CallID: input.CallID, + } + wrappedEndpoint, err := h.WrapEnhancedInvokableToolCall( + ctx, + func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + output, err := next(ctx, &compose.ToolInput{ + Name: input.Name, + CallID: input.CallID, + Arguments: toolArgument.Text, + CallOptions: opts, + }) + if err != nil { + return nil, err + } + return output.Result, nil + }, + tCtx, + ) + if err != nil { + return nil, err + } + result, err := wrappedEndpoint(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...) + if err != nil { + return nil, err + } + return &compose.EnhancedInvokableToolOutput{Result: result}, nil + } + } + + m.EnhancedStreamable = func(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + tCtx := &ToolContext{ + Name: input.Name, + CallID: input.CallID, + } + wrappedEndpoint, err := h.WrapEnhancedStreamableToolCall( + ctx, + func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + output, err := next(ctx, &compose.ToolInput{ + Name: input.Name, + CallID: input.CallID, + Arguments: toolArgument.Text, + CallOptions: opts, + }) + if err != nil { + return nil, err + } + return output.Result, nil + }, + tCtx, + ) + if err != nil { + return nil, err + } + result, err := wrappedEndpoint(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...) + if err != nil { + return nil, err + } + return &compose.EnhancedStreamableToolOutput{Result: result}, nil + } + } + + middlewares = append(middlewares, m) + } + return middlewares +} + +type typedEventSenderModelWrapper[M MessageType] struct { + *TypedBaseChatModelAgentMiddleware[M] +} + +// NewEventSenderModelWrapper creates a ChatModelAgentMiddleware that sends model output as agent events. +func NewEventSenderModelWrapper() ChatModelAgentMiddleware { + return &typedEventSenderModelWrapper[*schema.Message]{ + TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[*schema.Message]{}, + } +} + +func (w *typedEventSenderModelWrapper[M]) WrapModel(_ context.Context, m model.BaseModel[M], mc *TypedModelContext[M]) (model.BaseModel[M], error) { + inner := m + if mc != nil && mc.cancelContext != nil { + inner = &typedCancelMonitoredModel[M]{ + inner: inner, + cancelContext: mc.cancelContext, + } + } + var retryConfig *TypedModelRetryConfig[M] + if mc != nil { + retryConfig = mc.ModelRetryConfig + } + var failoverConfig *ModelFailoverConfig[M] + if mc != nil { + failoverConfig = mc.ModelFailoverConfig + } + return &typedEventSenderModel[M]{inner: inner, modelRetryConfig: retryConfig, modelFailoverConfig: failoverConfig}, nil +} + +type typedEventSenderModel[M MessageType] struct { + inner model.BaseModel[M] + modelRetryConfig *TypedModelRetryConfig[M] + modelFailoverConfig *ModelFailoverConfig[M] +} + +func (m *typedEventSenderModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + result, err := m.inner.Generate(ctx, input, opts...) + if err != nil { + var zero M + return zero, err + } + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx != nil && execCtx.suppressEventSend { + return result, nil + } + if execCtx == nil || execCtx.generator == nil { + var zero M + return zero, errors.New("generator is nil when sending event in Generate: ensure agent state is properly initialized") + } + + event := typedModelOutputEvent(copyMessage(result), nil) + execCtx.send(event) + + return result, nil +} + +func (m *typedEventSenderModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + result, err := m.inner.Stream(ctx, input, opts...) + if err != nil { + return nil, err + } + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil || execCtx.generator == nil { + result.Close() + return nil, errors.New("generator is nil when sending event in Stream: ensure agent state is properly initialized") + } + + streams := result.Copy(2) + + eventStream := streams[0] + if convertOpts := m.buildStreamConvertOptions(ctx); len(convertOpts) > 0 { + eventStream = schema.StreamReaderWithConvert(streams[0], + func(msg M) (M, error) { return msg, nil }, + convertOpts...) + } + + var zero M + event := typedModelOutputEvent[M](zero, eventStream) + execCtx.send(event) + + return streams[1], nil +} + +// buildStreamConvertOptions constructs ConvertOption hooks that gate stream termination behind +// the retry verdict signal protocol. +// +// Verdict signal lifecycle: +// - streamWithShouldRetry creates a new retryVerdictSignal per retry attempt, stores it in +// execCtx.retryVerdictSignal, and sends exactly one retryVerdict after ShouldRetry decides. +// - The closures below capture a *retryVerdictSignal that is nil at closure-creation time; they +// read the live value from execCtx.retryVerdictSignal, which is set before each model call. +// +// Two hooks cooperate to cover all stream termination paths: +// - WithErrWrapper intercepts mid-stream errors. It blocks on the verdict to decide +// whether to wrap the error as WillRetryError (rejected attempt) or pass it through (accepted). +// - WithOnEOF intercepts clean EOF (successful stream). It blocks on the verdict to +// either inject a WillRetryError (rejected) or pass through io.EOF (accepted). +// +// Both hooks share a sync.Once-guarded reader so the verdict channel is read at most once. +// This prevents a goroutine leak when a mid-stream error is followed by EOF: errWrapper fires +// first (caching the verdict), and onEOF reuses the cached value instead of blocking on a +// drained channel. +func (m *typedEventSenderModel[M]) buildStreamConvertOptions(ctx context.Context) []schema.ConvertOption { + var retryAttempt int + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + retryAttempt = st.getRetryAttempt() + return nil + }) + + wrapWithCancelGuard := func(inner func(error) error) func(error) error { + return func(err error) error { + if errors.Is(err, ErrStreamCanceled) { + return err + } + return inner(err) + } + } + + var opts []schema.ConvertOption + + var retryWrapper func(error) error + if m.modelRetryConfig != nil { + if m.modelRetryConfig.ShouldRetry != nil { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + signal := (*retryVerdictSignal)(nil) + if execCtx != nil { + signal = execCtx.retryVerdictSignal + } + if signal != nil { + var ( + verdictOnce sync.Once + cachedVerdict retryVerdict + ) + readVerdict := func() retryVerdict { + verdictOnce.Do(func() { + cachedVerdict = <-signal.ch + }) + return cachedVerdict + } + + retryWrapper = wrapWithCancelGuard(func(err error) error { + verdict := readVerdict() + if verdict.WillRetry { + return &WillRetryError{ + ErrStr: err.Error(), + RetryAttempt: verdict.RetryAttempt, + rejectReason: verdict.RejectReason, + err: err, + } + } + return err + }) + + opts = append(opts, schema.WithOnEOF(func() (any, error) { + verdict := readVerdict() + if verdict.WillRetry { + return nil, &WillRetryError{ + ErrStr: verdict.Err.Error(), + RetryAttempt: verdict.RetryAttempt, + rejectReason: verdict.RejectReason, + err: verdict.Err, + } + } + return nil, io.EOF + })) + } + } else { + retryWrapper = wrapWithCancelGuard( + genErrWrapper(ctx, m.modelRetryConfig.MaxRetries, retryAttempt, m.modelRetryConfig.IsRetryAble), + ) + } + } + + hasFailover := m.modelFailoverConfig != nil + // failoverHasMoreAttempts is set by failoverModelWrapper before each inner call. + // It is true when additional failover attempts remain after the current one, + // meaning stream errors should be wrapped as WillRetryError so the flow layer + // skips them. On the final attempt it is false, so the error propagates normally. + failoverHasMore := getFailoverHasMoreAttempts(ctx) + + if retryWrapper == nil && !(hasFailover && failoverHasMore) { + return opts + } + + combinedErrWrapper := func(err error) error { + // If retry is configured and will retry this error, use the retry wrapper's WillRetryError. + if retryWrapper != nil { + wrapped := retryWrapper(err) + if errors.As(wrapped, new(*WillRetryError)) { + return wrapped + } + } + // Retry won't handle this error (either exhausted or not configured), but + // failover still has more attempts remaining. Wrap it as WillRetryError so + // the flow layer skips this event from the failed attempt. + if hasFailover && failoverHasMore { + if errors.Is(err, ErrStreamCanceled) { + return err + } + return &WillRetryError{ErrStr: err.Error(), err: err} + } + return err + } + opts = append(opts, schema.WithErrWrapper(combinedErrWrapper)) + + return opts +} + +func copyMessage[M MessageType](msg M) M { + switch v := any(msg).(type) { + case *schema.Message: + cp := *v + return any(&cp).(M) + case *schema.AgenticMessage: + cp := *v + return any(&cp).(M) + default: + return msg + } +} + +// typedSetMessageID sets a specific message ID in Extra. +func typedSetMessageID[M MessageType](msg M, id string) { + switch v := any(msg).(type) { + case *schema.Message: + v.Extra = internal.SetMessageID(v.Extra, id) + case *schema.AgenticMessage: + v.Extra = internal.SetMessageID(v.Extra, id) + } +} + +// GetMessageID returns the eino-internal message ID from the given message, or "". +func GetMessageID[M MessageType](msg M) string { + switch v := any(msg).(type) { + case *schema.Message: + return internal.GetMessageID(v.Extra) + case *schema.AgenticMessage: + return internal.GetMessageID(v.Extra) + default: + return "" + } +} + +// EnsureMessageID assigns a UUID v4 message ID if the message doesn't have one. +// Idempotent: if ID already set, no-op. +// Middleware authors should call this before SendEvent if they create messages. +func EnsureMessageID[M MessageType](msg M) { + switch v := any(msg).(type) { + case *schema.Message: + v.Extra = internal.EnsureMessageID(v.Extra) + case *schema.AgenticMessage: + v.Extra = internal.EnsureMessageID(v.Extra) + } +} + +func typedPopToolGenAction[M MessageType](ctx context.Context, toolName string) *AgentAction { + toolCallID := compose.GetToolCallID(ctx) + + var action *AgentAction + _ = compose.ProcessState(ctx, func(ctx context.Context, st *typedState[M]) error { + if len(toolCallID) > 0 { + if a := st.popToolGenAction(toolCallID); a != nil { + action = a + return nil + } + } + + if a := st.popToolGenAction(toolName); a != nil { + action = a + } + + return nil + }) + + return action +} + +type typedEventSenderToolWrapper[M MessageType] struct { + *TypedBaseChatModelAgentMiddleware[M] +} + +func (*typedEventSenderToolWrapper[M]) isEventSenderToolWrapper() {} + +// eventSenderToolWrapperMarker enables cross-type detection of eventSenderToolWrapper +// in generic contexts. hasUserEventSenderToolWrapper[M] receives +// []TypedChatModelAgentMiddleware[M], so when M is *schema.AgenticMessage, a direct +// type assertion to *eventSenderToolWrapper (which implements the *schema.Message alias) +// would fail. The marker interface bridges this gap. +type eventSenderToolWrapperMarker interface{ isEventSenderToolWrapper() } + +// NewEventSenderToolWrapper returns a ChatModelAgentMiddleware that sends tool result events. +// By default, the framework places this before all user middlewares (outermost), so events +// reflect the fully processed tool output. To control exactly where events are emitted, +// include this in ChatModelAgentConfig.Handlers at the desired position. +// When detected in Handlers, the framework skips the default event sender to avoid duplicates. +func NewEventSenderToolWrapper() ChatModelAgentMiddleware { + return newTypedEventSenderToolWrapper[*schema.Message]() +} + +// newTypedEventSenderToolWrapper creates a typed event sender wrapper for the given message type. +// This is used internally to ensure the default event sender matches the agent's message type +// (e.g. *schema.AgenticMessage agents need an AgenticMessage-typed wrapper so that +// compose.ProcessState can access the correct state type). +func newTypedEventSenderToolWrapper[M MessageType]() *typedEventSenderToolWrapper[M] { + return &typedEventSenderToolWrapper[M]{ + TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[M]{}, + } +} + +// textToFunctionToolResultBlocks wraps a plain text string into FunctionToolResultBlocks. +func textToFunctionToolResultBlocks(text string) []*schema.FunctionToolResultContentBlock { + if text == "" { + return nil + } + return []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: text}}, + } +} + +// functionToolResultAgenticMessage constructs a function tool result message with AgenticRoleType "user". +func functionToolResultAgenticMessage(callID, name string, content []*schema.FunctionToolResultContentBlock) *schema.AgenticMessage { + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.FunctionToolResult{ + CallID: callID, + Name: name, + Content: content, + }), + }, + } +} + +func markAgenticMessageStreamingMeta(msg *schema.AgenticMessage, index int) { + if msg == nil { + return + } + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + block.StreamingMeta = &schema.StreamingMeta{Index: index} + } +} + +func toolSearchResultAgenticMessage(callID, name string, tr *schema.ToolResult) (*schema.AgenticMessage, bool) { + if tr == nil || len(tr.Parts) != 1 { + return nil, false + } + + part := tr.Parts[0] + if part.Type != schema.ToolPartTypeToolSearchResult || part.ToolSearchResult == nil { + return nil, false + } + + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.ToolSearchFunctionToolResult{ + CallID: callID, + Name: name, + Result: part.ToolSearchResult, + }), + }, + }, true +} + +func toolResultAgenticMessage(callID, name string, tr *schema.ToolResult) *schema.AgenticMessage { + if msg, ok := toolSearchResultAgenticMessage(callID, name, tr); ok { + return msg + } + return functionToolResultAgenticMessage(callID, name, toolResultToBlocks(tr)) +} + +// toolResultToBlocks converts a ToolResult's multimodal parts into FunctionToolResultBlocks. +// This preserves all media types (text, image, audio, video, file), unlike toolResultText +// which only extracts text. +func toolResultToBlocks(tr *schema.ToolResult) []*schema.FunctionToolResultContentBlock { + if tr == nil || len(tr.Parts) == 0 { + return nil + } + blocks := make([]*schema.FunctionToolResultContentBlock, 0, len(tr.Parts)) + for _, p := range tr.Parts { + var block *schema.FunctionToolResultContentBlock + switch p.Type { + case schema.ToolPartTypeText: + block = &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeText, + Text: &schema.UserInputText{Text: p.Text}, + Extra: p.Extra, + } + case schema.ToolPartTypeImage: + if p.Image != nil { + block = &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeImage, + Image: &schema.UserInputImage{ + URL: derefString(p.Image.URL), + Base64Data: derefString(p.Image.Base64Data), + MIMEType: p.Image.MIMEType, + }, + Extra: p.Extra, + } + } + case schema.ToolPartTypeAudio: + if p.Audio != nil { + block = &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeAudio, + Audio: &schema.UserInputAudio{ + URL: derefString(p.Audio.URL), + Base64Data: derefString(p.Audio.Base64Data), + MIMEType: p.Audio.MIMEType, + }, + Extra: p.Extra, + } + } + case schema.ToolPartTypeVideo: + if p.Video != nil { + block = &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeVideo, + Video: &schema.UserInputVideo{ + URL: derefString(p.Video.URL), + Base64Data: derefString(p.Video.Base64Data), + MIMEType: p.Video.MIMEType, + }, + Extra: p.Extra, + } + } + case schema.ToolPartTypeFile: + if p.File != nil { + block = &schema.FunctionToolResultContentBlock{ + Type: schema.FunctionToolResultContentBlockTypeFile, + File: &schema.UserInputFile{ + URL: derefString(p.File.URL), + Base64Data: derefString(p.File.Base64Data), + MIMEType: p.File.MIMEType, + }, + Extra: p.Extra, + } + } + } + if block != nil { + blocks = append(blocks, block) + } + } + return blocks +} + +func derefString(s *string) string { + if s == nil { + return "" + } + return *s +} + +// typedToolInvokeEvent constructs the tool result event for the invoke path, +// dispatching on M to create the correct message and event types. +func typedToolInvokeEvent[M MessageType](callID, toolName, result, toolMsgID string) *TypedAgentEvent[M] { + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.ToolMessage(result, callID, schema.WithToolName(toolName)) + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + event := EventFromMessage(msg, nil, schema.Tool, toolName) + return any(event).(*TypedAgentEvent[M]) + case *schema.AgenticMessage: + msg := functionToolResultAgenticMessage(callID, toolName, textToFunctionToolResultBlocks(result)) + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + event := EventFromAgenticMessage(msg, nil, schema.AgenticRoleTypeUser) + return any(event).(*TypedAgentEvent[M]) + default: + return nil + } +} + +// typedToolStreamEvent constructs the tool result event for the stream path, +// dispatching on M to create the correct message stream and event types. +func typedToolStreamEvent[M MessageType](callID, toolName, toolMsgID string, stream *schema.StreamReader[string]) *TypedAgentEvent[M] { + var zero M + switch any(zero).(type) { + case *schema.Message: + first := true + cvt := func(in string) (Message, error) { + msg := schema.ToolMessage(in, callID, schema.WithToolName(toolName)) + if first { + first = false + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + } + return msg, nil + } + msgStream := schema.StreamReaderWithConvert(stream, cvt) + event := EventFromMessage(nil, msgStream, schema.Tool, toolName) + return any(event).(*TypedAgentEvent[M]) + case *schema.AgenticMessage: + first := true + cvt := func(in string) (*schema.AgenticMessage, error) { + msg := functionToolResultAgenticMessage(callID, toolName, textToFunctionToolResultBlocks(in)) + markAgenticMessageStreamingMeta(msg, 0) + if first { + first = false + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + } + return msg, nil + } + msgStream := schema.StreamReaderWithConvert(stream, cvt) + event := EventFromAgenticMessage(nil, msgStream, schema.AgenticRoleTypeUser) + return any(event).(*TypedAgentEvent[M]) + default: + return nil + } +} + +// typedToolEnhancedInvokeEvent constructs the tool result event for the enhanced invoke path. +// For *schema.Message it builds a multimodal tool message; for *schema.AgenticMessage it +// uses the string content of the result (AgenticToolsNode only uses the string path). +func typedToolEnhancedInvokeEvent[M MessageType](callID, toolName, toolMsgID string, result *schema.ToolResult) (*TypedAgentEvent[M], error) { + var zero M + switch any(zero).(type) { + case *schema.Message: + msg := schema.ToolMessage("", callID, schema.WithToolName(toolName)) + var err error + msg.UserInputMultiContent, err = result.ToMessageInputParts() + if err != nil { + return nil, err + } + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + event := EventFromMessage(msg, nil, schema.Tool, toolName) + return any(event).(*TypedAgentEvent[M]), nil + case *schema.AgenticMessage: + msg := toolResultAgenticMessage(callID, toolName, result) + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + event := EventFromAgenticMessage(msg, nil, schema.AgenticRoleTypeUser) + return any(event).(*TypedAgentEvent[M]), nil + default: + return nil, nil + } +} + +// typedToolEnhancedStreamEvent constructs the tool result event for the enhanced stream path. +// For *schema.Message it builds multimodal tool messages; for *schema.AgenticMessage it +// converts each chunk's multimodal parts into FunctionToolResultBlocks. +func typedToolEnhancedStreamEvent[M MessageType](callID, toolName, toolMsgID string, stream *schema.StreamReader[*schema.ToolResult]) *TypedAgentEvent[M] { + var zero M + switch any(zero).(type) { + case *schema.Message: + first := true + cvt := func(in *schema.ToolResult) (Message, error) { + msg := schema.ToolMessage("", callID, schema.WithToolName(toolName)) + var cvtErr error + msg.UserInputMultiContent, cvtErr = in.ToMessageInputParts() + if cvtErr != nil { + return nil, cvtErr + } + if first { + first = false + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + } + return msg, nil + } + msgStream := schema.StreamReaderWithConvert(stream, cvt) + event := EventFromMessage(nil, msgStream, schema.Tool, toolName) + return any(event).(*TypedAgentEvent[M]) + case *schema.AgenticMessage: + first := true + cvt := func(in *schema.ToolResult) (*schema.AgenticMessage, error) { + msg := toolResultAgenticMessage(callID, toolName, in) + markAgenticMessageStreamingMeta(msg, 0) + if first { + first = false + msg.Extra = internal.SetMessageID(msg.Extra, toolMsgID) + } + return msg, nil + } + msgStream := schema.StreamReaderWithConvert(stream, cvt) + event := EventFromAgenticMessage(nil, msgStream, schema.AgenticRoleTypeUser) + return any(event).(*TypedAgentEvent[M]) + default: + return nil + } +} + +func (w *typedEventSenderToolWrapper[M]) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, tCtx *ToolContext) (InvokableToolCallEndpoint, error) { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + result, err := endpoint(ctx, argumentsInJSON, opts...) + if err != nil { + return "", err + } + + toolName := tCtx.Name + callID := tCtx.CallID + + prePopAction := typedPopToolGenAction[M](ctx, toolName) + toolMsgID := uuid.NewString() + event := typedToolInvokeEvent[M](callID, toolName, result, toolMsgID) + if prePopAction != nil { + event.Action = prePopAction + } + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setToolMsgID(toolName, callID, toolMsgID) + if st.getReturnDirectlyToolCallID() == callID { + st.setReturnDirectlyEvent(event) + } else { + execCtx.send(event) + } + return nil + }) + + return result, nil + }, nil +} + +func (w *typedEventSenderToolWrapper[M]) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, tCtx *ToolContext) (StreamableToolCallEndpoint, error) { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + result, err := endpoint(ctx, argumentsInJSON, opts...) + if err != nil { + return nil, err + } + + toolName := tCtx.Name + callID := tCtx.CallID + + prePopAction := typedPopToolGenAction[M](ctx, toolName) + streams := result.Copy(2) + + toolMsgID := uuid.NewString() + event := typedToolStreamEvent[M](callID, toolName, toolMsgID, streams[0]) + event.Action = prePopAction + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setToolMsgID(toolName, callID, toolMsgID) + if st.getReturnDirectlyToolCallID() == callID { + st.setReturnDirectlyEvent(event) + } else { + execCtx.send(event) + } + return nil + }) + + return streams[1], nil + }, nil +} + +func (w *typedEventSenderToolWrapper[M]) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error) { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + result, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + + toolName := tCtx.Name + callID := tCtx.CallID + + prePopAction := typedPopToolGenAction[M](ctx, toolName) + toolMsgID := uuid.NewString() + event, eventErr := typedToolEnhancedInvokeEvent[M](callID, toolName, toolMsgID, result) + if eventErr != nil { + return nil, eventErr + } + if prePopAction != nil { + event.Action = prePopAction + } + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setToolMsgID(toolName, callID, toolMsgID) + if st.getReturnDirectlyToolCallID() == callID { + st.setReturnDirectlyEvent(event) + } else { + execCtx.send(event) + } + return nil + }) + + return result, nil + }, nil +} + +func (w *typedEventSenderToolWrapper[M]) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error) { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + result, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + + toolName := tCtx.Name + callID := tCtx.CallID + + prePopAction := typedPopToolGenAction[M](ctx, toolName) + streams := result.Copy(2) + + toolMsgID := uuid.NewString() + event := typedToolEnhancedStreamEvent[M](callID, toolName, toolMsgID, streams[0]) + event.Action = prePopAction + + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.setToolMsgID(toolName, callID, toolMsgID) + if st.getReturnDirectlyToolCallID() == callID { + st.setReturnDirectlyEvent(event) + } else { + execCtx.send(event) + } + return nil + }) + + return streams[1], nil + }, nil +} + +func hasUserEventSenderToolWrapper[M MessageType](handlers []TypedChatModelAgentMiddleware[M]) bool { + for _, handler := range handlers { + if _, ok := any(handler).(eventSenderToolWrapperMarker); ok { + return true + } + } + return false +} + +type typedStateModelWrapper[M MessageType] struct { + inner model.BaseModel[M] + original model.BaseModel[M] + handlers []TypedChatModelAgentMiddleware[M] + middlewares []AgentMiddleware + toolInfos []*schema.ToolInfo + modelRetryConfig *TypedModelRetryConfig[M] + modelFailoverConfig *ModelFailoverConfig[M] + cancelContext *cancelContext +} + +type stateModelWrapper = typedStateModelWrapper[*schema.Message] + +func (w *typedStateModelWrapper[M]) IsCallbacksEnabled() bool { + return true +} + +func (w *typedStateModelWrapper[M]) GetType() string { + if typer, ok := any(w.original).(components.Typer); ok { + return typer.GetType() + } + return generic.ParseTypeName(reflect.ValueOf(w.original)) +} + +func (w *typedStateModelWrapper[M]) hasUserEventSender() bool { + for _, handler := range w.handlers { + if _, ok := any(handler).(*typedEventSenderModelWrapper[M]); ok { + return true + } + } + return false +} + +func (w *typedStateModelWrapper[M]) wrapGenerateEndpoint(endpoint typedGenerateEndpoint[M]) typedGenerateEndpoint[M] { + // === ID Assignment layer (innermost, framework-controlled) === + // Ensures model output has a message ID before any WrapModel handler or event sender sees it. + // Copies the result to avoid mutating a potentially shared pointer returned by the model. + { + realInner := endpoint + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) { + result, err := realInner(ctx, input, opts...) + if err != nil { + return result, err + } + if GetMessageID(result) == "" { + result = copyMessage(result) + EnsureMessageID(result) + } + return result, nil + } + } + + hasUserEventSender := w.hasUserEventSender() + retryConfig := w.modelRetryConfig + failoverConfig := w.modelFailoverConfig + cc := w.cancelContext + + for i := len(w.handlers) - 1; i >= 0; i-- { + handler := w.handlers[i] + innerEndpoint := endpoint + baseToolInfos := w.toolInfos + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) { + baseOpts := &model.Options{Tools: baseToolInfos} + commonOpts := model.GetCommonOptions(baseOpts, opts...) + mc := &TypedModelContext[M]{Tools: commonOpts.Tools, ModelRetryConfig: retryConfig, cancelContext: cc} + wrappedModel, err := handler.WrapModel(ctx, &typedEndpointModel[M]{generate: innerEndpoint}, mc) + if err != nil { + var zero M + return zero, err + } + return wrappedModel.Generate(ctx, input, opts...) + } + } + + if !hasUserEventSender { + innerEndpoint := endpoint + eventSender := &typedEventSenderModelWrapper[M]{ + TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[M]{}, + } + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil || execCtx.generator == nil { + return innerEndpoint(ctx, input, opts...) + } + mc := &TypedModelContext[M]{ModelRetryConfig: retryConfig, ModelFailoverConfig: failoverConfig, cancelContext: cc} + wrappedModel, err := eventSender.WrapModel(ctx, &typedEndpointModel[M]{generate: innerEndpoint}, mc) + if err != nil { + var zero M + return zero, err + } + return wrappedModel.Generate(ctx, input, opts...) + } + } + + if w.modelRetryConfig != nil { + innerEndpoint := endpoint + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) { + retryWrapper := newTypedRetryModelWrapper[M](&typedEndpointModel[M]{generate: innerEndpoint}, w.modelRetryConfig) + return retryWrapper.Generate(ctx, input, opts...) + } + } + + if w.modelFailoverConfig != nil { + config := w.modelFailoverConfig + innerEndpoint := endpoint + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (M, error) { + failoverWrapper := newFailoverModelWrapper[M](&typedEndpointModel[M]{generate: innerEndpoint}, config) + return failoverWrapper.Generate(ctx, input, opts...) + } + } + + return endpoint +} + +func (w *typedStateModelWrapper[M]) wrapStreamEndpoint(endpoint typedStreamEndpoint[M]) typedStreamEndpoint[M] { + // === ID Assignment layer (innermost, framework-controlled) === + // Pre-allocates a UUID and injects it into the first chunk only. + // Only the first chunk carries the ID in Extra to avoid concatStrings corruption + // during ConcatMessages (which string-concatenates duplicate Extra keys). + { + realInner := endpoint + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + reader, err := realInner(ctx, input, opts...) + if err != nil { + return nil, err + } + msgID := uuid.NewString() + first := true + return schema.StreamReaderWithConvert(reader, func(msg M) (M, error) { + if first { + first = false + if GetMessageID(msg) == "" { + typedSetMessageID(msg, msgID) + } + } + return msg, nil + }), nil + } + } + + hasUserEventSender := w.hasUserEventSender() + retryConfig := w.modelRetryConfig + failoverConfig := w.modelFailoverConfig + cc := w.cancelContext + + for i := len(w.handlers) - 1; i >= 0; i-- { + handler := w.handlers[i] + innerEndpoint := endpoint + baseToolInfos := w.toolInfos + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + baseOpts := &model.Options{Tools: baseToolInfos} + commonOpts := model.GetCommonOptions(baseOpts, opts...) + mc := &TypedModelContext[M]{Tools: commonOpts.Tools, ModelRetryConfig: retryConfig, cancelContext: cc} + wrappedModel, err := handler.WrapModel(ctx, &typedEndpointModel[M]{stream: innerEndpoint}, mc) + if err != nil { + return nil, err + } + return wrappedModel.Stream(ctx, input, opts...) + } + } + + if !hasUserEventSender { + innerEndpoint := endpoint + eventSender := &typedEventSenderModelWrapper[M]{ + TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[M]{}, + } + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + execCtx := getTypedChatModelAgentExecCtx[M](ctx) + if execCtx == nil || execCtx.generator == nil { + return innerEndpoint(ctx, input, opts...) + } + mc := &TypedModelContext[M]{ModelRetryConfig: retryConfig, ModelFailoverConfig: failoverConfig, cancelContext: cc} + wrappedModel, err := eventSender.WrapModel(ctx, &typedEndpointModel[M]{stream: innerEndpoint}, mc) + if err != nil { + return nil, err + } + return wrappedModel.Stream(ctx, input, opts...) + } + } + + if w.modelRetryConfig != nil { + innerEndpoint := endpoint + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + retryWrapper := newTypedRetryModelWrapper[M](&typedEndpointModel[M]{stream: innerEndpoint}, w.modelRetryConfig) + return retryWrapper.Stream(ctx, input, opts...) + } + } + + if w.modelFailoverConfig != nil { + config := w.modelFailoverConfig + innerEndpoint := endpoint + endpoint = func(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + failoverWrapper := newFailoverModelWrapper[M](&typedEndpointModel[M]{stream: innerEndpoint}, config) + return failoverWrapper.Stream(ctx, input, opts...) + } + } + + return endpoint +} + +func (w *typedStateModelWrapper[M]) Generate(ctx context.Context, _ []M, opts ...model.Option) (M, error) { + var ( + stateMessages []M + stateToolInfos []*schema.ToolInfo + stateDeferredToolInfos []*schema.ToolInfo + ) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + stateMessages = st.Messages + stateToolInfos = st.ToolInfos + stateDeferredToolInfos = st.DeferredToolInfos + return nil + }) + + // Backfill: old checkpoints or fresh starts have nil ToolInfos. + // Use compose-level tools from opts (which always reflects the latest bc.toolInfos) + // rather than w.toolInfos (which may be stale if the graph was reused). + if stateToolInfos == nil { + composeLevelOpts := model.GetCommonOptions(&model.Options{}, opts...) + if composeLevelOpts.Tools != nil { + stateToolInfos = composeLevelOpts.Tools + } else { + stateToolInfos = w.toolInfos + } + } + + state := &TypedChatModelAgentState[M]{ + Messages: stateMessages, + ToolInfos: stateToolInfos, + DeferredToolInfos: stateDeferredToolInfos, + } + + if msgState, ok := any(state).(*ChatModelAgentState); ok { + for _, m := range w.middlewares { + if m.BeforeChatModel != nil { + if err := m.BeforeChatModel(ctx, msgState); err != nil { + var zero M + return zero, err + } + } + } + } + + baseOpts := &model.Options{Tools: w.toolInfos} + commonOpts := model.GetCommonOptions(baseOpts, opts...) + mc := &TypedModelContext[M]{Tools: commonOpts.Tools, ModelRetryConfig: w.modelRetryConfig, cancelContext: w.cancelContext} + for _, handler := range w.handlers { + var err error + ctx, state, err = handler.BeforeModelRewriteState(ctx, state, mc) + if err != nil { + var zero M + return zero, err + } + } + + // Persist state (including tool infos) after BeforeModelRewriteState. + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.Messages = state.Messages + st.ToolInfos = state.ToolInfos + st.DeferredToolInfos = state.DeferredToolInfos + return nil + }) + + // Derive model options from state. Append after caller opts so state takes precedence + // (model.GetCommonOptions applies left-to-right, last wins). + // Use explicit copy to avoid mutating the caller's opts slice. + derivedOpts := make([]model.Option, len(opts), len(opts)+2) + copy(derivedOpts, opts) + derivedOpts = append(derivedOpts, model.WithTools(state.ToolInfos)) + if state.DeferredToolInfos != nil { + derivedOpts = append(derivedOpts, model.WithDeferredTools(state.DeferredToolInfos)) + } + + wrappedEndpoint := w.wrapGenerateEndpoint(w.inner.Generate) + result, err := wrappedEndpoint(ctx, state.Messages, derivedOpts...) + if err != nil { + var zero M + return zero, err + } + + // Re-read State.Messages after Generate completes: when ShouldRetry uses + // PersistModifiedInputMessages, applyDecisionForRetry writes modified messages to State. + // We must pick up those changes before appending the model result. + if w.modelRetryConfig != nil && w.modelRetryConfig.ShouldRetry != nil { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + state.Messages = st.Messages + return nil + }) + } + + state.Messages = append(state.Messages, result) + + for _, handler := range w.handlers { + ctx, state, err = handler.AfterModelRewriteState(ctx, state, mc) + if err != nil { + var zero M + return zero, err + } + } + + if msgState, ok := any(state).(*ChatModelAgentState); ok { + for _, m := range w.middlewares { + if m.AfterChatModel != nil { + if err := m.AfterChatModel(ctx, msgState); err != nil { + var zero M + return zero, err + } + } + } + } + + // Persist state (including tool infos) after AfterModelRewriteState. + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.Messages = state.Messages + st.ToolInfos = state.ToolInfos + st.DeferredToolInfos = state.DeferredToolInfos + return nil + }) + + if len(state.Messages) == 0 { + var zero M + return zero, errors.New("no messages left in state after model call") + } + return state.Messages[len(state.Messages)-1], nil +} + +func (w *typedStateModelWrapper[M]) Stream(ctx context.Context, _ []M, opts ...model.Option) (*schema.StreamReader[M], error) { + var ( + stateMessages []M + stateToolInfos []*schema.ToolInfo + stateDeferredToolInfos []*schema.ToolInfo + ) + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + stateMessages = st.Messages + stateToolInfos = st.ToolInfos + stateDeferredToolInfos = st.DeferredToolInfos + return nil + }) + + // Backfill: old checkpoints or fresh starts have nil ToolInfos. + // Use compose-level tools from opts (which always reflects the latest bc.toolInfos) + // rather than w.toolInfos (which may be stale if the graph was reused). + if stateToolInfos == nil { + composeLevelOpts := model.GetCommonOptions(&model.Options{}, opts...) + if composeLevelOpts.Tools != nil { + stateToolInfos = composeLevelOpts.Tools + } else { + stateToolInfos = w.toolInfos + } + } + + state := &TypedChatModelAgentState[M]{ + Messages: stateMessages, + ToolInfos: stateToolInfos, + DeferredToolInfos: stateDeferredToolInfos, + } + + if msgState, ok := any(state).(*ChatModelAgentState); ok { + for _, m := range w.middlewares { + if m.BeforeChatModel != nil { + if err := m.BeforeChatModel(ctx, msgState); err != nil { + return nil, err + } + } + } + } + + baseOpts := &model.Options{Tools: w.toolInfos} + commonOpts := model.GetCommonOptions(baseOpts, opts...) + mc := &TypedModelContext[M]{Tools: commonOpts.Tools, ModelRetryConfig: w.modelRetryConfig, cancelContext: w.cancelContext} + for _, handler := range w.handlers { + var err error + ctx, state, err = handler.BeforeModelRewriteState(ctx, state, mc) + if err != nil { + return nil, err + } + } + + // Persist state (including tool infos) after BeforeModelRewriteState. + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.Messages = state.Messages + st.ToolInfos = state.ToolInfos + st.DeferredToolInfos = state.DeferredToolInfos + return nil + }) + + // Derive model options from state. Append after caller opts so state takes precedence + // (model.GetCommonOptions applies left-to-right, last wins). + // Use explicit copy to avoid mutating the caller's opts slice. + derivedOpts := make([]model.Option, len(opts), len(opts)+2) + copy(derivedOpts, opts) + derivedOpts = append(derivedOpts, model.WithTools(state.ToolInfos)) + if state.DeferredToolInfos != nil { + derivedOpts = append(derivedOpts, model.WithDeferredTools(state.DeferredToolInfos)) + } + + wrappedEndpoint := w.wrapStreamEndpoint(w.inner.Stream) + stream, err := wrappedEndpoint(ctx, state.Messages, derivedOpts...) + if err != nil { + return nil, err + } + result, err := concatMessageStream(stream) + if err != nil { + return nil, err + } + + // Re-read State.Messages after Stream completes: same rationale as in Generate above. + if w.modelRetryConfig != nil && w.modelRetryConfig.ShouldRetry != nil { + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + state.Messages = st.Messages + return nil + }) + } + + state.Messages = append(state.Messages, result) + + for _, handler := range w.handlers { + ctx, state, err = handler.AfterModelRewriteState(ctx, state, mc) + if err != nil { + return nil, err + } + } + + if msgState, ok := any(state).(*ChatModelAgentState); ok { + for _, m := range w.middlewares { + if m.AfterChatModel != nil { + if err := m.AfterChatModel(ctx, msgState); err != nil { + return nil, err + } + } + } + } + + // Persist state (including tool infos) after AfterModelRewriteState. + _ = compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error { + st.Messages = state.Messages + st.ToolInfos = state.ToolInfos + st.DeferredToolInfos = state.DeferredToolInfos + return nil + }) + + if len(state.Messages) == 0 { + return nil, errors.New("no messages left in state after model call") + } + return schema.StreamReaderFromArray([]M{state.Messages[len(state.Messages)-1]}), nil +} + +type typedEndpointModel[M MessageType] struct { + generate typedGenerateEndpoint[M] + stream typedStreamEndpoint[M] +} + +func (m *typedEndpointModel[M]) Generate(ctx context.Context, input []M, opts ...model.Option) (M, error) { + if m.generate != nil { + return m.generate(ctx, input, opts...) + } + var zero M + return zero, errors.New("generate endpoint not set") +} + +func (m *typedEndpointModel[M]) Stream(ctx context.Context, input []M, opts ...model.Option) (*schema.StreamReader[M], error) { + if m.stream != nil { + return m.stream(ctx, input, opts...) + } + return nil, errors.New("stream endpoint not set") +} diff --git a/adk/wrappers_failover_test.go b/adk/wrappers_failover_test.go new file mode 100644 index 0000000..45fb0c2 --- /dev/null +++ b/adk/wrappers_failover_test.go @@ -0,0 +1,215 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +func TestBuildModelWrappers_FailoverProxyInner(t *testing.T) { + base := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil + }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 0, + ShouldFailover: func(context.Context, *schema.Message, error) bool { return false }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return base, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](base, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + smw, ok := wrapped.(*stateModelWrapper) + require.True(t, ok) + _, ok = smw.inner.(*failoverProxyModel) + require.True(t, ok) + require.Same(t, base, smw.original) + require.Same(t, failoverCfg, smw.modelFailoverConfig) +} + +func TestStateModelWrapper_Generate_WithFailover(t *testing.T) { + wantErr := errors.New("first failed") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return schema.AssistantMessage("partial", nil), wantErr + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok", nil), nil + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + require.ErrorIs(t, err, wantErr) + require.NotNil(t, out) + require.Equal(t, "partial", out.Content) + return true + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + got, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.NotNil(t, got) + require.Equal(t, "ok", got.Content) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) +} + +func TestStateModelWrapper_Stream_WithFailover(t *testing.T) { + streamErr := errors.New("mid error") + var shouldCalls int32 + var m1Calls int32 + var m2Calls int32 + + m1 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p1", nil), + schema.AssistantMessage("p2", nil), + }, streamErr), nil + }, + } + m2 := &fakeChatModel{ + callbacksEnabled: true, + generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + }, + stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final", nil)}), nil + }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool { + atomic.AddInt32(&shouldCalls, 1) + require.ErrorIs(t, err, streamErr) + require.NotNil(t, out) + require.Equal(t, "p1p2", out.Content) + return true + }, + GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.Equal(t, uint(1), failoverCtx.FailoverAttempt) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "final", msgs[0].Content) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls)) +} + +func TestFailoverAcceptsAgenticAgent(t *testing.T) { + ctx := context.Background() + + m := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("ok"), nil + }, + } + + fallbackModel := &mockAgenticModel{ + generateFn: func(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.AgenticMessage, error) { + return agenticMsg("fallback"), nil + }, + } + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "FailoverAgent", + Description: "Agent with failover config", + Model: m, + ModelFailoverConfig: &ModelFailoverConfig[*schema.AgenticMessage]{ + MaxRetries: 1, + ShouldFailover: func(ctx context.Context, outputMessage *schema.AgenticMessage, outputErr error) bool { + return true + }, + GetFailoverModel: func(ctx context.Context, failoverCtx *FailoverContext[*schema.AgenticMessage]) (model.BaseModel[*schema.AgenticMessage], []*schema.AgenticMessage, error) { + return fallbackModel, nil, nil + }, + }, + }) + require.NoError(t, err) + assert.NotNil(t, agent) +} diff --git a/adk/wrappers_retry_failover_test.go b/adk/wrappers_retry_failover_test.go new file mode 100644 index 0000000..c1a291d --- /dev/null +++ b/adk/wrappers_retry_failover_test.go @@ -0,0 +1,613 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +func newFakeChatModel( + gen func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error), + stream func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error), +) *fakeChatModel { + if gen == nil { + gen = func(context.Context, []*schema.Message, ...model.Option) (*schema.Message, error) { + return nil, errors.New("unused") + } + } + if stream == nil { + stream = func(context.Context, []*schema.Message, ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return nil, errors.New("unused") + } + } + return &fakeChatModel{callbacksEnabled: true, generate: gen, stream: stream} +} + +func TestRetryThenFailover(t *testing.T) { + t.Run("Generate_RetryExhaustedTriggersFailover", func(t *testing.T) { + modelErr := errors.New("model error") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, modelErr + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok from m2", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 2, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, fc *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.NotNil(t, fc.LastErr) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok from m2", msg.Content) + + // m1: 1 (lastSuccess) + 2 retries = 3 calls on lastSuccess attempt, + // then failover to m2 which also goes through retry wrapper: 1 call succeeds. + require.Equal(t, int32(3), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Generate_AllExhausted", func(t *testing.T) { + modelErr := errors.New("always fails") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, modelErr + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return nil, modelErr + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + + // Should be RetryExhaustedError from m2's retry wrapper + var retryErr *RetryExhaustedError + require.True(t, errors.As(err, &retryErr)) + + // m1: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + // m2: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Generate_RetrySucceedsNoFailover", func(t *testing.T) { + var m1Calls int32 + var failoverCalled int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + n := atomic.AddInt32(&m1Calls, 1) + if n == 1 { + return nil, errors.New("transient error") + } + return schema.AssistantMessage("ok on retry", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 2, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&failoverCalled, 1) + return true + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + t.Fatal("GetFailoverModel should not be called when retry succeeds") + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok on retry", msg.Content) + + // 2 calls: first fails, second succeeds via retry + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + // ShouldFailover should never be called + require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) + }) + + t.Run("Generate_NonRetryableErrorTriggersFailover", func(t *testing.T) { + nonRetryableErr := errors.New("non-retryable") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, nonRetryableErr + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("ok from m2", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 3, + IsRetryAble: func(_ context.Context, err error) bool { + // Only non-retryable errors + return !errors.Is(err, nonRetryableErr) + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "ok from m2", msg.Content) + + // m1 called only once — non-retryable error skips retry + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Stream_RetryExhaustedTriggersFailover", func(t *testing.T) { + streamErr := errors.New("stream mid error") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("partial", nil), + }, streamErr), nil + }) + m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok from m2", nil)}), nil + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, fc *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + require.NotNil(t, fc.LastErr) + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "ok from m2", msgs[0].Content) + + // m1: 1 initial + 1 retry = 2 calls on lastSuccess attempt + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Stream_AllExhausted", func(t *testing.T) { + streamErr := errors.New("always fails mid-stream") + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p", nil), + }, streamErr), nil + }) + m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("p", nil), + }, streamErr), nil + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + IsRetryAble: func(_ context.Context, err error) bool { return true }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + + var retryErr *RetryExhaustedError + require.True(t, errors.As(err, &retryErr)) + + // m1: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + // m2: 1 initial + 1 retry = 2 calls + require.Equal(t, int32(2), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("ShouldRetry_Stream_TriggersFailover", func(t *testing.T) { + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("bad from m1", nil)}), nil + }) + m2 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m2Calls, 1) + return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("good from m2", nil)}), nil + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad from m1" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + sr, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + msgs, err := drainMessageStream(sr) + require.NoError(t, err) + require.Len(t, msgs, 1) + require.Equal(t, "good from m2", msgs[0].Content) + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("ShouldRetry_Generate_TriggersFailover", func(t *testing.T) { + var m1Calls int32 + var m2Calls int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return schema.AssistantMessage("bad from m1", nil), nil + }, nil) + m2 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m2Calls, 1) + return schema.AssistantMessage("good from m2", nil), nil + }, nil) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 1, + ShouldRetry: func(_ context.Context, retryCtx *RetryContext) *RetryDecision { + if retryCtx.OutputMessage != nil && retryCtx.OutputMessage.Content == "bad from m1" { + return &RetryDecision{Retry: true} + } + return &RetryDecision{Retry: false} + }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return m2, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + msg, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.NoError(t, err) + require.Equal(t, "good from m2", msg.Content) + require.Equal(t, int32(2), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls)) + }) + + t.Run("Stream_GetFailoverModelReturnsNilModel", func(t *testing.T) { + streamErr := errors.New("m1 always fails") + var m1Calls int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return nil, streamErr + }) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 0, + IsRetryAble: func(_ context.Context, err error) bool { return false }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 1, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.Contains(t, err.Error(), "returned nil model at attempt") + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + }) + + t.Run("Stream_ContextCanceledDuringFailover", func(t *testing.T) { + streamErr := errors.New("m1 fails") + var m1Calls int32 + var failoverModelCalled int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return nil, streamErr + }) + + ctx, cancel := context.WithCancel(context.Background()) + + retryCfg := &ModelRetryConfig{ + MaxRetries: 0, + IsRetryAble: func(_ context.Context, err error) bool { return false }, + BackoffFunc: func(_ context.Context, _ int) time.Duration { return 0 }, + } + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 3, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + cancel() + return err != nil + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + atomic.AddInt32(&failoverModelCalled, 1) + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + retryConfig: retryCfg, + failoverConfig: failoverCfg, + }) + + ctx = withTypedChatModelAgentExecCtx(ctx, &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&failoverModelCalled)) + }) +} + +func TestErrStreamCanceled_Failover(t *testing.T) { + t.Run("Stream_NeverFailedOver", func(t *testing.T) { + var m1Calls int32 + var failoverCalled int32 + + m1 := newFakeChatModel(nil, func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) { + atomic.AddInt32(&m1Calls, 1) + return streamWithMidError([]*schema.Message{ + schema.AssistantMessage("partial", nil), + }, ErrStreamCanceled), nil + }) + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 2, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&failoverCalled, 1) + return true + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + t.Fatal("GetFailoverModel should not be called for ErrStreamCanceled") + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Stream(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.True(t, errors.Is(err, ErrStreamCanceled)) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) + }) + + t.Run("Generate_NeverFailedOver", func(t *testing.T) { + var m1Calls int32 + var failoverCalled int32 + + m1 := newFakeChatModel(func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + atomic.AddInt32(&m1Calls, 1) + return nil, ErrStreamCanceled + }, nil) + + failoverCfg := &ModelFailoverConfig[*schema.Message]{ + MaxRetries: 2, + ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool { + atomic.AddInt32(&failoverCalled, 1) + return true + }, + GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) { + t.Fatal("GetFailoverModel should not be called for ErrStreamCanceled") + return nil, nil, nil + }, + } + + wrapped := buildModelWrappers[*schema.Message](m1, &modelWrapperConfig{ + failoverConfig: failoverCfg, + }) + + ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{ + failoverLastSuccessModel: m1, + }) + _, err := wrapped.Generate(ctx, []*schema.Message{schema.UserMessage("hi")}) + require.Error(t, err) + require.True(t, errors.Is(err, ErrStreamCanceled)) + require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls)) + require.Equal(t, int32(0), atomic.LoadInt32(&failoverCalled)) + }) +} diff --git a/adk/wrappers_test.go b/adk/wrappers_test.go new file mode 100644 index 0000000..11dacbd --- /dev/null +++ b/adk/wrappers_test.go @@ -0,0 +1,2188 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package adk + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type testEnhancedToolWrapperHandler struct { + *BaseChatModelAgentMiddleware + wrapEnhancedInvokableFn func(context.Context, EnhancedInvokableToolCallEndpoint, *ToolContext) EnhancedInvokableToolCallEndpoint + wrapEnhancedStreamableFn func(context.Context, EnhancedStreamableToolCallEndpoint, *ToolContext) EnhancedStreamableToolCallEndpoint +} + +func (h *testEnhancedToolWrapperHandler) WrapEnhancedInvokableToolCall(ctx context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) (EnhancedInvokableToolCallEndpoint, error) { + if h.wrapEnhancedInvokableFn != nil { + return h.wrapEnhancedInvokableFn(ctx, endpoint, tCtx), nil + } + return endpoint, nil +} + +func (h *testEnhancedToolWrapperHandler) WrapEnhancedStreamableToolCall(ctx context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) (EnhancedStreamableToolCallEndpoint, error) { + if h.wrapEnhancedStreamableFn != nil { + return h.wrapEnhancedStreamableFn(ctx, endpoint, tCtx), nil + } + return endpoint, nil +} + +func newTestEnhancedInvokableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, EnhancedInvokableToolCallEndpoint, *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + if beforeFn != nil { + beforeFn() + } + result, err := endpoint(ctx, toolArgument, opts...) + if afterFn != nil { + afterFn() + } + return result, err + } + } +} + +func newTestEnhancedStreamableToolCallWrapper(beforeFn, afterFn func()) func(context.Context, EnhancedStreamableToolCallEndpoint, *ToolContext) EnhancedStreamableToolCallEndpoint { + return func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + if beforeFn != nil { + beforeFn() + } + result, err := endpoint(ctx, toolArgument, opts...) + if afterFn != nil { + afterFn() + } + return result, err + } + } +} + +func TestHandlersToToolMiddlewaresEnhanced(t *testing.T) { + t.Run("OnlyEnhancedInvokableHandler", func(t *testing.T) { + var called bool + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + called = true + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 1) + assert.NotNil(t, middlewares[0].EnhancedInvokable) + assert.NotNil(t, middlewares[0].Invokable) + assert.NotNil(t, middlewares[0].Streamable) + assert.NotNil(t, middlewares[0].EnhancedStreamable) + + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{ + Result: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "test"}}, + }, + }, nil + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + _, err := wrapped(context.Background(), &compose.ToolInput{ + Name: "test_tool", + CallID: "call-1", + Arguments: `{"input": "test"}`, + }) + assert.NoError(t, err) + assert.True(t, called) + }) + + t.Run("OnlyEnhancedStreamableHandler", func(t *testing.T) { + var called bool + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + called = true + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 1) + assert.NotNil(t, middlewares[0].EnhancedStreamable) + assert.NotNil(t, middlewares[0].Invokable) + assert.NotNil(t, middlewares[0].Streamable) + assert.NotNil(t, middlewares[0].EnhancedInvokable) + + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{ + Result: schema.StreamReaderFromArray([]*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "test"}}}, + }), + }, nil + } + + wrapped := middlewares[0].EnhancedStreamable(mockEndpoint) + _, err := wrapped(context.Background(), &compose.ToolInput{ + Name: "test_tool", + CallID: "call-1", + Arguments: `{"input": "test"}`, + }) + assert.NoError(t, err) + assert.True(t, called) + }) + + t.Run("MixedHandlers", func(t *testing.T) { + var invokableCalled, streamableCalled, enhancedInvokableCalled, enhancedStreamableCalled bool + + handlers := []ChatModelAgentMiddleware{ + &testToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapInvokableFn: func(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) InvokableToolCallEndpoint { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + invokableCalled = true + return endpoint(ctx, argumentsInJSON, opts...) + } + }, + }, + &testToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapStreamableFn: func(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) StreamableToolCallEndpoint { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + streamableCalled = true + return endpoint(ctx, argumentsInJSON, opts...) + } + }, + }, + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + enhancedInvokableCalled = true + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + enhancedStreamableCalled = true + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 4) + + invokableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + return &compose.ToolOutput{Result: "test"}, nil + } + _, _ = middlewares[0].Invokable(invokableEndpoint)(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + streamableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + return &compose.StreamToolOutput{Result: schema.StreamReaderFromArray([]string{"test"})}, nil + } + _, _ = middlewares[1].Streamable(streamableEndpoint)(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + enhancedInvokableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil + } + _, _ = middlewares[2].EnhancedInvokable(enhancedInvokableEndpoint)(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + enhancedStreamableEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{{}})}, nil + } + _, _ = middlewares[3].EnhancedStreamable(enhancedStreamableEndpoint)(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + assert.True(t, invokableCalled) + assert.True(t, streamableCalled) + assert.True(t, enhancedInvokableCalled) + assert.True(t, enhancedStreamableCalled) + }) + + t.Run("NoHandlers", func(t *testing.T) { + handlers := []ChatModelAgentMiddleware{} + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 0) + }) + + t.Run("HandlerWithNoToolWrappers", func(t *testing.T) { + handlers := []ChatModelAgentMiddleware{ + &BaseChatModelAgentMiddleware{}, + } + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 1) + }) + + t.Run("EnhancedInvokableToolCallErrorPropagation", func(t *testing.T) { + expectedErr := errors.New("test error") + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + return nil, expectedErr + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + _, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + assert.Error(t, err) + assert.Equal(t, expectedErr, err) + }) + + t.Run("EnhancedStreamableToolCallErrorPropagation", func(t *testing.T) { + expectedErr := errors.New("test error") + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return nil, expectedErr + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{})}, nil + } + + wrapped := middlewares[0].EnhancedStreamable(mockEndpoint) + _, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + assert.Error(t, err) + assert.Equal(t, expectedErr, err) + }) + + t.Run("MultipleEnhancedInvokableWrappers", func(t *testing.T) { + var executionOrder []string + var mu sync.Mutex + + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: newTestEnhancedInvokableToolCallWrapper( + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler1-before") + mu.Unlock() + }, + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler1-after") + mu.Unlock() + }, + ), + }, + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: newTestEnhancedInvokableToolCallWrapper( + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler2-before") + mu.Unlock() + }, + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler2-after") + mu.Unlock() + }, + ), + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 2) + + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil + } + + wrapped := middlewares[0].EnhancedInvokable(middlewares[1].EnhancedInvokable(mockEndpoint)) + _, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + assert.NoError(t, err) + assert.Equal(t, []string{"handler1-before", "handler2-before", "handler2-after", "handler1-after"}, executionOrder) + }) + + t.Run("MultipleEnhancedStreamableWrappers", func(t *testing.T) { + var executionOrder []string + var mu sync.Mutex + + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: newTestEnhancedStreamableToolCallWrapper( + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler1-before") + mu.Unlock() + }, + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler1-after") + mu.Unlock() + }, + ), + }, + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: newTestEnhancedStreamableToolCallWrapper( + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler2-before") + mu.Unlock() + }, + func() { + mu.Lock() + executionOrder = append(executionOrder, "handler2-after") + mu.Unlock() + }, + ), + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + assert.Len(t, middlewares, 2) + + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{{}})}, nil + } + + wrapped := middlewares[0].EnhancedStreamable(middlewares[1].EnhancedStreamable(mockEndpoint)) + _, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + assert.NoError(t, err) + assert.Equal(t, []string{"handler1-before", "handler2-before", "handler2-after", "handler1-after"}, executionOrder) + }) +} + +func TestEnhancedToolContextPropagation(t *testing.T) { + t.Run("ToolContextContainsCorrectInfo", func(t *testing.T) { + var capturedCtx *ToolContext + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, tCtx *ToolContext) EnhancedInvokableToolCallEndpoint { + capturedCtx = tCtx + return endpoint + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + _, _ = wrapped(context.Background(), &compose.ToolInput{ + Name: "my_tool", + CallID: "call-123", + Arguments: `{"key": "value"}`, + }) + + assert.NotNil(t, capturedCtx) + assert.Equal(t, "my_tool", capturedCtx.Name) + assert.Equal(t, "call-123", capturedCtx.CallID) + }) + + t.Run("StreamableToolContextContainsCorrectInfo", func(t *testing.T) { + var capturedCtx *ToolContext + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, tCtx *ToolContext) EnhancedStreamableToolCallEndpoint { + capturedCtx = tCtx + return endpoint + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return &compose.EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{{}})}, nil + } + + wrapped := middlewares[0].EnhancedStreamable(mockEndpoint) + _, _ = wrapped(context.Background(), &compose.ToolInput{ + Name: "stream_tool", + CallID: "call-456", + Arguments: `{"data": "test"}`, + }) + + assert.NotNil(t, capturedCtx) + assert.Equal(t, "stream_tool", capturedCtx.Name) + assert.Equal(t, "call-456", capturedCtx.CallID) + }) +} + +func TestBaseChatModelAgentMiddlewareEnhancedDefaults(t *testing.T) { + t.Run("DefaultEnhancedInvokableReturnsEndpoint", func(t *testing.T) { + base := &BaseChatModelAgentMiddleware{} + + var called bool + endpoint := func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + called = true + return &schema.ToolResult{}, nil + } + + wrapped, wrapErr := base.WrapEnhancedInvokableToolCall(context.Background(), endpoint, &ToolContext{Name: "test", CallID: "1"}) + assert.NoError(t, wrapErr) + _, err := wrapped(context.Background(), &schema.ToolArgument{Text: "{}"}) + + assert.NoError(t, err) + assert.True(t, called) + }) + + t.Run("DefaultEnhancedStreamableReturnsEndpoint", func(t *testing.T) { + base := &BaseChatModelAgentMiddleware{} + + var called bool + endpoint := func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + called = true + return schema.StreamReaderFromArray([]*schema.ToolResult{}), nil + } + + wrapped, wrapErr := base.WrapEnhancedStreamableToolCall(context.Background(), endpoint, &ToolContext{Name: "test", CallID: "1"}) + assert.NoError(t, wrapErr) + _, err := wrapped(context.Background(), &schema.ToolArgument{Text: "{}"}) + + assert.NoError(t, err) + assert.True(t, called) + }) +} + +func TestEnhancedToolArgumentsPropagation(t *testing.T) { + t.Run("ArgumentsPassedCorrectly", func(t *testing.T) { + var capturedArgs string + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + capturedArgs = toolArgument.Text + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{}}, nil + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + _, _ = wrapped(context.Background(), &compose.ToolInput{ + Name: "test_tool", + CallID: "call-1", + Arguments: `{"name": "test", "value": 123}`, + }) + + assert.Equal(t, `{"name": "test", "value": 123}`, capturedArgs) + }) +} + +func TestEnhancedToolResultPropagation(t *testing.T) { + t.Run("ResultPassedThroughMiddleware", func(t *testing.T) { + expectedResult := &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "original result"}, + }, + } + + var capturedResult *schema.ToolResult + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + result, err := endpoint(ctx, toolArgument, opts...) + capturedResult = result + return result, err + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: expectedResult}, nil + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + output, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + assert.NoError(t, err) + assert.Equal(t, expectedResult, capturedResult) + assert.Equal(t, expectedResult, output.Result) + }) + + t.Run("ModifiedResultPropagated", func(t *testing.T) { + modifiedResult := &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "modified result"}, + }, + } + + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + _, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + return modifiedResult, nil + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return &compose.EnhancedInvokableToolOutput{Result: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "original"}}, + }}, nil + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + output, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + assert.NoError(t, err) + assert.Equal(t, modifiedResult, output.Result) + assert.Equal(t, "modified result", output.Result.Parts[0].Text) + }) +} + +func TestEnhancedToolEndpointErrorFromNext(t *testing.T) { + t.Run("EnhancedInvokableNextError", func(t *testing.T) { + expectedErr := errors.New("next endpoint error") + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedInvokableFn: func(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) EnhancedInvokableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + return nil, expectedErr + } + + wrapped := middlewares[0].EnhancedInvokable(mockEndpoint) + _, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + assert.Error(t, err) + assert.Equal(t, expectedErr, err) + }) + + t.Run("EnhancedStreamableNextError", func(t *testing.T) { + expectedErr := errors.New("next endpoint error") + handlers := []ChatModelAgentMiddleware{ + &testEnhancedToolWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + wrapEnhancedStreamableFn: func(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) EnhancedStreamableToolCallEndpoint { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return endpoint(ctx, toolArgument, opts...) + } + }, + }, + } + + middlewares := handlersToToolMiddlewares(handlers) + mockEndpoint := func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + return nil, expectedErr + } + + wrapped := middlewares[0].EnhancedStreamable(mockEndpoint) + _, err := wrapped(context.Background(), &compose.ToolInput{Name: "test", CallID: "1", Arguments: "{}"}) + + assert.Error(t, err) + assert.Equal(t, expectedErr, err) + }) +} + +func TestWrapModelStreamChunksPreserved(t *testing.T) { + t.Run("AgentEventMessageStreamShouldPreserveChunksWithNoopWrapModel", func(t *testing.T) { + ctx := context.Background() + + chunk1 := schema.AssistantMessage("Hello ", nil) + chunk2 := schema.AssistantMessage("World", nil) + + mockModel := &mockStreamingModel{ + chunks: []*schema.Message{chunk1, chunk2}, + } + + noopWrapModelHandler := &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return m + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + Handlers: []ChatModelAgentMiddleware{noopWrapModelHandler}, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var streamingEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && + event.Output.MessageOutput.Role == schema.Assistant { + streamingEvents = append(streamingEvents, event) + } + } + + assert.GreaterOrEqual(t, len(streamingEvents), 1, "Should have at least one streaming event") + + if len(streamingEvents) > 0 { + event := streamingEvents[0] + assert.NotNil(t, event.Output.MessageOutput.MessageStream, "Event should have message stream") + + var receivedChunks []*schema.Message + for { + chunk, recvErr := event.Output.MessageOutput.MessageStream.Recv() + if recvErr != nil { + break + } + receivedChunks = append(receivedChunks, chunk) + } + + assert.Equal(t, 2, len(receivedChunks), + "AgentEvent's MessageStream should contain 2 separate chunks, not 1 concatenated chunk. "+ + "Got %d chunks instead. This indicates the stream is being concatenated before being sent to AgentEvent.", + len(receivedChunks)) + + if len(receivedChunks) >= 2 { + assert.Equal(t, "Hello ", receivedChunks[0].Content, "First chunk content should be preserved") + assert.Equal(t, "World", receivedChunks[1].Content, "Second chunk content should be preserved") + } + } + }) + + t.Run("AgentEventMessageStreamShouldReflectUserMiddlewareModifications", func(t *testing.T) { + ctx := context.Background() + + chunk1 := schema.AssistantMessage("Hello ", nil) + chunk2 := schema.AssistantMessage("World", nil) + + mockModel := &mockStreamingModel{ + chunks: []*schema.Message{chunk1, chunk2}, + } + + streamConsumingWrapModelHandler := &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return &streamConsumingModelWrapper{inner: m} + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + Handlers: []ChatModelAgentMiddleware{streamConsumingWrapModelHandler}, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var streamingEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && + event.Output.MessageOutput.Role == schema.Assistant { + streamingEvents = append(streamingEvents, event) + } + } + + assert.GreaterOrEqual(t, len(streamingEvents), 1, "Should have at least one streaming event") + + if len(streamingEvents) > 0 { + event := streamingEvents[0] + assert.NotNil(t, event.Output.MessageOutput.MessageStream, "Event should have message stream") + + var receivedChunks []*schema.Message + for { + chunk, recvErr := event.Output.MessageOutput.MessageStream.Recv() + if recvErr != nil { + break + } + receivedChunks = append(receivedChunks, chunk) + } + + assert.Equal(t, 1, len(receivedChunks), + "AgentEvent's MessageStream should contain 1 concatenated chunk (modified by user middleware). "+ + "Got %d chunks instead.", + len(receivedChunks)) + + if len(receivedChunks) >= 1 { + assert.Equal(t, "Hello World", receivedChunks[0].Content, "Chunk content should be concatenated by user middleware") + } + } + }) + + t.Run("AgentEventMessageStreamShouldReflectMultipleUserMiddlewareModifications", func(t *testing.T) { + ctx := context.Background() + + chunk1 := schema.AssistantMessage("Hello ", nil) + chunk2 := schema.AssistantMessage("World", nil) + + mockModel := &mockStreamingModel{ + chunks: []*schema.Message{chunk1, chunk2}, + } + + handler1 := &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return &streamConsumingModelWrapper{inner: m} + }, + } + + handler2 := &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return m + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + Handlers: []ChatModelAgentMiddleware{handler1, handler2}, + ModelRetryConfig: &ModelRetryConfig{ + MaxRetries: 3, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var streamingEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && + event.Output.MessageOutput.Role == schema.Assistant { + streamingEvents = append(streamingEvents, event) + } + } + + assert.GreaterOrEqual(t, len(streamingEvents), 1, "Should have at least one streaming event") + + if len(streamingEvents) > 0 { + event := streamingEvents[0] + assert.NotNil(t, event.Output.MessageOutput.MessageStream, "Event should have message stream") + + var receivedChunks []*schema.Message + for { + chunk, recvErr := event.Output.MessageOutput.MessageStream.Recv() + if recvErr != nil { + break + } + receivedChunks = append(receivedChunks, chunk) + } + + assert.Equal(t, 1, len(receivedChunks), + "AgentEvent's MessageStream should contain 1 concatenated chunk (modified by user middleware). "+ + "Got %d chunks instead.", + len(receivedChunks)) + + if len(receivedChunks) >= 1 { + assert.Equal(t, "Hello World", receivedChunks[0].Content, "Chunk content should be concatenated by user middleware") + } + } + }) +} + +type mockStreamingModel struct { + chunks []*schema.Message +} + +func (m *mockStreamingModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return schema.ConcatMessages(m.chunks) +} + +func (m *mockStreamingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](len(m.chunks)) + go func() { + defer sw.Close() + for _, chunk := range m.chunks { + sw.Send(chunk, nil) + } + }() + return sr, nil +} + +type streamConsumingModelWrapper struct { + inner model.BaseChatModel +} + +func (m *streamConsumingModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return m.inner.Generate(ctx, input, opts...) +} + +func (m *streamConsumingModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + stream, err := m.inner.Stream(ctx, input, opts...) + if err != nil { + return nil, err + } + result, err := schema.ConcatMessageStream(stream) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{result}), nil +} + +func TestEventSenderModelWrapperCustomPosition(t *testing.T) { + t.Run("UserConfiguredEventSenderSkipsDefaultEventSender", func(t *testing.T) { + ctx := context.Background() + + chunk1 := schema.AssistantMessage("Hello ", nil) + chunk2 := schema.AssistantMessage("World", nil) + + mockModel := &mockStreamingModel{ + chunks: []*schema.Message{chunk1, chunk2}, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + Handlers: []ChatModelAgentMiddleware{NewEventSenderModelWrapper()}, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: true, + }) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var streamingEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.IsStreaming && + event.Output.MessageOutput.Role == schema.Assistant { + streamingEvents = append(streamingEvents, event) + } + } + + assert.Equal(t, 1, len(streamingEvents), "Should have exactly one streaming event (no duplicate from default event sender)") + }) + + t.Run("EventSenderAfterUserMiddlewareByDefault", func(t *testing.T) { + ctx := context.Background() + + mockModel := &mockStreamingModel{ + chunks: []*schema.Message{ + schema.AssistantMessage("Original", nil), + }, + } + + modifiedContent := "Modified" + contentModifyingHandler := &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return &contentModifyingModelWrapper{inner: m, newContent: modifiedContent} + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + Handlers: []ChatModelAgentMiddleware{contentModifyingHandler}, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + }) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var assistantEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.Role == schema.Assistant { + assistantEvents = append(assistantEvents, event) + } + } + + assert.GreaterOrEqual(t, len(assistantEvents), 1, "Should have at least one assistant event") + if len(assistantEvents) > 0 { + msg := assistantEvents[0].Output.MessageOutput.Message + assert.Equal(t, modifiedContent, msg.Content, "Event should contain modified content from user middleware") + } + }) + + t.Run("EventSenderInnermostGetsOriginalOutput", func(t *testing.T) { + ctx := context.Background() + + originalContent := "Original" + mockModel := &mockStreamingModel{ + chunks: []*schema.Message{ + schema.AssistantMessage(originalContent, nil), + }, + } + + modifiedContent := "Modified" + contentModifyingHandler := &testModelWrapperHandler{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + fn: func(_ context.Context, m model.BaseChatModel, _ *ModelContext) model.BaseChatModel { + return &contentModifyingModelWrapper{inner: m, newContent: modifiedContent} + }, + } + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + Handlers: []ChatModelAgentMiddleware{ + contentModifyingHandler, + NewEventSenderModelWrapper(), + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{ + Agent: agent, + EnableStreaming: false, + }) + iter := r.Run(ctx, []Message{schema.UserMessage("test")}) + + var assistantEvents []*AgentEvent + for { + event, ok := iter.Next() + if !ok { + break + } + if event.Output != nil && event.Output.MessageOutput != nil && + event.Output.MessageOutput.Role == schema.Assistant { + assistantEvents = append(assistantEvents, event) + } + } + + assert.GreaterOrEqual(t, len(assistantEvents), 1, "Should have at least one assistant event") + if len(assistantEvents) > 0 { + msg := assistantEvents[0].Output.MessageOutput.Message + assert.Equal(t, originalContent, msg.Content, "Event should contain original content (EventSenderModelWrapper is innermost)") + } + }) +} + +type contentModifyingModelWrapper struct { + inner model.BaseChatModel + newContent string +} + +func (m *contentModifyingModelWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + result, err := m.inner.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + result.Content = m.newContent + return result, nil +} + +func (m *contentModifyingModelWrapper) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + stream, err := m.inner.Stream(ctx, input, opts...) + if err != nil { + return nil, err + } + result, err := schema.ConcatMessageStream(stream) + if err != nil { + return nil, err + } + result.Content = m.newContent + return schema.StreamReaderFromArray([]*schema.Message{result}), nil +} + +type mockToolCallingModel struct { + mu sync.Mutex + generateCalls int + toolCallName string +} + +func (m *mockToolCallingModel) Generate(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) { + m.mu.Lock() + m.generateCalls++ + calls := m.generateCalls + m.mu.Unlock() + if calls == 1 { + return schema.AssistantMessage("calling tool", []schema.ToolCall{ + {ID: "tc-1", Function: schema.FunctionCall{Name: m.toolCallName, Arguments: `{"input":"test"}`}}, + }), nil + } + return schema.AssistantMessage("done", nil), nil +} + +func (m *mockToolCallingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + return schema.StreamReaderFromArray([]*schema.Message{msg}), nil +} + +func (m *mockToolCallingModel) WithTools(_ []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + return m, nil +} + +type invokableTestTool struct { + name string + result string +} + +func (t *invokableTestTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "test tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Desc: "input", Required: true, Type: schema.String}, + }), + }, nil +} + +func (t *invokableTestTool) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return t.result, nil +} + +type streamableTestTool struct { + name string + result string +} + +func (t *streamableTestTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "test tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Desc: "input", Required: true, Type: schema.String}, + }), + }, nil +} + +func (t *streamableTestTool) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) { + return schema.StreamReaderFromArray([]string{t.result}), nil +} + +type enhancedInvokableTestTool struct { + name string + result string +} + +func (t *enhancedInvokableTestTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "test tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Desc: "input", Required: true, Type: schema.String}, + }), + }, nil +} + +func (t *enhancedInvokableTestTool) InvokableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: t.result}}, + }, nil +} + +type enhancedStreamableTestTool struct { + name string + result string +} + +func (t *enhancedStreamableTestTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, + Desc: "test tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Desc: "input", Required: true, Type: schema.String}, + }), + }, nil +} + +func (t *enhancedStreamableTestTool) StreamableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return schema.StreamReaderFromArray([]*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: t.result}}}, + }), nil +} + +type invokableResultModifier struct { + *BaseChatModelAgentMiddleware + modifiedResult string +} + +func (h *invokableResultModifier) WrapInvokableToolCall(_ context.Context, endpoint InvokableToolCallEndpoint, _ *ToolContext) (InvokableToolCallEndpoint, error) { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + _, err := endpoint(ctx, argumentsInJSON, opts...) + if err != nil { + return "", err + } + return h.modifiedResult, nil + }, nil +} + +type streamableResultModifier struct { + *BaseChatModelAgentMiddleware + modifiedResult string +} + +func (h *streamableResultModifier) WrapStreamableToolCall(_ context.Context, endpoint StreamableToolCallEndpoint, _ *ToolContext) (StreamableToolCallEndpoint, error) { + return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + sr, err := endpoint(ctx, argumentsInJSON, opts...) + if err != nil { + return nil, err + } + sr.Close() + return schema.StreamReaderFromArray([]string{h.modifiedResult}), nil + }, nil +} + +type enhancedInvokableResultModifier struct { + *BaseChatModelAgentMiddleware + modifiedResult string +} + +func (h *enhancedInvokableResultModifier) WrapEnhancedInvokableToolCall(_ context.Context, endpoint EnhancedInvokableToolCallEndpoint, _ *ToolContext) (EnhancedInvokableToolCallEndpoint, error) { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + _, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: h.modifiedResult}}, + }, nil + }, nil +} + +type enhancedStreamableResultModifier struct { + *BaseChatModelAgentMiddleware + modifiedResult string +} + +func (h *enhancedStreamableResultModifier) WrapEnhancedStreamableToolCall(_ context.Context, endpoint EnhancedStreamableToolCallEndpoint, _ *ToolContext) (EnhancedStreamableToolCallEndpoint, error) { + return func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + sr, err := endpoint(ctx, toolArgument, opts...) + if err != nil { + return nil, err + } + sr.Close() + return schema.StreamReaderFromArray([]*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: h.modifiedResult}}}, + }), nil + }, nil +} + +func collectToolEvents(it *AsyncIterator[*AgentEvent]) []*AgentEvent { + var toolEvents []*AgentEvent + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + mo := ev.Output.MessageOutput + if mo.Message != nil && mo.Message.Role == schema.Tool { + toolEvents = append(toolEvents, ev) + continue + } + if mo.IsStreaming && mo.Role == schema.Tool && mo.MessageStream != nil { + toolEvents = append(toolEvents, ev) + } + } + return toolEvents +} + +func collectToolContent(events []*AgentEvent) []string { + var contents []string + for _, ev := range events { + mo := ev.Output.MessageOutput + if !mo.IsStreaming && mo.Message != nil { + if mo.Message.Content != "" { + contents = append(contents, mo.Message.Content) + } else if len(mo.Message.UserInputMultiContent) > 0 { + for _, part := range mo.Message.UserInputMultiContent { + if part.Text != "" { + contents = append(contents, part.Text) + } + } + } + continue + } + if mo.IsStreaming && mo.MessageStream != nil { + var msgs []*schema.Message + for { + msg, err := mo.MessageStream.Recv() + if err != nil { + break + } + msgs = append(msgs, msg) + } + if len(msgs) > 0 { + concated, err := schema.ConcatMessages(msgs) + if err == nil { + if concated.Content != "" { + contents = append(contents, concated.Content) + } else if len(concated.UserInputMultiContent) > 0 { + for _, part := range concated.UserInputMultiContent { + if part.Text != "" { + contents = append(contents, part.Text) + } + } + } + } + } + } + } + return contents +} + +func TestEventSenderToolHandler(t *testing.T) { + t.Run("Invokable", func(t *testing.T) { + t.Run("DefaultSendsEvent", func(t *testing.T) { + ctx := context.Background() + testTool := &invokableTestTool{name: "test_tool", result: "invokable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, "invokable_output") + }) + + t.Run("UserConfiguredSkipsDefault", func(t *testing.T) { + ctx := context.Background() + testTool := &invokableTestTool{name: "test_tool", result: "invokable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{NewEventSenderToolWrapper()}, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + }) + + t.Run("InnermostGetsOriginalOutput", func(t *testing.T) { + ctx := context.Background() + originalResult := "original_invokable_output" + modifiedResult := "modified_invokable_output" + testTool := &invokableTestTool{name: "test_tool", result: originalResult} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &invokableResultModifier{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + modifiedResult: modifiedResult, + }, + NewEventSenderToolWrapper(), + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.GreaterOrEqual(t, len(toolEvents), 1) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, originalResult) + }) + }) + + t.Run("Streamable", func(t *testing.T) { + t.Run("DefaultSendsEvent", func(t *testing.T) { + ctx := context.Background() + testTool := &streamableTestTool{name: "test_tool", result: "streamable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, "streamable_output") + }) + + t.Run("UserConfiguredSkipsDefault", func(t *testing.T) { + ctx := context.Background() + testTool := &streamableTestTool{name: "test_tool", result: "streamable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{NewEventSenderToolWrapper()}, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + }) + + t.Run("InnermostGetsOriginalOutput", func(t *testing.T) { + ctx := context.Background() + originalResult := "original_streamable_output" + modifiedResult := "modified_streamable_output" + testTool := &streamableTestTool{name: "test_tool", result: originalResult} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &streamableResultModifier{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + modifiedResult: modifiedResult, + }, + NewEventSenderToolWrapper(), + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.GreaterOrEqual(t, len(toolEvents), 1) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, originalResult) + }) + }) + + t.Run("EnhancedInvokable", func(t *testing.T) { + t.Run("DefaultSendsEvent", func(t *testing.T) { + ctx := context.Background() + testTool := &enhancedInvokableTestTool{name: "test_tool", result: "enhanced_invokable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, "enhanced_invokable_output") + }) + + t.Run("UserConfiguredSkipsDefault", func(t *testing.T) { + ctx := context.Background() + testTool := &enhancedInvokableTestTool{name: "test_tool", result: "enhanced_invokable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{NewEventSenderToolWrapper()}, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + }) + + t.Run("InnermostGetsOriginalOutput", func(t *testing.T) { + ctx := context.Background() + originalResult := "original_enhanced_invokable_output" + modifiedResult := "modified_enhanced_invokable_output" + testTool := &enhancedInvokableTestTool{name: "test_tool", result: originalResult} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &enhancedInvokableResultModifier{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + modifiedResult: modifiedResult, + }, + NewEventSenderToolWrapper(), + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: false}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.GreaterOrEqual(t, len(toolEvents), 1) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, originalResult) + }) + }) + + t.Run("EnhancedStreamable", func(t *testing.T) { + t.Run("DefaultSendsEvent", func(t *testing.T) { + ctx := context.Background() + testTool := &enhancedStreamableTestTool{name: "test_tool", result: "enhanced_streamable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, "enhanced_streamable_output") + }) + + t.Run("UserConfiguredSkipsDefault", func(t *testing.T) { + ctx := context.Background() + testTool := &enhancedStreamableTestTool{name: "test_tool", result: "enhanced_streamable_output"} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{NewEventSenderToolWrapper()}, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + }) + + t.Run("InnermostGetsOriginalOutput", func(t *testing.T) { + ctx := context.Background() + originalResult := "original_enhanced_streamable_output" + modifiedResult := "modified_enhanced_streamable_output" + testTool := &enhancedStreamableTestTool{name: "test_tool", result: originalResult} + mockModel := &mockToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewChatModelAgent(ctx, &ChatModelAgentConfig{ + Name: "TestAgent", + Description: "Test agent", + Model: mockModel, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + }, + }, + Handlers: []ChatModelAgentMiddleware{ + &enhancedStreamableResultModifier{ + BaseChatModelAgentMiddleware: &BaseChatModelAgentMiddleware{}, + modifiedResult: modifiedResult, + }, + NewEventSenderToolWrapper(), + }, + }) + assert.NoError(t, err) + + r := NewRunner(ctx, RunnerConfig{Agent: agent, EnableStreaming: true}) + it := r.Run(ctx, []Message{schema.UserMessage("test")}) + + toolEvents := collectToolEvents(it) + assert.GreaterOrEqual(t, len(toolEvents), 1) + contents := collectToolContent(toolEvents) + assert.Contains(t, contents, originalResult) + }) + }) +} + +// mockAgenticToolCallingModel is a model.BaseModel[*schema.AgenticMessage] that +// returns a tool call on the first Generate, then a final answer on the second. +type mockAgenticToolCallingModel struct { + toolCallName string + callCount int32 +} + +func (m *mockAgenticToolCallingModel) Generate(_ context.Context, _ []*schema.AgenticMessage, _ ...model.Option) (*schema.AgenticMessage, error) { + idx := atomic.AddInt32(&m.callCount, 1) + if idx == 1 { + return agenticToolCallMsg(m.toolCallName, "tc-1", `{"input":"test"}`), nil + } + return agenticMsg("done"), nil +} + +func (m *mockAgenticToolCallingModel) Stream(ctx context.Context, input []*schema.AgenticMessage, opts ...model.Option) (*schema.StreamReader[*schema.AgenticMessage], error) { + msg, err := m.Generate(ctx, input, opts...) + if err != nil { + return nil, err + } + r, w := schema.Pipe[*schema.AgenticMessage](1) + go func() { defer w.Close(); w.Send(msg, nil) }() + return r, nil +} + +// collectAgenticToolEvents filters tool result events from the agentic iterator. +// Agentic tool results have AgenticRole == AgenticRoleTypeUser and contain +// FunctionToolResult content blocks. +func collectAgenticToolEvents(it *AsyncIterator[*agenticAgentEvent]) []*agenticAgentEvent { + var toolEvents []*agenticAgentEvent + for { + ev, ok := it.Next() + if !ok { + break + } + if ev.Output == nil || ev.Output.MessageOutput == nil { + continue + } + mo := ev.Output.MessageOutput + if mo.AgenticRole == schema.AgenticRoleTypeUser { + toolEvents = append(toolEvents, ev) + } + } + return toolEvents +} + +// collectAgenticToolContent extracts text from agentic tool result events. +func collectAgenticToolContent(events []*agenticAgentEvent) []string { + var contents []string + for _, ev := range events { + mo := ev.Output.MessageOutput + if !mo.IsStreaming && mo.Message != nil { + for _, cb := range mo.Message.ContentBlocks { + if cb.FunctionToolResult != nil { + for _, b := range cb.FunctionToolResult.Content { + if b.Text != nil { + contents = append(contents, b.Text.Text) + } + } + } + } + continue + } + if mo.IsStreaming && mo.MessageStream != nil { + for { + msg, err := mo.MessageStream.Recv() + if err != nil { + break + } + for _, cb := range msg.ContentBlocks { + if cb.FunctionToolResult != nil { + for _, b := range cb.FunctionToolResult.Content { + if b.Text != nil { + contents = append(contents, b.Text.Text) + } + } + } + } + } + } + } + return contents +} + +func newAgenticEventSenderToolWrapper() TypedChatModelAgentMiddleware[*schema.AgenticMessage] { + return &typedEventSenderToolWrapper[*schema.AgenticMessage]{ + TypedBaseChatModelAgentMiddleware: &TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{}, + } +} + +// TestAgenticEventSenderToolHandler exercises the *schema.AgenticMessage branches +// in typedToolInvokeEvent, typedToolStreamEvent, typedToolEnhancedInvokeEvent, +// typedToolEnhancedStreamEvent, plus the helpers textToFunctionToolResultBlocks, +// toolResultToBlocks, and derefString. +func TestAgenticEventSenderToolHandler(t *testing.T) { + t.Run("Invokable", func(t *testing.T) { + ctx := context.Background() + testTool := &invokableTestTool{name: "test_tool", result: "invokable_output"} + mdl := &mockAgenticToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newAgenticEventSenderToolWrapper()}, + }) + require.NoError(t, err) + + r := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent, EnableStreaming: false}) + it := r.Query(ctx, "test") + + toolEvents := collectAgenticToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectAgenticToolContent(toolEvents) + assert.Contains(t, contents, "invokable_output") + }) + + t.Run("Streamable", func(t *testing.T) { + ctx := context.Background() + testTool := &streamableTestTool{name: "test_tool", result: "streamable_output"} + mdl := &mockAgenticToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newAgenticEventSenderToolWrapper()}, + }) + require.NoError(t, err) + + r := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent, EnableStreaming: true}) + it := r.Query(ctx, "test") + + toolEvents := collectAgenticToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectAgenticToolContent(toolEvents) + assert.Contains(t, contents, "streamable_output") + }) + + t.Run("EnhancedInvokable", func(t *testing.T) { + ctx := context.Background() + testTool := &enhancedInvokableTestTool{name: "test_tool", result: "enhanced_output"} + mdl := &mockAgenticToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newAgenticEventSenderToolWrapper()}, + }) + require.NoError(t, err) + + r := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent, EnableStreaming: false}) + it := r.Query(ctx, "test") + + toolEvents := collectAgenticToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectAgenticToolContent(toolEvents) + assert.Contains(t, contents, "enhanced_output") + }) + + t.Run("EnhancedStreamable", func(t *testing.T) { + ctx := context.Background() + testTool := &enhancedStreamableTestTool{name: "test_tool", result: "enhanced_stream_output"} + mdl := &mockAgenticToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newAgenticEventSenderToolWrapper()}, + }) + require.NoError(t, err) + + r := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent, EnableStreaming: true}) + it := r.Query(ctx, "test") + + toolEvents := collectAgenticToolEvents(it) + assert.Equal(t, 1, len(toolEvents)) + contents := collectAgenticToolContent(toolEvents) + assert.Contains(t, contents, "enhanced_stream_output") + }) + + t.Run("EnhancedInvokableMultimodal", func(t *testing.T) { + ctx := context.Background() + imgURL := "https://example.com/img.png" + testTool := &multimodalEnhancedInvokableTestTool{ + name: "test_tool", + result: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "caption"}, + {Type: schema.ToolPartTypeImage, Image: &schema.ToolOutputImage{MessagePartCommon: schema.MessagePartCommon{URL: &imgURL}}}, + }, + }, + } + mdl := &mockAgenticToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newAgenticEventSenderToolWrapper()}, + }) + require.NoError(t, err) + + r := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent, EnableStreaming: false}) + it := r.Query(ctx, "test") + + toolEvents := collectAgenticToolEvents(it) + require.Equal(t, 1, len(toolEvents)) + + // Verify multimodal content + msg := toolEvents[0].Output.MessageOutput.Message + require.NotNil(t, msg) + require.Len(t, msg.ContentBlocks, 1) + ftr := msg.ContentBlocks[0].FunctionToolResult + require.NotNil(t, ftr) + require.Len(t, ftr.Content, 2) + assert.Equal(t, "caption", ftr.Content[0].Text.Text) + assert.Equal(t, "https://example.com/img.png", ftr.Content[1].Image.URL) + }) + + t.Run("EnhancedStreamableMultimodal", func(t *testing.T) { + ctx := context.Background() + audioURL := "https://example.com/audio.mp3" + testTool := &multimodalEnhancedStreamableTestTool{ + name: "test_tool", + result: &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "transcript"}, + {Type: schema.ToolPartTypeAudio, Audio: &schema.ToolOutputAudio{MessagePartCommon: schema.MessagePartCommon{URL: &audioURL}}}, + }, + }, + } + mdl := &mockAgenticToolCallingModel{toolCallName: "test_tool"} + + agent, err := NewTypedChatModelAgent(ctx, &TypedChatModelAgentConfig[*schema.AgenticMessage]{ + Name: "TestAgent", + Description: "test", + Model: mdl, + ToolsConfig: ToolsConfig{ + ToolsNodeConfig: compose.ToolsNodeConfig{Tools: []tool.BaseTool{testTool}}, + }, + Handlers: []TypedChatModelAgentMiddleware[*schema.AgenticMessage]{newAgenticEventSenderToolWrapper()}, + }) + require.NoError(t, err) + + r := NewTypedRunner(TypedRunnerConfig[*schema.AgenticMessage]{Agent: agent, EnableStreaming: true}) + it := r.Query(ctx, "test") + + toolEvents := collectAgenticToolEvents(it) + require.Equal(t, 1, len(toolEvents)) + + // Drain the stream and verify multimodal content + mo := toolEvents[0].Output.MessageOutput + require.True(t, mo.IsStreaming) + var allBlocks []*schema.FunctionToolResultContentBlock + for { + msg, err := mo.MessageStream.Recv() + if err != nil { + break + } + for _, cb := range msg.ContentBlocks { + if cb.FunctionToolResult != nil { + allBlocks = append(allBlocks, cb.FunctionToolResult.Content...) + } + } + } + require.Len(t, allBlocks, 2) + assert.Equal(t, "transcript", allBlocks[0].Text.Text) + assert.Equal(t, "https://example.com/audio.mp3", allBlocks[1].Audio.URL) + }) +} + +func TestTypedToolStreamEventAgenticMessageSetsStreamingMeta(t *testing.T) { + event := typedToolStreamEvent[*schema.AgenticMessage]( + "call_1", + "execute", + "msg_1", + schema.StreamReaderFromArray([]string{"first\n", "second\n"}), + ) + require.NotNil(t, event) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.True(t, event.Output.MessageOutput.IsStreaming) + require.NotNil(t, event.Output.MessageOutput.MessageStream) + + first, err := event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + require.Len(t, first.ContentBlocks, 1) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, first.ContentBlocks[0].StreamingMeta) + + second, err := event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + require.Len(t, second.ContentBlocks, 1) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, second.ContentBlocks[0].StreamingMeta) + + result, err := schema.ConcatAgenticMessages([]*schema.AgenticMessage{first, second}) + require.NoError(t, err) + require.Len(t, result.ContentBlocks, 1) + assert.Nil(t, result.ContentBlocks[0].StreamingMeta) + require.NotNil(t, result.ContentBlocks[0].FunctionToolResult) + require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 2) + assert.Equal(t, "first\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) + assert.Equal(t, "second\n", result.ContentBlocks[0].FunctionToolResult.Content[1].Text.Text) +} + +func TestTypedToolEnhancedStreamEventAgenticMessageSetsStreamingMeta(t *testing.T) { + event := typedToolEnhancedStreamEvent[*schema.AgenticMessage]( + "call_1", + "execute", + "msg_1", + schema.StreamReaderFromArray([]*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "first\n"}}}, + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "second\n"}}}, + }), + ) + require.NotNil(t, event) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + require.True(t, event.Output.MessageOutput.IsStreaming) + require.NotNil(t, event.Output.MessageOutput.MessageStream) + + first, err := event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + require.Len(t, first.ContentBlocks, 1) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, first.ContentBlocks[0].StreamingMeta) + + second, err := event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + require.Len(t, second.ContentBlocks, 1) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, second.ContentBlocks[0].StreamingMeta) + + result, err := schema.ConcatAgenticMessages([]*schema.AgenticMessage{first, second}) + require.NoError(t, err) + require.Len(t, result.ContentBlocks, 1) + assert.Nil(t, result.ContentBlocks[0].StreamingMeta) + require.NotNil(t, result.ContentBlocks[0].FunctionToolResult) + require.Len(t, result.ContentBlocks[0].FunctionToolResult.Content, 2) + assert.Equal(t, "first\n", result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) + assert.Equal(t, "second\n", result.ContentBlocks[0].FunctionToolResult.Content[1].Text.Text) +} + +// multimodalEnhancedInvokableTestTool returns a pre-built multimodal ToolResult. +type multimodalEnhancedInvokableTestTool struct { + name string + result *schema.ToolResult +} + +func (t *multimodalEnhancedInvokableTestTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, Desc: "multimodal test tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Desc: "input", Required: true, Type: schema.String}, + }), + }, nil +} + +func (t *multimodalEnhancedInvokableTestTool) InvokableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return t.result, nil +} + +// multimodalEnhancedStreamableTestTool returns a pre-built multimodal ToolResult as a stream. +type multimodalEnhancedStreamableTestTool struct { + name string + result *schema.ToolResult +} + +func (t *multimodalEnhancedStreamableTestTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: t.name, Desc: "multimodal streaming test tool", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Desc: "input", Required: true, Type: schema.String}, + }), + }, nil +} + +func (t *multimodalEnhancedStreamableTestTool) StreamableRun(_ context.Context, _ *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return schema.StreamReaderFromArray([]*schema.ToolResult{t.result}), nil +} + +func Test_functionToolResultAgenticMessage(t *testing.T) { + t.Run("basic", func(t *testing.T) { + blocks := []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "result_str"}}, + } + msg := functionToolResultAgenticMessage("call_1", "tool_name", blocks) + assert.Equal(t, schema.AgenticRoleTypeUser, msg.Role) + assert.Len(t, msg.ContentBlocks, 1) + assert.Equal(t, schema.ContentBlockTypeFunctionToolResult, msg.ContentBlocks[0].Type) + ftr := msg.ContentBlocks[0].FunctionToolResult + assert.Equal(t, "call_1", ftr.CallID) + assert.Equal(t, "tool_name", ftr.Name) + assert.Len(t, ftr.Content, 1) + assert.Equal(t, "result_str", ftr.Content[0].Text.Text) + }) + + t.Run("multimodal", func(t *testing.T) { + blocks := []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "description"}}, + {Type: schema.FunctionToolResultContentBlockTypeImage, Image: &schema.UserInputImage{URL: "https://example.com/img.png"}}, + } + msg := functionToolResultAgenticMessage("call_2", "vision_tool", blocks) + assert.Equal(t, schema.AgenticRoleTypeUser, msg.Role) + ftr := msg.ContentBlocks[0].FunctionToolResult + assert.Equal(t, "call_2", ftr.CallID) + assert.Equal(t, "vision_tool", ftr.Name) + assert.Len(t, ftr.Content, 2) + assert.Equal(t, "description", ftr.Content[0].Text.Text) + assert.Equal(t, "https://example.com/img.png", ftr.Content[1].Image.URL) + }) +} + +func TestTypedToolEnhancedEventAgenticToolSearchResult(t *testing.T) { + result := &schema.ToolResult{Parts: []schema.ToolOutputPart{ + { + Type: schema.ToolPartTypeToolSearchResult, + ToolSearchResult: &schema.ToolSearchResult{Tools: []*schema.ToolInfo{ + {Name: "dynamic_tool", Desc: "dynamic tool"}, + }}, + }, + }} + + t.Run("invoke", func(t *testing.T) { + event, err := typedToolEnhancedInvokeEvent[*schema.AgenticMessage]("call_1", "tool_search", "msg_1", result) + require.NoError(t, err) + require.NotNil(t, event) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg := event.Output.MessageOutput.Message + require.NotNil(t, msg) + require.Len(t, msg.ContentBlocks, 1) + block := msg.ContentBlocks[0] + assert.Equal(t, schema.ContentBlockTypeToolSearchResult, block.Type) + require.NotNil(t, block.ToolSearchFunctionToolResult) + assert.Equal(t, "call_1", block.ToolSearchFunctionToolResult.CallID) + assert.Equal(t, "tool_search", block.ToolSearchFunctionToolResult.Name) + require.NotNil(t, block.ToolSearchFunctionToolResult.Result) + require.Len(t, block.ToolSearchFunctionToolResult.Result.Tools, 1) + assert.Equal(t, "dynamic_tool", block.ToolSearchFunctionToolResult.Result.Tools[0].Name) + }) + + t.Run("stream", func(t *testing.T) { + event := typedToolEnhancedStreamEvent[*schema.AgenticMessage]( + "call_2", + "tool_search", + "msg_2", + schema.StreamReaderFromArray([]*schema.ToolResult{result}), + ) + require.NotNil(t, event) + require.NotNil(t, event.Output) + require.NotNil(t, event.Output.MessageOutput) + + msg, err := event.Output.MessageOutput.MessageStream.Recv() + require.NoError(t, err) + require.NotNil(t, msg) + require.Len(t, msg.ContentBlocks, 1) + block := msg.ContentBlocks[0] + assert.Equal(t, schema.ContentBlockTypeToolSearchResult, block.Type) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, block.StreamingMeta) + require.NotNil(t, block.ToolSearchFunctionToolResult) + assert.Equal(t, "call_2", block.ToolSearchFunctionToolResult.CallID) + assert.Equal(t, "tool_search", block.ToolSearchFunctionToolResult.Name) + require.NotNil(t, block.ToolSearchFunctionToolResult.Result) + require.Len(t, block.ToolSearchFunctionToolResult.Result.Tools, 1) + assert.Equal(t, "dynamic_tool", block.ToolSearchFunctionToolResult.Result.Tools[0].Name) + }) +} + +func TestExtractToolIdentifiersToolSearchResult(t *testing.T) { + msg := &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + schema.NewContentBlock(&schema.ToolSearchFunctionToolResult{ + CallID: "call_1", + Name: "tool_search", + Result: &schema.ToolSearchResult{}, + }), + }, + } + + toolName, callID := extractToolIdentifiers(msg) + assert.Equal(t, "tool_search", toolName) + assert.Equal(t, "call_1", callID) +} diff --git a/callbacks/aspect_inject.go b/callbacks/aspect_inject.go new file mode 100644 index 0000000..2b91e25 --- /dev/null +++ b/callbacks/aspect_inject.go @@ -0,0 +1,174 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/internal/callbacks" + "github.com/cloudwego/eino/schema" +) + +// OnStart Fast inject callback input / output aspect for component developer +// e.g. +// +// func (t *testChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (resp *schema.Message, err error) { +// defer func() { +// if err != nil { +// callbacks.OnError(ctx, err) +// } +// }() +// +// ctx = callbacks.OnStart(ctx, &model.CallbackInput{ +// Messages: input, +// Tools: nil, +// Extra: nil, +// }) +// +// // do smt +// +// ctx = callbacks.OnEnd(ctx, &model.CallbackOutput{ +// Message: resp, +// Extra: nil, +// }) +// +// return resp, nil +// } + +// OnStart invokes the OnStart timing for all registered handlers in the +// context. This is called by component implementations that manage their own +// callbacks (i.e. implement [components.Checker] and return true from +// IsCallbacksEnabled). The returned context must be propagated to subsequent +// OnEnd/OnError calls so handlers can correlate start and end events. +// +// Handlers are invoked in reverse registration order (last registered = first +// called) to match the middleware wrapping convention. +// +// Example — typical usage inside a component's Generate method: +// +// func (m *myChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { +// ctx = callbacks.OnStart(ctx, &model.CallbackInput{Messages: input}) +// resp, err := m.doGenerate(ctx, input, opts...) +// if err != nil { +// callbacks.OnError(ctx, err) +// return nil, err +// } +// callbacks.OnEnd(ctx, &model.CallbackOutput{Message: resp}) +// return resp, nil +// } +func OnStart[T any](ctx context.Context, input T) context.Context { + ctx, _ = callbacks.On(ctx, input, callbacks.OnStartHandle[T], TimingOnStart, true) + + return ctx +} + +// OnEnd invokes the OnEnd timing for all registered handlers. Call this after +// the component produces a successful result. Handlers run in registration +// order (first registered = first called). +// +// Do not call both OnEnd and OnError for the same invocation — OnEnd signals +// success; OnError signals failure. +func OnEnd[T any](ctx context.Context, output T) context.Context { + ctx, _ = callbacks.On(ctx, output, callbacks.OnEndHandle[T], TimingOnEnd, false) + + return ctx +} + +// OnStartWithStreamInput invokes the OnStartWithStreamInput timing. Use this +// when the component's input is itself a stream (Collect / Transform +// paradigms). The framework automatically copies the stream so each handler +// receives an independent reader; handlers MUST close their copy or the +// underlying goroutine will leak. +// +// Returns the updated context and a new StreamReader that the component should +// use going forward (the original is consumed by the framework). +func OnStartWithStreamInput[T any](ctx context.Context, input *schema.StreamReader[T]) ( + nextCtx context.Context, newStreamReader *schema.StreamReader[T]) { + + return callbacks.On(ctx, input, callbacks.OnStartWithStreamInputHandle[T], TimingOnStartWithStreamInput, true) +} + +// OnEndWithStreamOutput invokes the OnEndWithStreamOutput timing. Use this +// when the component produces a streaming output (Stream / Transform +// paradigms). Like OnStartWithStreamInput, stream copies are made per +// handler; each handler must close its copy. +// +// Returns the updated context and the StreamReader the component should return +// to its caller. +func OnEndWithStreamOutput[T any](ctx context.Context, output *schema.StreamReader[T]) ( + nextCtx context.Context, newStreamReader *schema.StreamReader[T]) { + + return callbacks.On(ctx, output, callbacks.OnEndWithStreamOutputHandle[T], TimingOnEndWithStreamOutput, false) +} + +// OnError invokes the OnError timing for all registered handlers. Call this +// when the component returns an error. Errors that occur mid-stream (after the +// StreamReader has been returned) are NOT routed through OnError; they surface +// as errors inside Recv. +// +// Handlers run in registration order (same as OnEnd). +func OnError(ctx context.Context, err error) context.Context { + ctx, _ = callbacks.On(ctx, err, callbacks.OnErrorHandle, TimingOnError, false) + + return ctx +} + +// EnsureRunInfo ensures the context carries a [RunInfo] for the given type and +// component kind. If the context already has a matching RunInfo, it is +// returned unchanged. Otherwise, a new callback manager is created that +// inherits the global handlers plus any handlers already in ctx. +// +// Component implementations that set IsCallbacksEnabled() = true should call +// this at the start of every public method (Generate, Stream, etc.) before +// calling [OnStart], so that the RunInfo is never missing from callbacks. +func EnsureRunInfo(ctx context.Context, typ string, comp components.Component) context.Context { + return callbacks.EnsureRunInfo(ctx, typ, comp) +} + +// ReuseHandlers creates a new context that inherits all handlers already +// present in ctx and sets a new RunInfo. Global handlers are added if ctx +// carries none yet. +// +// Use this when a component calls another component internally and wants the +// inner component's callbacks to share the same set of handlers as the outer +// component, but with the inner component's own identity in RunInfo: +// +// innerCtx := callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{ +// Type: "InnerChatModel", +// Component: components.ComponentOfChatModel, +// Name: "inner-cm", +// }) +func ReuseHandlers(ctx context.Context, info *RunInfo) context.Context { + return callbacks.ReuseHandlers(ctx, info) +} + +// InitCallbacks creates a new context with the given RunInfo and handlers, +// completely replacing any RunInfo and handlers already in ctx. +// +// Use this when running a component standalone outside a Graph — the Graph +// normally manages RunInfo injection automatically, but standalone callers must +// set it up themselves: +// +// ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{ +// Type: myModel.GetType(), +// Component: components.ComponentOfChatModel, +// Name: "my-model", +// }, myHandler) +func InitCallbacks(ctx context.Context, info *RunInfo, handlers ...Handler) context.Context { + return callbacks.InitCallbacks(ctx, info, handlers...) +} diff --git a/callbacks/aspect_inject_test.go b/callbacks/aspect_inject_test.go new file mode 100644 index 0000000..76d73f7 --- /dev/null +++ b/callbacks/aspect_inject_test.go @@ -0,0 +1,330 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + "fmt" + "io" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/callbacks" + "github.com/cloudwego/eino/schema" +) + +func TestAspectInject(t *testing.T) { + t.Run("ctx without manager", func(t *testing.T) { + ctx := context.Background() + ctx = OnStart(ctx, 1) + ctx = OnEnd(ctx, 2) + ctx = OnError(ctx, fmt.Errorf("3")) + isr, isw := schema.Pipe[int](2) + go func() { + for i := 0; i < 10; i++ { + isw.Send(i, nil) + } + isw.Close() + }() + + var nisr *schema.StreamReader[int] + ctx, nisr = OnStartWithStreamInput(ctx, isr) + j := 0 + for { + i, err := nisr.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, j, i) + j++ + } + nisr.Close() + + osr, osw := schema.Pipe[int](2) + go func() { + for i := 0; i < 10; i++ { + osw.Send(i, nil) + } + osw.Close() + }() + + var nosr *schema.StreamReader[int] + ctx, nosr = OnEndWithStreamOutput(ctx, osr) + j = 0 + for { + i, err := nosr.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, j, i) + j++ + } + nosr.Close() + }) + + t.Run("ctx with manager", func(t *testing.T) { + ctx := context.Background() + cnt := 0 + + hb := NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { + cnt += input.(int) + return ctx + }). + OnEndFn(func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context { + cnt += output.(int) + return ctx + }). + OnErrorFn(func(ctx context.Context, info *RunInfo, err error) context.Context { + v, _ := strconv.ParseInt(err.Error(), 10, 64) + cnt += int(v) + return ctx + }). + OnStartWithStreamInputFn(func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context { + for { + i, err := input.Recv() + if err == io.EOF { + break + } + + cnt += i.(int) + } + + input.Close() + return ctx + }). + OnEndWithStreamOutputFn(func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context { + for { + o, err := output.Recv() + if err == io.EOF { + break + } + + cnt += o.(int) + } + + output.Close() + return ctx + }).Build() + + ctx = InitCallbacks(ctx, nil, hb) + + ctx = OnStart(ctx, 1) + ctx = OnEnd(ctx, 2) + ctx = OnError(ctx, fmt.Errorf("3")) + isr, isw := schema.Pipe[int](2) + go func() { + for i := 0; i < 10; i++ { + isw.Send(i, nil) + } + isw.Close() + }() + + ctx = ReuseHandlers(ctx, &RunInfo{}) + var nisr *schema.StreamReader[int] + ctx, nisr = OnStartWithStreamInput(ctx, isr) + j := 0 + for { + i, err := nisr.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, j, i) + j++ + cnt += i + } + nisr.Close() + + osr, osw := schema.Pipe[int](2) + go func() { + for i := 0; i < 10; i++ { + osw.Send(i, nil) + } + osw.Close() + }() + + var nosr *schema.StreamReader[int] + ctx, nosr = OnEndWithStreamOutput(ctx, osr) + j = 0 + for { + i, err := nosr.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, j, i) + j++ + cnt += i + } + nosr.Close() + assert.Equal(t, 186, cnt) + }) +} + +func TestGlobalCallbacksRepeated(t *testing.T) { + times := 0 + testHandler := NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + times++ + return ctx + }).Build() + callbacks.GlobalHandlers = append(callbacks.GlobalHandlers, testHandler) + + ctx := context.Background() + ctx = callbacks.AppendHandlers(ctx, &RunInfo{}) + ctx = callbacks.AppendHandlers(ctx, &RunInfo{}) + + callbacks.On(ctx, "test", callbacks.OnStartHandle[string], TimingOnStart, true) + assert.Equal(t, times, 1) +} + +func TestEnsureRunInfo(t *testing.T) { + ctx := context.Background() + + var name, typ, comp string + ctx = InitCallbacks(ctx, &RunInfo{Name: "name", Type: "type", Component: "component"}, NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { + name = info.Name + typ = info.Type + comp = string(info.Component) + return ctx + }).Build()) + + ctx = OnStart(ctx, "") + assert.Equal(t, "name", name) + assert.Equal(t, "type", typ) + assert.Equal(t, "component", comp) + ctx2 := EnsureRunInfo(ctx, "type2", "component2") + OnStart(ctx2, "") + assert.Equal(t, "", name) + assert.Equal(t, "type2", typ) + assert.Equal(t, "component2", comp) + + // EnsureRunInfo on an empty Context + AppendGlobalHandlers(NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { + typ = info.Type + comp = string(info.Component) + return ctx + }).Build()) + ctx3 := EnsureRunInfo(context.Background(), "type3", "component3") + OnStart(ctx3, 0) + assert.Equal(t, "type3", typ) + assert.Equal(t, "component3", comp) + callbacks.GlobalHandlers = []Handler{} +} + +func TestNesting(t *testing.T) { + ctx := context.Background() + cb := &myCallback{t: t} + ctx = InitCallbacks(ctx, &RunInfo{ + Name: "test", + }, cb) + + // jumped + ctx1 := OnStart(ctx, 0) + ctx2 := OnStart(ctx1, 1) + OnEnd(ctx2, 1) + OnEnd(ctx1, 0) + assert.Equal(t, 4, cb.times) + + // reused + cb.times = 0 + ctx1 = OnStart(ctx, 0) + ctx2 = ReuseHandlers(ctx1, &RunInfo{Name: "test2"}) + ctx3 := OnStart(ctx2, 1) + OnEnd(ctx3, 1) + OnEnd(ctx1, 0) + assert.Equal(t, 4, cb.times) + +} + +func TestReuseHandlersOnEmptyCtx(t *testing.T) { + callbacks.GlobalHandlers = []Handler{} + cb := &myCallback{t: t} + AppendGlobalHandlers(cb) + ctx := ReuseHandlers(context.Background(), &RunInfo{Name: "test"}) + OnStart(ctx, 0) + assert.Equal(t, 1, cb.times) +} + +func TestAppendHandlersTwiceOnSameCtx(t *testing.T) { + callbacks.GlobalHandlers = []Handler{} + cb := &myCallback{t: t} + cb1 := &myCallback{t: t} + cb2 := &myCallback{t: t} + ctx := InitCallbacks(context.Background(), &RunInfo{Name: "test"}, cb) + ctx1 := callbacks.AppendHandlers(ctx, &RunInfo{Name: "test"}, cb1) + ctx2 := callbacks.AppendHandlers(ctx, &RunInfo{Name: "test"}, cb2) + OnStart(ctx1, 0) + OnStart(ctx2, 0) + assert.Equal(t, 2, cb.times) + assert.Equal(t, 1, cb1.times) + assert.Equal(t, 1, cb2.times) +} + +type myCallback struct { + t *testing.T + times int +} + +func (m *myCallback) OnStart(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { + m.times++ + if info == nil { + assert.Equal(m.t, 2, m.times) + return ctx + } + if info.Name == "test" { + assert.Equal(m.t, 0, input) + } else { + assert.Equal(m.t, 1, input) + } + + return ctx +} + +func (m *myCallback) OnEnd(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context { + m.times++ + if info == nil { + assert.Equal(m.t, 3, m.times) + return ctx + } + if info.Name == "test" { + assert.Equal(m.t, 0, output) + } else { + assert.Equal(m.t, 1, output) + } + return ctx +} + +func (m *myCallback) OnError(ctx context.Context, info *RunInfo, err error) context.Context { + panic("implement me") +} + +func (m *myCallback) OnStartWithStreamInput(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context { + panic("implement me") +} + +func (m *myCallback) OnEndWithStreamOutput(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context { + panic("implement me") +} diff --git a/callbacks/doc.go b/callbacks/doc.go new file mode 100644 index 0000000..e2a3ca0 --- /dev/null +++ b/callbacks/doc.go @@ -0,0 +1,129 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package callbacks provides observability hooks for component execution in Eino. +// +// Callbacks fire at five lifecycle timings around every component invocation: +// - [TimingOnStart] / [TimingOnEnd]: non-streaming input and output. +// - [TimingOnStartWithStreamInput] / [TimingOnEndWithStreamOutput]: streaming +// variants — handlers receive a copy of the stream and MUST close it. +// - [TimingOnError]: component returned a non-nil error (stream-internal +// errors are NOT reported here). +// +// # Attaching Handlers +// +// Global handlers (observe every node in every graph): +// +// callbacks.AppendGlobalHandlers(myHandler) // call once, at startup — NOT thread-safe +// +// Per-invocation handlers (observe one graph run): +// +// runnable.Invoke(ctx, input, compose.WithCallbacks(myHandler)) +// +// Target a specific node: +// +// compose.WithCallbacks(myHandler).DesignateNode("nodeName") +// +// Handler inheritance: if the context passed to a graph run already carries +// handlers (e.g. from a parent graph), those handlers are inherited by the +// entire child run automatically. +// +// # Building Handlers +// +// Option 1 — [NewHandlerBuilder]: register raw functions for the timings you +// need. Input/output are untyped; use the component package's ConvCallbackInput +// helper to cast to a concrete type: +// +// handler := callbacks.NewHandlerBuilder(). +// OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { +// // Handle component start +// return ctx +// }). +// OnEndFn(func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context { +// // Handle component end +// return ctx +// }). +// OnErrorFn(func(ctx context.Context, info *RunInfo, err error) context.Context { +// // Handle component error +// return ctx +// }). +// OnStartWithStreamInputFn(func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context { +// defer input.Close() // MUST close — failure causes pipeline goroutine leak +// return ctx +// }). +// OnEndWithStreamOutputFn(func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context { +// defer output.Close() // MUST close +// return ctx +// }). +// Build() +// +// Option 2 — utils/callbacks.NewHandlerHelper: dispatches by component type, so +// each handler function receives the concrete typed input/output directly: +// +// handler := callbacks.NewHandlerHelper(). +// ChatModel(&model.CallbackHandler{ +// OnStart: func(ctx context.Context, info *RunInfo, input *model.CallbackInput) context.Context { +// log.Printf("Model started: %s, messages: %d", info.Name, len(input.Messages)) +// return ctx +// }, +// }). +// Prompt(&prompt.CallbackHandler{ +// OnEnd: func(ctx context.Context, info *RunInfo, output *prompt.CallbackOutput) context.Context { +// log.Printf("Prompt completed") +// return ctx +// }, +// }). +// Handler() +// +// # Passing State Within a Handler +// +// The ctx returned by one timing is passed to the next timing of the SAME +// handler, enabling OnStart→OnEnd state transfer via context.WithValue: +// +// NewHandlerBuilder(). +// OnStartFn(func(ctx context.Context, info *RunInfo, _ CallbackInput) context.Context { +// return context.WithValue(ctx, startTimeKey{}, time.Now()) +// }). +// OnEndFn(func(ctx context.Context, info *RunInfo, _ CallbackOutput) context.Context { +// start := ctx.Value(startTimeKey{}).(time.Time) +// log.Printf("duration: %v", time.Since(start)) +// return ctx +// }).Build() +// +// Between DIFFERENT handlers there is no guaranteed execution order and no +// context chain. To share state between handlers, store it in a +// concurrency-safe variable in the outermost context instead. +// +// # Common Pitfalls +// +// - Stream copies must be closed: when N handlers register for a streaming +// timing, the stream is copied N+1 times (one per handler + one for +// downstream). If any handler's copy is not closed, the original stream +// cannot be freed and the entire pipeline leaks. +// +// - Do NOT mutate Input/Output: all downstream nodes and handlers share the +// same pointer. Mutations cause data races in concurrent graph execution. +// +// - AppendGlobalHandlers is NOT thread-safe: call only during initialization, +// never concurrently with graph execution. +// +// - Stream errors are invisible to OnError: errors that occur while a +// consumer reads from a StreamReader are not routed through OnError. +// +// - RunInfo may be nil: always nil-check before dereferencing in handlers, +// especially when a component is used standalone outside a graph without +// InitCallbacks being called. +package callbacks diff --git a/callbacks/handler_builder.go b/callbacks/handler_builder.go new file mode 100644 index 0000000..438b808 --- /dev/null +++ b/callbacks/handler_builder.go @@ -0,0 +1,163 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// HandlerBuilder constructs a [Handler] by registering callback functions for +// individual timings. Only set the timings you care about; the built handler +// implements [TimingChecker] and returns false for unregistered timings, so +// the framework skips those timings with no overhead. +// +// The input/output values are untyped (CallbackInput / CallbackOutput). To +// work with a specific component's payload, use the component package's +// ConvCallbackInput / ConvCallbackOutput helpers inside your function. For a +// higher-level API that dispatches by component type automatically, see +// utils/callbacks.NewHandlerHelper. +// +// Example: +// +// handler := callbacks.NewHandlerBuilder(). +// OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { +// mi := model.ConvCallbackInput(input) +// if mi != nil { +// log.Printf("[%s] model start: %d messages", info.Name, len(mi.Messages)) +// } +// return ctx +// }). +// OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { +// mo := model.ConvCallbackOutput(output) +// if mo != nil && mo.Message.ResponseMeta != nil { +// log.Printf("[%s] tokens: %d", info.Name, mo.Message.ResponseMeta.Usage.TotalTokens) +// } +// return ctx +// }). +// Build() +type HandlerBuilder struct { + onStartFn func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context + onEndFn func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context + onErrorFn func(ctx context.Context, info *RunInfo, err error) context.Context + onStartWithStreamInputFn func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context + onEndWithStreamOutputFn func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context +} + +type handlerImpl struct { + HandlerBuilder +} + +func (hb *handlerImpl) OnStart(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { + return hb.onStartFn(ctx, info, input) +} + +func (hb *handlerImpl) OnEnd(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context { + return hb.onEndFn(ctx, info, output) +} + +func (hb *handlerImpl) OnError(ctx context.Context, info *RunInfo, err error) context.Context { + return hb.onErrorFn(ctx, info, err) +} + +func (hb *handlerImpl) OnStartWithStreamInput(ctx context.Context, info *RunInfo, + input *schema.StreamReader[CallbackInput]) context.Context { + + return hb.onStartWithStreamInputFn(ctx, info, input) +} + +func (hb *handlerImpl) OnEndWithStreamOutput(ctx context.Context, info *RunInfo, + output *schema.StreamReader[CallbackOutput]) context.Context { + + return hb.onEndWithStreamOutputFn(ctx, info, output) +} + +func (hb *handlerImpl) Needed(_ context.Context, _ *RunInfo, timing CallbackTiming) bool { + switch timing { + case TimingOnStart: + return hb.onStartFn != nil + case TimingOnEnd: + return hb.onEndFn != nil + case TimingOnError: + return hb.onErrorFn != nil + case TimingOnStartWithStreamInput: + return hb.onStartWithStreamInputFn != nil + case TimingOnEndWithStreamOutput: + return hb.onEndWithStreamOutputFn != nil + default: + return false + } +} + +// NewHandlerBuilder creates and returns a new HandlerBuilder instance. +// HandlerBuilder is used to construct a Handler with custom callback functions +func NewHandlerBuilder() *HandlerBuilder { + return &HandlerBuilder{} +} + +// OnStartFn sets the handler for the start timing. +func (hb *HandlerBuilder) OnStartFn( + fn func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context) *HandlerBuilder { + + hb.onStartFn = fn + return hb +} + +// OnEndFn sets the handler for the end timing. +func (hb *HandlerBuilder) OnEndFn( + fn func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context) *HandlerBuilder { + + hb.onEndFn = fn + return hb +} + +// OnErrorFn sets the handler for the error timing. +func (hb *HandlerBuilder) OnErrorFn( + fn func(ctx context.Context, info *RunInfo, err error) context.Context) *HandlerBuilder { + + hb.onErrorFn = fn + return hb +} + +// OnStartWithStreamInputFn sets the callback invoked when a component receives +// streaming input. The handler receives a [*schema.StreamReader] that is a +// private copy; it MUST close the reader after consuming it to avoid goroutine +// and memory leaks. +func (hb *HandlerBuilder) OnStartWithStreamInputFn( + fn func(ctx context.Context, info *RunInfo, input *schema.StreamReader[CallbackInput]) context.Context) *HandlerBuilder { + + hb.onStartWithStreamInputFn = fn + return hb +} + +// OnEndWithStreamOutputFn sets the callback invoked when a component produces +// streaming output. Like OnStartWithStreamInputFn, the handler receives a +// private copy of the stream and MUST close it after reading to prevent +// goroutine and memory leaks. This is the right place to implement streaming +// token-usage accounting or streaming log capture. +func (hb *HandlerBuilder) OnEndWithStreamOutputFn( + fn func(ctx context.Context, info *RunInfo, output *schema.StreamReader[CallbackOutput]) context.Context) *HandlerBuilder { + + hb.onEndWithStreamOutputFn = fn + return hb +} + +// Build returns a Handler with the functions set in the builder. +func (hb *HandlerBuilder) Build() Handler { + return &handlerImpl{*hb} +} diff --git a/callbacks/interface.go b/callbacks/interface.go new file mode 100644 index 0000000..b80554b --- /dev/null +++ b/callbacks/interface.go @@ -0,0 +1,145 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "github.com/cloudwego/eino/internal/callbacks" +) + +// RunInfo describes the entity that triggered a callback. Always nil-check +// before dereferencing — a component that calls OnStart without first calling +// EnsureRunInfo or InitCallbacks will leave RunInfo absent in the context. +// +// Fields: +// - Name: business-meaningful name specified by the user. For nodes in a +// graph this is the node name (compose.WithNodeName). For standalone +// components it must be set explicitly via [InitCallbacks] or +// [ReuseHandlers]; it is empty string if not set. +// - Type: implementation identity, e.g. "OpenAI". Set by the component via +// [components.Typer]; falls back to reflection (struct/func name) if the +// interface is not implemented. Empty for Graph itself. +// - Component: category constant, e.g. components.ComponentOfChatModel. +// Fixed value "Lambda" for lambdas, "Graph"/"Chain"/"Workflow" for graphs. +// Use this to branch on component kind without caring about implementation. +// +// Handlers should filter using RunInfo rather than assuming a fixed execution +// order — there is no guaranteed ordering between different Handlers. +type RunInfo = callbacks.RunInfo + +// CallbackInput is the value passed to OnStart and OnStartWithStreamInput +// handlers. The concrete type is defined by the component — for example, +// ChatModel callbacks carry *model.CallbackInput. Use the component package's +// ConvCallbackInput helper (e.g. model.ConvCallbackInput) to cast safely; it +// returns nil if the type does not match, so you can ignore irrelevant +// component types: +// +// modelInput := model.ConvCallbackInput(in) +// if modelInput == nil { +// return ctx // not a model invocation, skip +// } +// log.Printf("prompt: %v", modelInput.Messages) +type CallbackInput = callbacks.CallbackInput + +// CallbackOutput is the value passed to OnEnd and OnEndWithStreamOutput +// handlers. Like CallbackInput, the concrete type is component-defined. +// Use the component package's ConvCallbackOutput helper to cast safely. +type CallbackOutput = callbacks.CallbackOutput + +// Handler is the unified callback handler interface. Implement all five +// methods (OnStart, OnEnd, OnError, OnStartWithStreamInput, +// OnEndWithStreamOutput) or use [NewHandlerBuilder] to set only the timings +// you care about. +// +// Each method receives the context returned by the previous timing of the +// SAME handler, which lets a single handler pass state between its OnStart +// and OnEnd calls via context.WithValue. There is NO guaranteed execution +// order between DIFFERENT handlers, and the context chain does not flow +// from one handler to the next — do not rely on handler ordering. +// +// Implement [TimingChecker] (the Needed method) on your handler so the +// framework can skip timings you have not registered; this avoids unnecessary +// stream copies and goroutine allocations on every component invocation. +// +// Stream handlers (OnStartWithStreamInput, OnEndWithStreamOutput) receive a +// [*schema.StreamReader] that has already been copied; they MUST close their +// copy after reading. If any handler's copy is not closed, the original stream +// cannot be freed, causing a goroutine/memory leak for the entire pipeline. +// +// Important: do NOT mutate the Input or Output values. All downstream nodes +// and handlers share the same pointer (direct assignment, not a deep copy). +// Mutations cause data races in concurrent graph execution. +type Handler = callbacks.Handler + +// InitCallbackHandlers sets the global callback handlers. +// It should be called BEFORE any callback handler by user. +// It's useful when you want to inject some basic callbacks to all nodes. +// Deprecated: Use AppendGlobalHandlers instead. +func InitCallbackHandlers(handlers []Handler) { + callbacks.GlobalHandlers = handlers +} + +// AppendGlobalHandlers appends handlers to the process-wide list of callback +// handlers. Global handlers run before per-invocation handlers provided via +// compose.WithCallbacks, giving them higher priority for instrumentation that +// must observe every component invocation (e.g. distributed tracing, metrics). +// +// This function is NOT thread-safe. Call it once during program initialization +// (e.g. in main or TestMain), before any graph executions begin. +// Calling it concurrently with ongoing graph executions leads to data races. +func AppendGlobalHandlers(handlers ...Handler) { + callbacks.GlobalHandlers = append(callbacks.GlobalHandlers, handlers...) +} + +// CallbackTiming enumerates the lifecycle moments at which a callback handler +// is invoked. Implement [TimingChecker] on your handler and return false for +// timings you do not handle, so the framework skips the overhead of stream +// copying and goroutine spawning for those timings. +type CallbackTiming = callbacks.CallbackTiming + +// Callback timing constants. +const ( + // TimingOnStart fires just before the component begins processing. + // Receives a fully-formed input value (non-streaming). + TimingOnStart CallbackTiming = iota + // TimingOnEnd fires after the component returns a result successfully. + // Receives the output value. Only fires on success — not on error. + TimingOnEnd + // TimingOnError fires when the component returns a non-nil error. + // Stream errors (mid-stream panics) are NOT reported here; they surface + // as errors inside the stream reader. + TimingOnError + // TimingOnStartWithStreamInput fires when the component receives a + // streaming input (Collect / Transform paradigms). The handler receives a + // copy of the input stream and must close it after reading. + TimingOnStartWithStreamInput + // TimingOnEndWithStreamOutput fires after the component returns a + // streaming output (Stream / Transform paradigms). The handler receives a + // copy of the output stream and must close it after reading. This is + // typically where you implement streaming metrics or logging. + TimingOnEndWithStreamOutput +) + +// TimingChecker is an optional interface for [Handler] implementations. +// When a handler implements Needed, the framework calls it before each +// component invocation to decide whether to set up callback infrastructure +// (stream copying, goroutine allocation) for that timing. Returning false +// avoids unnecessary overhead. +// +// Handlers built with [NewHandlerBuilder] or +// utils/callbacks.NewHandlerHelper automatically implement TimingChecker +// based on which callback functions were set. +type TimingChecker = callbacks.TimingChecker diff --git a/callbacks/interface_test.go b/callbacks/interface_test.go new file mode 100644 index 0000000..6bb7e0f --- /dev/null +++ b/callbacks/interface_test.go @@ -0,0 +1,56 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/callbacks" +) + +func TestAppendGlobalHandlers(t *testing.T) { + // Clear global handlers before test + callbacks.GlobalHandlers = nil + + // Create test handlers + handler1 := NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *RunInfo, input CallbackInput) context.Context { + return ctx + }).Build() + handler2 := NewHandlerBuilder(). + OnEndFn(func(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context { + return ctx + }).Build() + + // Test appending first handler + AppendGlobalHandlers(handler1) + assert.Equal(t, 1, len(callbacks.GlobalHandlers)) + assert.Contains(t, callbacks.GlobalHandlers, handler1) + + // Test appending second handler + AppendGlobalHandlers(handler2) + assert.Equal(t, 2, len(callbacks.GlobalHandlers)) + assert.Contains(t, callbacks.GlobalHandlers, handler1) + assert.Contains(t, callbacks.GlobalHandlers, handler2) + + // Test appending nil handler + AppendGlobalHandlers([]Handler{}...) + assert.Equal(t, 2, len(callbacks.GlobalHandlers)) +} diff --git a/components/document/callback_extra_loader.go b/components/document/callback_extra_loader.go new file mode 100644 index 0000000..ee1a167 --- /dev/null +++ b/components/document/callback_extra_loader.go @@ -0,0 +1,71 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package document + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// LoaderCallbackInput is the input for the loader callback. +type LoaderCallbackInput struct { + // Source is the source of the documents. + Source Source + + // Extra is the extra information for the callback. + Extra map[string]any +} + +// LoaderCallbackOutput is the output for the loader callback. +type LoaderCallbackOutput struct { + // Source is the source of the documents. + Source Source + + // Docs is the documents to be loaded. + Docs []*schema.Document + + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvLoaderCallbackInput converts the callback input to the loader callback input. +func ConvLoaderCallbackInput(src callbacks.CallbackInput) *LoaderCallbackInput { + switch t := src.(type) { + case *LoaderCallbackInput: + return t + case Source: + return &LoaderCallbackInput{ + Source: t, + } + default: + return nil + } +} + +// ConvLoaderCallbackOutput converts the callback output to the loader callback output. +func ConvLoaderCallbackOutput(src callbacks.CallbackOutput) *LoaderCallbackOutput { + switch t := src.(type) { + case *LoaderCallbackOutput: + return t + case []*schema.Document: + return &LoaderCallbackOutput{ + Docs: t, + } + default: + return nil + } +} diff --git a/components/document/callback_extra_transformer.go b/components/document/callback_extra_transformer.go new file mode 100644 index 0000000..cd98f4a --- /dev/null +++ b/components/document/callback_extra_transformer.go @@ -0,0 +1,68 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package document + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// TransformerCallbackInput is the input for the transformer callback. +type TransformerCallbackInput struct { + // Input is the input documents. + Input []*schema.Document + + // Extra is the extra information for the callback. + Extra map[string]any +} + +// TransformerCallbackOutput is the output for the transformer callback. +type TransformerCallbackOutput struct { + // Output is the output documents. + Output []*schema.Document + + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvTransformerCallbackInput converts the callback input to the transformer callback input. +func ConvTransformerCallbackInput(src callbacks.CallbackInput) *TransformerCallbackInput { + switch t := src.(type) { + case *TransformerCallbackInput: + return t + case []*schema.Document: + return &TransformerCallbackInput{ + Input: t, + } + default: + return nil + } +} + +// ConvTransformerCallbackOutput converts the callback output to the transformer callback output. +func ConvTransformerCallbackOutput(src callbacks.CallbackOutput) *TransformerCallbackOutput { + switch t := src.(type) { + case *TransformerCallbackOutput: + return t + case []*schema.Document: + return &TransformerCallbackOutput{ + Output: t, + } + default: + return nil + } +} diff --git a/components/document/doc.go b/components/document/doc.go new file mode 100644 index 0000000..2e91f5b --- /dev/null +++ b/components/document/doc.go @@ -0,0 +1,41 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package document defines the Loader and Transformer component interfaces +// for ingesting and processing documents in an eino pipeline. +// +// # Components +// +// - [Loader]: reads raw content from an external source (file, URL, S3, …) +// and returns [schema.Document] values. Parsing is typically delegated to +// a [parser.Parser] configured on the loader. +// - [Transformer]: takes a slice of [schema.Document] values and transforms +// them — splitting, filtering, merging, re-ranking, etc. +// +// Concrete implementations live in eino-ext: +// +// github.com/cloudwego/eino-ext/components/document/ +// +// # Document Metadata +// +// [schema.Document].MetaData is the primary mechanism for carrying contextual +// information (source URI, scores, chunk indices, embeddings) through the +// pipeline. Transformers should preserve existing metadata and merge rather +// than replace when adding their own keys. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/document_loader_guide/ +// See https://www.cloudwego.io/docs/eino/core_modules/components/document_transformer_guide/ +package document diff --git a/components/document/interface.go b/components/document/interface.go new file mode 100644 index 0000000..6be8d73 --- /dev/null +++ b/components/document/interface.go @@ -0,0 +1,55 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package document + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// Source identifies the external location of a document. +// URI can be a local file path or a remote URL reachable by the loader. +type Source struct { + URI string +} + +//go:generate mockgen -destination ../../internal/mock/components/document/document_mock.go --package document -source interface.go + +// Loader reads raw content from an external source and returns it as a slice +// of [schema.Document] values. +// +// The Source.URI may be a local file path or a remote URL. The loader is +// responsible for fetching the raw bytes; actual format parsing is typically +// delegated to a [parser.Parser] configured on the loader via +// [WithParserOptions]. +// +// Document metadata ([schema.Document].MetaData) should be populated with at +// least the source URI so that downstream nodes can trace document provenance. +type Loader interface { + Load(ctx context.Context, src Source, opts ...LoaderOption) ([]*schema.Document, error) +} + +// Transformer converts a slice of [schema.Document] values into another slice, +// applying operations such as splitting, filtering, merging, or re-ranking. +// +// Implementations should preserve existing MetaData keys and merge rather than +// replace when adding their own metadata. Downstream nodes (e.g. Indexer, +// Retriever) may depend on metadata set by earlier pipeline stages. +type Transformer interface { + Transform(ctx context.Context, src []*schema.Document, opts ...TransformerOption) ([]*schema.Document, error) +} diff --git a/components/document/option.go b/components/document/option.go new file mode 100644 index 0000000..777b2cc --- /dev/null +++ b/components/document/option.go @@ -0,0 +1,167 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package document + +import "github.com/cloudwego/eino/components/document/parser" + +// LoaderOptions configures document loaders, including parser options. +type LoaderOptions struct { + ParserOptions []parser.Option +} + +// LoaderOption defines call option for Loader component, which is part of the component interface signature. +// Each Loader implementation could define its own options struct and option funcs within its own package, +// then wrap the impl specific option funcs into this type, before passing to Load. +type LoaderOption struct { + apply func(opts *LoaderOptions) + + implSpecificOptFn any +} + +// WrapLoaderImplSpecificOptFn wraps the impl specific option functions into LoaderOption type. +// T: the type of the impl specific options struct. +// Loader implementations are required to use this function to convert its own option functions into the unified LoaderOption type. +// For example, if the Loader impl defines its own options struct: +// +// type customOptions struct { +// conf string +// } +// +// Then the impl needs to provide an option function as such: +// +// func WithConf(conf string) Option { +// return WrapLoaderImplSpecificOptFn(func(o *customOptions) { +// o.conf = conf +// } +// } +func WrapLoaderImplSpecificOptFn[T any](optFn func(*T)) LoaderOption { + return LoaderOption{ + implSpecificOptFn: optFn, + } +} + +// GetLoaderImplSpecificOptions provides Loader author the ability to extract their own custom options from the unified LoaderOption type. +// T: the type of the impl specific options struct. +// This function should be used within the Loader implementation's Load function. +// It is recommended to provide a base T as the first argument, within which the Loader author can provide default values for the impl specific options. +// eg. +// +// myOption := &MyOption{ +// Field1: "default_value", +// } +// myOption := loader.GetLoaderImplSpecificOptions(myOption, opts...) +func GetLoaderImplSpecificOptions[T any](base *T, opts ...LoaderOption) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + s, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + s(base) + } + } + } + + return base +} + +// GetLoaderCommonOptions extract loader Options from Option list, optionally providing a base Options with default values. +func GetLoaderCommonOptions(base *LoaderOptions, opts ...LoaderOption) *LoaderOptions { + if base == nil { + base = &LoaderOptions{} + } + + for i := range opts { + opt := opts[i] + if opt.apply != nil { + opt.apply(base) + } + } + + return base +} + +// WithParserOptions attaches parser options to a loader request. +func WithParserOptions(opts ...parser.Option) LoaderOption { + return LoaderOption{ + apply: func(o *LoaderOptions) { + o.ParserOptions = opts + }, + } +} + +// TransformerOption defines call option for Transformer component, which is part of the component interface signature. +// Each Transformer implementation could define its own options struct and option funcs within its own package, +// then wrap the impl specific option funcs into this type, before passing to Transform. +type TransformerOption struct { + implSpecificOptFn any +} + +// WrapTransformerImplSpecificOptFn wraps the impl specific option functions into TransformerOption type. +// T: the type of the impl specific options struct. +// Transformer implementations are required to use this function to convert its own option functions into the unified TransformerOption type. +// For example, if the Transformer impl defines its own options struct: +// +// type customOptions struct { +// conf string +// } +// +// Then the impl needs to provide an option function as such: +// +// func WithConf(conf string) TransformerOption { +// return WrapTransformerImplSpecificOptFn(func(o *customOptions) { +// o.conf = conf +// } +// } +// +// . +func WrapTransformerImplSpecificOptFn[T any](optFn func(*T)) TransformerOption { + return TransformerOption{ + implSpecificOptFn: optFn, + } +} + +// GetTransformerImplSpecificOptions provides Transformer author the ability to extract their own custom options from the unified TransformerOption type. +// T: the type of the impl specific options struct. +// This function should be used within the Transformer implementation's Transform function. +// It is recommended to provide a base T as the first argument, within which the Transformer author can provide default values for the impl specific options. +// eg. +// +// myOption := &MyOption{ +// Field1: "default_value", +// } +// myOption := transformer.GetTransformerImplSpecificOptions(myOption, opts...) +func GetTransformerImplSpecificOptions[T any](base *T, opts ...TransformerOption) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + s, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + s(base) + } + } + } + + return base +} diff --git a/components/document/option_test.go b/components/document/option_test.go new file mode 100644 index 0000000..2bff574 --- /dev/null +++ b/components/document/option_test.go @@ -0,0 +1,78 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package document + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" + + "github.com/cloudwego/eino/components/document/parser" +) + +func TestImplSpecificOpts(t *testing.T) { + type implSpecificOptions struct { + conf string + index int + } + + withConf := func(conf string) func(o *implSpecificOptions) { + return func(o *implSpecificOptions) { + o.conf = conf + } + } + + withIndex := func(index int) func(o *implSpecificOptions) { + return func(o *implSpecificOptions) { + o.index = index + } + } + + convey.Convey("TestLoaderImplSpecificOpts", t, func() { + documentOption1 := WrapLoaderImplSpecificOptFn(withConf("test_conf")) + documentOption2 := WrapLoaderImplSpecificOptFn(withIndex(1)) + + implSpecificOpts := GetLoaderImplSpecificOptions(&implSpecificOptions{}, documentOption1, documentOption2) + + convey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{ + conf: "test_conf", + index: 1, + }) + }) + convey.Convey("TestTransformerImplSpecificOpts", t, func() { + documentOption1 := WrapTransformerImplSpecificOptFn(withConf("test_conf")) + documentOption2 := WrapTransformerImplSpecificOptFn(withIndex(1)) + + implSpecificOpts := GetTransformerImplSpecificOptions(&implSpecificOptions{}, documentOption1, documentOption2) + + convey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{ + conf: "test_conf", + index: 1, + }) + }) +} + +func TestCommonOptions(t *testing.T) { + convey.Convey("TestCommonOptions", t, func() { + o := &LoaderOptions{ParserOptions: []parser.Option{{}}} + o1 := GetLoaderCommonOptions(o) + convey.So(len(o1.ParserOptions), convey.ShouldEqual, 1) + + o2 := GetLoaderCommonOptions(o, WithParserOptions(parser.Option{}, parser.Option{})) + convey.So(len(o2.ParserOptions), convey.ShouldEqual, 2) + }) +} diff --git a/components/document/parser/doc.go b/components/document/parser/doc.go new file mode 100644 index 0000000..541989f --- /dev/null +++ b/components/document/parser/doc.go @@ -0,0 +1,48 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package parser defines the Parser interface for converting raw byte streams +// into [schema.Document] values. +// +// # Overview +// +// A Parser is not a standalone pipeline component — it is used inside a +// [document.Loader] to handle format-specific decoding. The loader fetches +// raw bytes; the parser converts them into documents. +// +// # Built-in Implementations +// +// - TextParser: treats the entire reader as plain text, one document per call +// - ExtParser: selects a parser by file extension (from [Options.URI]), with +// a configurable fallback for unknown extensions +// +// Use ExtParser when you want format-agnostic loading: pass the source URI +// via [WithURI] and ExtParser picks the right sub-parser automatically. +// +// # Reader Contract +// +// The [io.Reader] passed to [Parser.Parse] is consumed during the call — +// it cannot be read again. Loaders must not reuse the same reader across +// multiple Parse calls. +// +// # Metadata Propagation +// +// Use [WithExtraMeta] to attach key-value pairs that are merged into every +// document's MetaData. This is the standard way to tag documents with source +// information (URI, content type, etc.) at parse time. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/document_loader_guide/document_parser_interface_guide/ +package parser diff --git a/components/document/parser/ext_parser.go b/components/document/parser/ext_parser.go new file mode 100644 index 0000000..7d1f2f7 --- /dev/null +++ b/components/document/parser/ext_parser.go @@ -0,0 +1,130 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parser + +import ( + "context" + "errors" + "io" + "path/filepath" + + "github.com/cloudwego/eino/schema" +) + +// ExtParserConfig defines the configuration for the ExtParser. +type ExtParserConfig struct { + // ext -> parser. + // eg: map[string]Parser{ + // ".pdf": &PDFParser{}, + // ".md": &MarkdownParser{}, + // } + Parsers map[string]Parser + + // Fallback parser to use when no other parser is found. + // Default is TextParser if not set. + FallbackParser Parser +} + +// ExtParser is a parser that uses the file extension to determine which parser to use. +// You can register your own parsers by calling RegisterParser. +// Default parser is TextParser. +// Note: +// +// parse 时，是通过 filepath.Ext(uri) 的方式找到对应的 parser，因此使用时需要： +// ① 必须使用 parser.WithURI 在请求时传入 URI +// ② URI 必须能通过 filepath.Ext 来解析出符合预期的 ext +// +// eg: +// +// pdf, _ := os.Open("./testdata/test.pdf") +// docs, err := ExtParser.Parse(ctx, pdf, parser.WithURI("./testdata/test.pdf")) +type ExtParser struct { + parsers map[string]Parser + + fallbackParser Parser +} + +// NewExtParser creates a new ExtParser. +func NewExtParser(ctx context.Context, conf *ExtParserConfig) (*ExtParser, error) { + if conf == nil { + conf = &ExtParserConfig{} + } + + p := &ExtParser{ + parsers: conf.Parsers, + fallbackParser: conf.FallbackParser, + } + + if p.fallbackParser == nil { + p.fallbackParser = TextParser{} + } + + if p.parsers == nil { + p.parsers = make(map[string]Parser) + } + + return p, nil +} + +// GetParsers returns a copy of the registered parsers. +// It is safe to modify the returned parsers. +func (p *ExtParser) GetParsers() map[string]Parser { + res := make(map[string]Parser, len(p.parsers)) + for k, v := range p.parsers { + res[k] = v + } + + return res +} + +// Parse parses the given reader and returns a list of documents. +func (p *ExtParser) Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) { + opt := GetCommonOptions(&Options{}, opts...) + + ext := filepath.Ext(opt.URI) + + parser, ok := p.parsers[ext] + + if !ok { + parser = p.fallbackParser + } + + if parser == nil { + return nil, errors.New("no parser found for extension " + ext) + } + + docs, err := parser.Parse(ctx, reader, opts...) + if err != nil { + return nil, err + } + + for _, doc := range docs { + if doc == nil { + continue + } + + if doc.MetaData == nil { + doc.MetaData = make(map[string]any) + } + + for k, v := range opt.ExtraMeta { + doc.MetaData[k] = v + } + } + + return docs, nil +} diff --git a/components/document/parser/interface.go b/components/document/parser/interface.go new file mode 100644 index 0000000..e2e9082 --- /dev/null +++ b/components/document/parser/interface.go @@ -0,0 +1,36 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parser + +import ( + "context" + "io" + + "github.com/cloudwego/eino/schema" +) + +// Parser converts raw content from an [io.Reader] into [schema.Document] values. +// +// Parse may return multiple documents from a single reader (e.g. a PDF with +// per-page splitting). The reader is consumed during Parse and must not be +// reused. +// +// Parsers are typically not called directly — they are configured on a +// [document.Loader] and invoked via [document.WithParserOptions]. +type Parser interface { + Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) +} diff --git a/components/document/parser/option.go b/components/document/parser/option.go new file mode 100644 index 0000000..ee9706b --- /dev/null +++ b/components/document/parser/option.go @@ -0,0 +1,116 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parser + +// Options configures the document parser with source URI and extra metadata. +type Options struct { + // uri of source. + URI string + + // extra metadata will merge to each document. + ExtraMeta map[string]any +} + +// Option defines call option for Parser component, which is part of the component interface signature. +// Each Parser implementation could define its own options struct and option funcs within its own package, +// then wrap the impl specific option funcs into this type, before passing to Transform. +type Option struct { + apply func(opts *Options) + + implSpecificOptFn any +} + +// WithURI specifies the source URI of the document. +// It will be used as to select parser in ExtParser. +func WithURI(uri string) Option { + return Option{ + apply: func(opts *Options) { + opts.URI = uri + }, + } +} + +// WithExtraMeta attaches extra metadata to the parsed document. +func WithExtraMeta(meta map[string]any) Option { + return Option{ + apply: func(opts *Options) { + opts.ExtraMeta = meta + }, + } +} + +// GetCommonOptions extract parser Options from Option list, optionally providing a base Options with default values. +func GetCommonOptions(base *Options, opts ...Option) *Options { + if base == nil { + base = &Options{} + } + + for i := range opts { + opt := opts[i] + if opt.apply != nil { + opt.apply(base) + } + } + + return base +} + +// WrapImplSpecificOptFn wraps the impl specific option functions into Option type. +// T: the type of the impl specific options struct. +// Parser implementations are required to use this function to convert its own option functions into the unified Option type. +// For example, if the Parser impl defines its own options struct: +// +// type customOptions struct { +// conf string +// } +// +// Then the impl needs to provide an option function as such: +// +// func WithConf(conf string) Option { +// return WrapImplSpecificOptFn(func(o *customOptions) { +// o.conf = conf +// } +// } +// +// . +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions provides Parser author the ability to extract their own custom options from the unified Option type. +// T: the type of the impl specific options struct. +// This function should be used within the Parser implementation's Transform function. +// It is recommended to provide a base T as the first argument, within which the Parser author can provide default values for the impl specific options. +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + s, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + s(base) + } + } + } + + return base +} diff --git a/components/document/parser/option_test.go b/components/document/parser/option_test.go new file mode 100644 index 0000000..d2ed699 --- /dev/null +++ b/components/document/parser/option_test.go @@ -0,0 +1,54 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parser + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" +) + +func TestImplSpecificOpts(t *testing.T) { + type implSpecificOptions struct { + conf string + index int + } + + withConf := func(conf string) func(o *implSpecificOptions) { + return func(o *implSpecificOptions) { + o.conf = conf + } + } + + withIndex := func(index int) func(o *implSpecificOptions) { + return func(o *implSpecificOptions) { + o.index = index + } + } + + convey.Convey("TestImplSpecificOpts", t, func() { + parserOption1 := WrapImplSpecificOptFn(withConf("test_conf")) + parserOption2 := WrapImplSpecificOptFn(withIndex(1)) + + implSpecificOpts := GetImplSpecificOptions(&implSpecificOptions{}, parserOption1, parserOption2) + + convey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{ + conf: "test_conf", + index: 1, + }) + }) +} diff --git a/components/document/parser/parser_test.go b/components/document/parser/parser_test.go new file mode 100644 index 0000000..2324e3c --- /dev/null +++ b/components/document/parser/parser_test.go @@ -0,0 +1,118 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parser + +import ( + "context" + "io" + "os" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +type ParserForTest struct { + mock func() ([]*schema.Document, error) +} + +func (p *ParserForTest) Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) { + return p.mock() +} + +func TestParser(t *testing.T) { + ctx := context.Background() + + t.Run("Test default parser", func(t *testing.T) { + conf := &ExtParserConfig{} + + p, err := NewExtParser(ctx, conf) + if err != nil { + t.Fatal(err) + } + + f, err := os.Open("testdata/test.md") + if err != nil { + t.Fatal(err) + } + defer f.Close() + + docs, err := p.Parse(ctx, f, WithURI("testdata/test.md")) + if err != nil { + t.Fatal(err) + } + + assert.Equal(t, 1, len(docs)) + assert.Equal(t, "# Title\nhello world", docs[0].Content) + }) + + t.Run("test types", func(t *testing.T) { + mockParser := &ParserForTest{ + mock: func() ([]*schema.Document, error) { + return []*schema.Document{ + { + Content: "hello world", + MetaData: map[string]any{ + "type": "text", + }, + }, + }, nil + }, + } + + conf := &ExtParserConfig{ + Parsers: map[string]Parser{ + ".md": mockParser, + }, + } + + p, err := NewExtParser(ctx, conf) + if err != nil { + t.Fatal(err) + } + + f, err := os.Open("testdata/test.md") + if err != nil { + t.Fatal(err) + } + defer f.Close() + + docs, err := p.Parse(ctx, f, WithURI("x/test.md")) + if err != nil { + t.Fatal(err) + } + + assert.Equal(t, 1, len(docs)) + assert.Equal(t, "hello world", docs[0].Content) + assert.Equal(t, "text", docs[0].MetaData["type"]) + }) + + t.Run("test get parsers", func(t *testing.T) { + p, err := NewExtParser(ctx, &ExtParserConfig{ + Parsers: map[string]Parser{ + ".md": &TextParser{}, + }, + }) + if err != nil { + t.Fatal(err) + } + + ps := p.GetParsers() + assert.Equal(t, 1, len(ps)) + }) +} diff --git a/components/document/parser/testdata/test.md b/components/document/parser/testdata/test.md new file mode 100644 index 0000000..fb76356 --- /dev/null +++ b/components/document/parser/testdata/test.md @@ -0,0 +1,2 @@ +# Title +hello world \ No newline at end of file diff --git a/components/document/parser/text_parser.go b/components/document/parser/text_parser.go new file mode 100644 index 0000000..3cd956e --- /dev/null +++ b/components/document/parser/text_parser.go @@ -0,0 +1,60 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parser + +import ( + "context" + "io" + + "github.com/cloudwego/eino/schema" +) + +const ( + // MetaKeySource is the metadata key storing the document's source URI. + MetaKeySource = "_source" +) + +// TextParser is a simple parser that reads the text from a reader and returns a single document. +// eg: +// +// docs, err := TextParser.Parse(ctx, strings.NewReader("hello world")) +// fmt.Println(docs[0].Content) // "hello world" +type TextParser struct{} + +// Parse reads the text from a reader and returns a single document. +func (dp TextParser) Parse(ctx context.Context, reader io.Reader, opts ...Option) ([]*schema.Document, error) { + data, err := io.ReadAll(reader) + if err != nil { + return nil, err + } + + opt := GetCommonOptions(&Options{}, opts...) + + meta := make(map[string]any) + meta[MetaKeySource] = opt.URI + + for k, v := range opt.ExtraMeta { + meta[k] = v + } + + doc := &schema.Document{ + Content: string(data), + MetaData: meta, + } + + return []*schema.Document{doc}, nil +} diff --git a/components/embedding/callback_extra.go b/components/embedding/callback_extra.go new file mode 100644 index 0000000..cab5151 --- /dev/null +++ b/components/embedding/callback_extra.go @@ -0,0 +1,97 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package embedding + +import ( + "github.com/cloudwego/eino/callbacks" +) + +// TokenUsage is the token usage for the embedding. +type TokenUsage struct { + // PromptTokens is the number of prompt tokens. + PromptTokens int + // CompletionTokens is the number of completion tokens. + CompletionTokens int + // TotalTokens is the total number of tokens. + TotalTokens int +} + +// Config is the config for the embedding. +type Config struct { + // Model is the model name. + Model string + // EncodingFormat is the encoding format. + EncodingFormat string +} + +// ComponentExtra is the extra information for the embedding. +type ComponentExtra struct { + // Config is the config for the embedding. + Config *Config + // TokenUsage is the token usage for the embedding. + TokenUsage *TokenUsage +} + +// CallbackInput is the input for the embedding callback. +type CallbackInput struct { + // Texts is the texts to be embedded. + Texts []string + // Config is the config for the embedding. + Config *Config + // Extra is the extra information for the callback. + Extra map[string]any +} + +// CallbackOutput is the output for the embedding callback. +type CallbackOutput struct { + // Embeddings is the embeddings. + Embeddings [][]float64 + // Config is the config for creating the embedding. + Config *Config + // TokenUsage is the token usage for the embedding. + TokenUsage *TokenUsage + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvCallbackInput converts the callback input to the embedding callback input. +func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput { + switch t := src.(type) { + case *CallbackInput: + return t + case []string: + return &CallbackInput{ + Texts: t, + } + default: + return nil + } +} + +// ConvCallbackOutput converts the callback output to the embedding callback output. +func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput { + switch t := src.(type) { + case *CallbackOutput: + return t + case [][]float64: + return &CallbackOutput{ + Embeddings: t, + } + default: + return nil + } +} diff --git a/components/embedding/callback_extra_test.go b/components/embedding/callback_extra_test.go new file mode 100644 index 0000000..90e5cb9 --- /dev/null +++ b/components/embedding/callback_extra_test.go @@ -0,0 +1,33 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package embedding + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestConvEmbedding(t *testing.T) { + assert.NotNil(t, ConvCallbackInput(&CallbackInput{})) + assert.NotNil(t, ConvCallbackInput([]string{})) + assert.Nil(t, ConvCallbackInput("asd")) + + assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{})) + assert.NotNil(t, ConvCallbackOutput([][]float64{})) + assert.Nil(t, ConvCallbackOutput("asd")) +} diff --git a/components/embedding/doc.go b/components/embedding/doc.go new file mode 100644 index 0000000..628cd96 --- /dev/null +++ b/components/embedding/doc.go @@ -0,0 +1,44 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package embedding defines the Embedder component interface for converting +// text into vector representations. +// +// # Overview +// +// An Embedder converts a batch of strings into dense float vectors. Semantically +// similar texts produce vectors that are close in the vector space, making +// embeddings the backbone of semantic search, RAG pipelines, and clustering. +// +// Concrete implementations (OpenAI, Ark, Ollama, …) live in eino-ext: +// +// github.com/cloudwego/eino-ext/components/embedding/ +// +// # Output Format +// +// [Embedder.EmbedStrings] returns `[][]float64` where: +// - outer index corresponds to the input text at the same position +// - inner slice is the embedding vector; its length (dimensions) is fixed by +// the model and is the same for every text +// +// # Consistency Requirement +// +// The same model must be used for both indexing and retrieval. Mixing models +// produces vectors in different spaces — similarity scores become meaningless +// and semantic search breaks silently. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/embedding_guide/ +package embedding diff --git a/components/embedding/interface.go b/components/embedding/interface.go new file mode 100644 index 0000000..dc0cbe7 --- /dev/null +++ b/components/embedding/interface.go @@ -0,0 +1,39 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package embedding + +import "context" + +// Embedder converts a batch of strings into dense vector representations. +// +// EmbedStrings returns one vector per input text, in the same order. The +// vector length (dimensions) is fixed by the underlying model and identical +// for every text in the batch. +// +// The returned [][]float64 maps as: +// +// embeddings[i] → vector for texts[i] +// len(embeddings[i]) → model's embedding dimension (e.g. 1536 for ada-002) +// +// Both [Indexer] and [Retriever] use an Embedder to convert documents and +// queries into vectors. They must share the exact same model — mismatched +// dimensions or model families break semantic similarity. +// +//go:generate mockgen -destination ../../internal/mock/components/embedding/Embedding_mock.go --package embedding -source interface.go +type Embedder interface { + EmbedStrings(ctx context.Context, texts []string, opts ...Option) ([][]float64, error) // invoke +} diff --git a/components/embedding/option.go b/components/embedding/option.go new file mode 100644 index 0000000..b8d7fe9 --- /dev/null +++ b/components/embedding/option.go @@ -0,0 +1,102 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package embedding + +// Options is the options for the embedding. +type Options struct { + // Model is the model name for the embedding. + Model *string +} + +// Option is a call-time option for an Embedder. +type Option struct { + apply func(opts *Options) + + implSpecificOptFn any +} + +// WithModel is the option to set the model for the embedding. +func WithModel(model string) Option { + return Option{ + apply: func(opts *Options) { + opts.Model = &model + }, + } +} + +// GetCommonOptions extract embedding Options from Option list, optionally providing a base Options with default values. +// eg. +// +// defaultModelName := "default_model" +// embeddingOption := &embedding.Options{ +// Model: &defaultModelName, +// } +// embeddingOption := embedding.GetCommonOptions(embeddingOption, opts...) +func GetCommonOptions(base *Options, opts ...Option) *Options { + if base == nil { + base = &Options{} + } + + for i := range opts { + opt := opts[i] + if opt.apply != nil { + opt.apply(base) + } + } + + return base +} + +// WrapImplSpecificOptFn wraps an implementation-specific option function so it +// can be passed alongside standard options. For use by Embedder implementors: +// +// func WithMyParam(v string) embedding.Option { +// return embedding.WrapImplSpecificOptFn(func(o *MyOptions) { +// o.MyParam = v +// }) +// } +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions extracts implementation-specific options from opts, +// merging them onto base. Call alongside [GetCommonOptions] inside EmbedStrings: +// +// func (e *MyEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) { +// common := embedding.GetCommonOptions(nil, opts...) +// mine := embedding.GetImplSpecificOptions(&MyOptions{}, opts...) +// // use common.Model, mine.MyParam, etc. +// } +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} diff --git a/components/embedding/option_test.go b/components/embedding/option_test.go new file mode 100644 index 0000000..63c71af --- /dev/null +++ b/components/embedding/option_test.go @@ -0,0 +1,30 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package embedding + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestOptions(t *testing.T) { + defaultModel := "default_model" + opts := GetCommonOptions(&Options{Model: &defaultModel}, WithModel("test_model")) + assert.NotNil(t, opts.Model) + assert.Equal(t, *opts.Model, "test_model") +} diff --git a/components/indexer/callback_extra.go b/components/indexer/callback_extra.go new file mode 100644 index 0000000..88de567 --- /dev/null +++ b/components/indexer/callback_extra.go @@ -0,0 +1,66 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package indexer + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// CallbackInput is the input for the indexer callback. +type CallbackInput struct { + // Docs is the documents to be indexed. + Docs []*schema.Document + // Extra is the extra information for the callback. + Extra map[string]any +} + +// CallbackOutput is the output for the indexer callback. +type CallbackOutput struct { + // IDs is the ids of the indexed documents returned by the indexer. + IDs []string + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvCallbackInput converts the callback input to the indexer callback input. +func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput { + switch t := src.(type) { + case *CallbackInput: + return t + case []*schema.Document: + return &CallbackInput{ + Docs: t, + } + default: + return nil + } +} + +// ConvCallbackOutput converts the callback output to the indexer callback output. +func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput { + switch t := src.(type) { + case *CallbackOutput: + return t + case []string: + return &CallbackOutput{ + IDs: t, + } + default: + return nil + } +} diff --git a/components/indexer/callback_extra_test.go b/components/indexer/callback_extra_test.go new file mode 100644 index 0000000..2ac79be --- /dev/null +++ b/components/indexer/callback_extra_test.go @@ -0,0 +1,35 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package indexer + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestConvIndexer(t *testing.T) { + assert.NotNil(t, ConvCallbackInput(&CallbackInput{})) + assert.NotNil(t, ConvCallbackInput([]*schema.Document{})) + assert.Nil(t, ConvCallbackInput("asd")) + + assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{})) + assert.NotNil(t, ConvCallbackOutput([]string{})) + assert.Nil(t, ConvCallbackOutput("asd")) +} diff --git a/components/indexer/doc.go b/components/indexer/doc.go new file mode 100644 index 0000000..968ff9e --- /dev/null +++ b/components/indexer/doc.go @@ -0,0 +1,39 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package indexer defines the Indexer component interface for storing documents +// and their vector representations in a backend store. +// +// # Overview +// +// An Indexer is the write path of a RAG pipeline. It takes [schema.Document] +// values, optionally generates vector embeddings, and persists them in a +// backend (vector DB, search engine, etc.) for later retrieval. +// +// Concrete implementations (VikingDB, Milvus, Elasticsearch, …) live in +// eino-ext: +// +// github.com/cloudwego/eino-ext/components/indexer/ +// +// # Vector Dimension Consistency +// +// When using the [Options.Embedding] option, the embedding model must be +// identical to the one used by the paired [retriever.Retriever]. Mismatched +// models produce vectors in different spaces — queries will not match stored +// documents. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/indexer_guide/ +package indexer diff --git a/components/indexer/interface.go b/components/indexer/interface.go new file mode 100644 index 0000000..59154fd --- /dev/null +++ b/components/indexer/interface.go @@ -0,0 +1,41 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package indexer + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// Indexer stores documents (and optionally their vector embeddings) in a +// backend for later retrieval. +// +// Store accepts a batch of [schema.Document] values and returns the IDs +// assigned to them by the backend. When [Options.Embedding] is provided, +// the implementation generates vectors before storing — the same embedder +// must be used by the paired [retriever.Retriever]. +// +// Use [Options.Index] to choose a backend index at call time, and +// [Options.SubIndexes] to write documents into logical sub-partitions within +// the same store. +// +//go:generate mockgen -destination ../../internal/mock/components/indexer/indexer_mock.go --package indexer -source interface.go +type Indexer interface { + // Store stores the documents and returns their assigned IDs. + Store(ctx context.Context, docs []*schema.Document, opts ...Option) (ids []string, err error) // invoke +} diff --git a/components/indexer/option.go b/components/indexer/option.go new file mode 100644 index 0000000..9d095ad --- /dev/null +++ b/components/indexer/option.go @@ -0,0 +1,113 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package indexer + +import "github.com/cloudwego/eino/components/embedding" + +// Options is the options for the indexer. +type Options struct { + // Index is the index for the indexer, index in different indexers may be different. + Index *string + // SubIndexes is the sub indexes to be indexed. + SubIndexes []string + // Embedding is the embedding component. + Embedding embedding.Embedder +} + +// WithIndex wraps the index option. +func WithIndex(index string) Option { + return Option{ + apply: func(opts *Options) { + opts.Index = &index + }, + } +} + +// WithSubIndexes is the option to set the sub indexes for the indexer. +func WithSubIndexes(subIndexes []string) Option { + return Option{ + apply: func(opts *Options) { + opts.SubIndexes = subIndexes + }, + } +} + +// WithEmbedding is the option to set the embedder for the indexer, which convert document to embeddings. +func WithEmbedding(emb embedding.Embedder) Option { + return Option{ + apply: func(opts *Options) { + opts.Embedding = emb + }, + } +} + +// Option is a call-time option for an Indexer. +type Option struct { + apply func(opts *Options) + + implSpecificOptFn any +} + +// GetCommonOptions extracts standard [Options] from opts, merging onto base. +// Implementors must call this inside Store: +// +// func (idx *MyIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) ([]string, error) { +// options := indexer.GetCommonOptions(nil, opts...) +// // use options.Embedding to generate vectors before storage +// } +func GetCommonOptions(base *Options, opts ...Option) *Options { + if base == nil { + base = &Options{} + } + + for i := range opts { + opt := opts[i] + if opt.apply != nil { + opt.apply(base) + } + } + + return base +} + +// WrapImplSpecificOptFn wraps an implementation-specific option function so it +// can be passed alongside standard options. For use by Indexer implementors. +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions extracts implementation-specific options from opts, +// merging onto base. Call alongside [GetCommonOptions] inside Store. +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} diff --git a/components/indexer/option_test.go b/components/indexer/option_test.go new file mode 100644 index 0000000..1ef37b5 --- /dev/null +++ b/components/indexer/option_test.go @@ -0,0 +1,48 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package indexer + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" + + "github.com/cloudwego/eino/internal/mock/components/embedding" +) + +func TestOptions(t *testing.T) { + convey.Convey("test options", t, func() { + var ( + index = "index" + subIndexes = []string{"index_1", "index_2"} + e = &embedding.MockEmbedder{} + ) + + opts := GetCommonOptions( + &Options{}, + WithIndex(index), + WithSubIndexes(subIndexes), + WithEmbedding(e), + ) + + convey.So(opts, convey.ShouldResemble, &Options{ + Index: &index, + SubIndexes: subIndexes, + Embedding: e, + }) + }) +} diff --git a/components/model/agentic_callback_extra.go b/components/model/agentic_callback_extra.go new file mode 100644 index 0000000..9a769cf --- /dev/null +++ b/components/model/agentic_callback_extra.go @@ -0,0 +1,94 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// AgenticConfig is the config for the agentic model. +type AgenticConfig struct { + // Model is the model name. + Model string + // MaxTokens is the max number of output tokens, if reached the max tokens, the model will stop generating. + MaxTokens int + // Temperature is the temperature, which controls the randomness of the agentic model. + Temperature float32 + // TopP is the top p, which controls the diversity of the agentic model. + TopP float32 +} + +// AgenticCallbackInput is the input for the agentic model callback. +type AgenticCallbackInput struct { + // Messages is the agentic messages to be sent to the agentic model. + Messages []*schema.AgenticMessage + // Tools is the tools to be used in the agentic model. + Tools []*schema.ToolInfo + // Config is the config for the agentic model. + Config *AgenticConfig + // Extra is the extra information for the callback. + Extra map[string]any +} + +// AgenticCallbackOutput is the output for the agentic model callback. +type AgenticCallbackOutput struct { + // Message is the agentic message generated by the agentic model. + Message *schema.AgenticMessage + // Config is the config for the agentic model. + Config *AgenticConfig + // TokenUsage is the token usage of this request. + TokenUsage *TokenUsage + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvAgenticCallbackInput converts the callback input to the agentic model callback input. +func ConvAgenticCallbackInput(src callbacks.CallbackInput) *AgenticCallbackInput { + switch t := src.(type) { + case *AgenticCallbackInput: + // when callback is triggered within component implementation, + // the input is usually already a typed *model.AgenticCallbackInput + return t + case []*schema.AgenticMessage: + // when callback is injected by graph node, not the component implementation itself, + // the input is the input of Agentic Model interface, which is []*schema.AgenticMessage + return &AgenticCallbackInput{ + Messages: t, + } + default: + return nil + } +} + +// ConvAgenticCallbackOutput converts the callback output to the agentic model callback output. +func ConvAgenticCallbackOutput(src callbacks.CallbackOutput) *AgenticCallbackOutput { + switch t := src.(type) { + case *AgenticCallbackOutput: + // when callback is triggered within component implementation, + // the output is usually already a typed *model.AgenticCallbackOutput + return t + case *schema.AgenticMessage: + // when callback is injected by graph node, not the component implementation itself, + // the output is the output of Agentic Model interface, which is *schema.AgenticMessage + return &AgenticCallbackOutput{ + Message: t, + } + default: + return nil + } +} diff --git a/components/model/agentic_callback_extra_test.go b/components/model/agentic_callback_extra_test.go new file mode 100644 index 0000000..9373674 --- /dev/null +++ b/components/model/agentic_callback_extra_test.go @@ -0,0 +1,35 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestConvAgenticModel(t *testing.T) { + assert.NotNil(t, ConvAgenticCallbackInput(&AgenticCallbackInput{})) + assert.NotNil(t, ConvAgenticCallbackInput([]*schema.AgenticMessage{})) + assert.Nil(t, ConvAgenticCallbackInput("asd")) + + assert.NotNil(t, ConvAgenticCallbackOutput(&AgenticCallbackOutput{})) + assert.NotNil(t, ConvAgenticCallbackOutput(&schema.AgenticMessage{})) + assert.Nil(t, ConvAgenticCallbackOutput("asd")) +} diff --git a/components/model/callback_extra.go b/components/model/callback_extra.go new file mode 100644 index 0000000..8591c43 --- /dev/null +++ b/components/model/callback_extra.go @@ -0,0 +1,117 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// TokenUsage is the token usage for the model. +type TokenUsage struct { + // PromptTokens is the number of prompt tokens, including all the input tokens of this request. + PromptTokens int + // PromptTokenDetails is a breakdown of the prompt tokens. + PromptTokenDetails PromptTokenDetails + // CompletionTokens is the number of completion tokens. + CompletionTokens int + // TotalTokens is the total number of tokens. + TotalTokens int + // CompletionTokensDetails is breakdown of completion tokens. + CompletionTokensDetails CompletionTokensDetails `json:"completion_token_details"` +} + +type CompletionTokensDetails struct { + // ReasoningTokens tokens generated by the model for reasoning. + // This is currently supported by OpenAI, Gemini, ARK and Qwen chat models. + // For other models, this field will be 0. + ReasoningTokens int `json:"reasoning_tokens,omitempty"` +} + +// PromptTokenDetails provides a breakdown of prompt token usage. +type PromptTokenDetails struct { + // Cached tokens present in the prompt. + CachedTokens int +} + +// Config is the config for the model. +type Config struct { + // Model is the model name. + Model string + // MaxTokens is the max number of tokens, if reached the max tokens, the model will stop generating, and mostly return an finish reason of "length". + MaxTokens int + // Temperature is the temperature, which controls the randomness of the model. + Temperature float32 + // TopP is the top p, which controls the diversity of the model. + TopP float32 + // Stop is the stop words, which controls the stopping condition of the model. + Stop []string +} + +// CallbackInput is the input for the model callback. +type CallbackInput struct { + // Messages is the messages to be sent to the model. + Messages []*schema.Message + // Tools is the tools to be used in the model. + Tools []*schema.ToolInfo + // ToolChoice is the tool choice, which controls the tool to be used in the model. + ToolChoice *schema.ToolChoice + // Config is the config for the model. + Config *Config + // Extra is the extra information for the callback. + Extra map[string]any +} + +// CallbackOutput is the output for the model callback. +type CallbackOutput struct { + // Message is the message generated by the model. + Message *schema.Message + // Config is the config for the model. + Config *Config + // TokenUsage is the token usage of this request. + TokenUsage *TokenUsage + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvCallbackInput converts the callback input to the model callback input. +func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput { + switch t := src.(type) { + case *CallbackInput: // when callback is triggered within component implementation, the input is usually already a typed *model.CallbackInput + return t + case []*schema.Message: // when callback is injected by graph node, not the component implementation itself, the input is the input of Chat Model interface, which is []*schema.Message + return &CallbackInput{ + Messages: t, + } + default: + return nil + } +} + +// ConvCallbackOutput converts the callback output to the model callback output. +func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput { + switch t := src.(type) { + case *CallbackOutput: // when callback is triggered within component implementation, the output is usually already a typed *model.CallbackOutput + return t + case *schema.Message: // when callback is injected by graph node, not the component implementation itself, the output is the output of Chat Model interface, which is *schema.Message + return &CallbackOutput{ + Message: t, + } + default: + return nil + } +} diff --git a/components/model/callback_extra_test.go b/components/model/callback_extra_test.go new file mode 100644 index 0000000..2fe8443 --- /dev/null +++ b/components/model/callback_extra_test.go @@ -0,0 +1,35 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestConvModel(t *testing.T) { + assert.NotNil(t, ConvCallbackInput(&CallbackInput{})) + assert.NotNil(t, ConvCallbackInput([]*schema.Message{})) + assert.Nil(t, ConvCallbackInput("asd")) + + assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{})) + assert.NotNil(t, ConvCallbackOutput(&schema.Message{})) + assert.Nil(t, ConvCallbackOutput("asd")) +} diff --git a/components/model/doc.go b/components/model/doc.go new file mode 100644 index 0000000..4c91d85 --- /dev/null +++ b/components/model/doc.go @@ -0,0 +1,59 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package model defines the ChatModel component interface for interacting with +// large language models (LLMs). +// +// # Overview +// +// A ChatModel takes a slice of [schema.Message] as input and returns a response +// message — either in full ([BaseChatModel.Generate]) or incrementally as a +// stream ([BaseChatModel.Stream]). It is the most fundamental building block in +// an eino pipeline: every application that talks to an LLM goes through this +// interface. +// +// Concrete implementations (OpenAI, Ark, Ollama, …) live in eino-ext: +// +// github.com/cloudwego/eino-ext/components/model/ +// +// # Interface Hierarchy +// +// BaseChatModel — Generate + Stream (all implementations) +// ├── ToolCallingChatModel — preferred; WithTools returns a new instance (concurrency-safe) +// └── ChatModel — deprecated; BindTools mutates state (avoid in new code) +// +// # Choosing Generate vs Stream +// +// Use [BaseChatModel.Generate] when the full response is needed before +// proceeding (e.g. structured extraction, classification). +// Use [BaseChatModel.Stream] when output should be forwarded to the caller +// incrementally (e.g. chat UI, long-form generation). Always close the +// [schema.StreamReader] returned by Stream — failing to do so leaks the +// underlying connection: +// +// reader, err := model.Stream(ctx, messages) +// if err != nil { ... } +// defer reader.Close() +// +// # Implementing a ChatModel +// +// Implementations must call [GetCommonOptions] to extract standard options and +// [GetImplSpecificOptions] to extract their own options from the Option list. +// Expose implementation-specific options via [WrapImplSpecificOptFn]. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/chat_model_guide/ +// for the full component guide. +package model diff --git a/components/model/interface.go b/components/model/interface.go new file mode 100644 index 0000000..d326e66 --- /dev/null +++ b/components/model/interface.go @@ -0,0 +1,109 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// messageType is the sealed type constraint for message types used in BaseModel. +// Only *schema.Message and *schema.AgenticMessage satisfy this constraint. +type messageType interface { + *schema.Message | *schema.AgenticMessage +} + +// BaseModel is the generic base model interface parameterized by message type M. +// It exposes two modes of interaction: +// - [BaseModel.Generate]: blocks until the model returns a complete response. +// - [BaseModel.Stream]: returns a [schema.StreamReader] that yields message +// chunks incrementally as the model generates them. +type BaseModel[M messageType] interface { + Generate(ctx context.Context, input []M, opts ...Option) (M, error) + Stream(ctx context.Context, input []M, opts ...Option) (*schema.StreamReader[M], error) +} + +// BaseChatModel is a backward-compatible type alias for BaseModel specialized +// with *schema.Message. All existing code using model.BaseChatModel continues +// to work without modification. +// +// It exposes two modes of interaction: +// - [BaseChatModel.Generate]: blocks until the model returns a complete response. +// - [BaseChatModel.Stream]: returns a [schema.StreamReader] that yields message +// chunks incrementally as the model generates them. +// +// The input is a slice of [schema.Message] representing the conversation history. +// Messages carry a role (system, user, assistant, tool) and support multimodal +// content (text, images, audio, video). Tool messages must include a ToolCallID +// that correlates them with a prior assistant tool-call message. +// +// Stream usage — the caller is responsible for closing the reader: +// +// reader, err := m.Stream(ctx, messages) +// if err != nil { ... } +// defer reader.Close() +// for { +// chunk, err := reader.Recv() +// if errors.Is(err, io.EOF) { break } +// if err != nil { ... } +// // handle chunk +// } +// +// Note: a [schema.StreamReader] can only be read once. If multiple consumers +// need the stream, it must be copied before reading. +// +//go:generate mockgen -destination ../../internal/mock/components/model/ChatModel_mock.go --package model github.com/cloudwego/eino/components/model BaseChatModel,ChatModel,ToolCallingChatModel +type BaseChatModel = BaseModel[*schema.Message] + +// Deprecated: Use [ToolCallingChatModel] instead. +// +// ChatModel extends [BaseChatModel] with tool binding via [ChatModel.BindTools]. +// BindTools mutates the instance in place, which causes a race condition when +// the same instance is used concurrently: one goroutine's tool list can +// overwrite another's. Prefer [ToolCallingChatModel.WithTools], which returns +// a new immutable instance and is safe for concurrent use. +type ChatModel interface { + BaseChatModel + + // BindTools bind tools to the model. + // BindTools before requesting ChatModel generally. + // notice the non-atomic problem of BindTools and Generate. + BindTools(tools []*schema.ToolInfo) error +} + +// ToolCallingChatModel extends [BaseChatModel] with safe tool binding. +// +// Unlike the deprecated [ChatModel.BindTools], [ToolCallingChatModel.WithTools] +// does not mutate the receiver — it returns a new instance with the given tools +// attached. This makes it safe to share a base model instance across goroutines +// and derive per-request variants with different tool sets: +// +// base, _ := openai.NewChatModel(ctx, cfg) // shared, no tools +// withSearch, _ := base.WithTools([]*schema.ToolInfo{searchTool}) +// withCalc, _ := base.WithTools([]*schema.ToolInfo{calcTool}) +type ToolCallingChatModel interface { + BaseChatModel + + WithTools(tools []*schema.ToolInfo) (ToolCallingChatModel, error) +} + +// AgenticModel is a type alias for BaseModel specialized with +// *schema.AgenticMessage. Unlike ToolCallingChatModel, agentic models do NOT +// expose a WithTools method; tools are passed at request time via the +// model.WithTools option, consistent with how ChatModelAgent binds tools. +type AgenticModel = BaseModel[*schema.AgenticMessage] diff --git a/components/model/option.go b/components/model/option.go new file mode 100644 index 0000000..2222e14 --- /dev/null +++ b/components/model/option.go @@ -0,0 +1,255 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import "github.com/cloudwego/eino/schema" + +// Options is the common options for the model. +type Options struct { + // Temperature is the temperature for the model, which controls the randomness of the model. + Temperature *float32 + // Model is the model name. + Model *string + // TopP is the top p for the model, which controls the diversity of the model. + TopP *float32 + // Tools is a list of tools the model may call. + Tools []*schema.ToolInfo + // DeferredTools is a list of tools to be registered with defer_loading=true + // for the model's built-in (server-side) tool search capability. + // These tools are sent to the model API but not loaded into context upfront — + // only their names and descriptions are visible to the model. The model's + // built-in tool search tool searches through them and loads matching ones + // on demand. + DeferredTools []*schema.ToolInfo + + ToolSearchTool *schema.ToolInfo + + // MaxTokens is the max number of tokens, if reached the max tokens, the model will stop generating, and mostly return a finish reason of "length". + MaxTokens *int + // Stop is the stop words for the model, which controls the stopping condition of the model. + Stop []string + + // Options only available for chat model. + + // ToolChoice controls which tool is called by the model. + ToolChoice *schema.ToolChoice + // AllowedToolNames specifies a list of tool names that the model is allowed to call. + // This allows for constraining the model to a specific subset of the available tools. + AllowedToolNames []string + + // Options only available for agentic model. + + // AgenticToolChoice controls how the agentic model calls tools. + AgenticToolChoice *schema.AgenticToolChoice +} + +// Option is a call-time option for a ChatModel. Options are immutable and +// composable: each Option carries either a common-option setter (applied via +// [GetCommonOptions]) or an implementation-specific setter (applied via +// [GetImplSpecificOptions]), never both. +type Option struct { + apply func(opts *Options) + + implSpecificOptFn any +} + +// WithTemperature is the option to set the temperature for the model. +func WithTemperature(temperature float32) Option { + return Option{ + apply: func(opts *Options) { + opts.Temperature = &temperature + }, + } +} + +// WithMaxTokens is the option to set the max tokens for the model. +func WithMaxTokens(maxTokens int) Option { + return Option{ + apply: func(opts *Options) { + opts.MaxTokens = &maxTokens + }, + } +} + +// WithModel is the option to set the model name. +func WithModel(name string) Option { + return Option{ + apply: func(opts *Options) { + opts.Model = &name + }, + } +} + +// WithTopP is the option to set the top p for the model. +func WithTopP(topP float32) Option { + return Option{ + apply: func(opts *Options) { + opts.TopP = &topP + }, + } +} + +// WithStop is the option to set the stop words for the model. +func WithStop(stop []string) Option { + return Option{ + apply: func(opts *Options) { + opts.Stop = stop + }, + } +} + +// WithTools is the option to set tools for the model. +func WithTools(tools []*schema.ToolInfo) Option { + if tools == nil { + tools = []*schema.ToolInfo{} + } + return Option{ + apply: func(opts *Options) { + opts.Tools = tools + }, + } +} + +// WithToolSearchTool is the option to register a tool search tool with the model. +// When set, the model uses this tool to discover and load deferred tools on demand. +// Note: The tool search tool should NOT be included in WithTools. +func WithToolSearchTool(tool *schema.ToolInfo) Option { + return Option{ + apply: func(opts *Options) { + opts.ToolSearchTool = tool + }, + } +} + +// WithDeferredTools is the option to set deferred tools for the model's +// built-in (server-side) tool search. These tools are registered with +// defer_loading=true so the model can discover and load them on demand +// via its native tool search capability. +// Note: Deferred tools should NOT be included in WithTools. +func WithDeferredTools(tools []*schema.ToolInfo) Option { + if tools == nil { + tools = []*schema.ToolInfo{} + } + return Option{ + apply: func(opts *Options) { + opts.DeferredTools = tools + }, + } +} + +// WithToolChoice sets the tool choice for the model. It also allows for providing a list of +// tool names to constrain the model to a specific subset of the available tools. +// Only available for ChatModel. +func WithToolChoice(toolChoice schema.ToolChoice, allowedToolNames ...string) Option { + return Option{ + apply: func(opts *Options) { + opts.ToolChoice = &toolChoice + opts.AllowedToolNames = allowedToolNames + }, + } +} + +// WithAgenticToolChoice is the option to set tool choice for the agentic model. +// Only available for AgenticModel. +func WithAgenticToolChoice(toolChoice *schema.AgenticToolChoice) Option { + return Option{ + apply: func(opts *Options) { + opts.AgenticToolChoice = toolChoice + }, + } +} + +// WrapImplSpecificOptFn is the option to wrap the implementation specific option function. +// WrapImplSpecificOptFn wraps an implementation-specific option function into +// an [Option] so it can be passed alongside standard options. +// +// This is intended for ChatModel implementors, not callers. Define a typed +// setter for your own config struct and expose it as an Option: +// +// // In your implementation package: +// func WithMyParam(v string) model.Option { +// return model.WrapImplSpecificOptFn(func(o *MyOptions) { +// o.MyParam = v +// }) +// } +// +// Callers can then mix standard and implementation-specific options freely: +// +// model.Generate(ctx, msgs, +// model.WithTemperature(0.7), +// mypkg.WithMyParam("value"), +// ) +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetCommonOptions extracts standard [Options] from an Option list, merging +// them onto base. If base is nil, a zero-value Options is used. +// +// Implementors must call this to honour options passed by callers: +// +// func (m *MyModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { +// options := model.GetCommonOptions(&model.Options{Temperature: &m.defaultTemp}, opts...) +// // use options.Temperature, options.Tools, etc. +// } +func GetCommonOptions(base *Options, opts ...Option) *Options { + if base == nil { + base = &Options{} + } + + for i := range opts { + opt := opts[i] + if opt.apply != nil { + opt.apply(base) + } + } + + return base +} + +// GetImplSpecificOptions extracts implementation-specific options from an +// Option list, merging them onto base. If base is nil, a zero-value T is used. +// +// Call this alongside [GetCommonOptions] to support both standard and custom +// options in your implementation: +// +// type MyOptions struct { MyParam string } +// +// func (m *MyModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { +// common := model.GetCommonOptions(nil, opts...) +// myOpts := model.GetImplSpecificOptions(&MyOptions{MyParam: "default"}, opts...) +// // use common.Temperature, myOpts.MyParam, etc. +// } +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} diff --git a/components/model/option_test.go b/components/model/option_test.go new file mode 100644 index 0000000..c836933 --- /dev/null +++ b/components/model/option_test.go @@ -0,0 +1,136 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package model + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" + + "github.com/cloudwego/eino/schema" +) + +func TestOptions(t *testing.T) { + convey.Convey("test options", t, func() { + var ( + modelName = "model" + temperature float32 = 0.9 + maxToken = 5000 + topP float32 = 0.8 + defaultModel = "default_model" + defaultTemperature float32 = 1.0 + defaultMaxTokens = 1000 + defaultTopP float32 = 0.5 + tools = []*schema.ToolInfo{{Name: "asd"}, {Name: "qwe"}} + toolChoice = schema.ToolChoiceForced + allowedToolNames = []string{"web_search"} + ) + + opts := GetCommonOptions( + &Options{ + Model: &defaultModel, + Temperature: &defaultTemperature, + MaxTokens: &defaultMaxTokens, + TopP: &defaultTopP, + }, + WithModel(modelName), + WithTemperature(temperature), + WithMaxTokens(maxToken), + WithTopP(topP), + WithStop([]string{"hello", "bye"}), + WithTools(tools), + WithToolChoice(toolChoice, allowedToolNames...), + ) + + convey.So(opts, convey.ShouldResemble, &Options{ + Model: &modelName, + Temperature: &temperature, + MaxTokens: &maxToken, + TopP: &topP, + Stop: []string{"hello", "bye"}, + Tools: tools, + ToolChoice: &toolChoice, + AllowedToolNames: allowedToolNames, + }) + }) + + convey.Convey("test nil tools option", t, func() { + opts := GetCommonOptions( + &Options{ + Tools: []*schema.ToolInfo{ + {Name: "asd"}, + {Name: "qwe"}, + }, + }, + WithTools(nil), + ) + + convey.So(opts.Tools, convey.ShouldNotBeNil) + convey.So(len(opts.Tools), convey.ShouldEqual, 0) + }) + + convey.Convey("test agentic tool choice option", t, func() { + var ( + toolChoice = schema.ToolChoiceForced + allowedTools = []*schema.AllowedTool{ + {FunctionName: "agentic_tool"}, + } + ) + opts := GetCommonOptions( + nil, + WithAgenticToolChoice(&schema.AgenticToolChoice{ + Type: toolChoice, + Forced: &schema.AgenticForcedToolChoice{ + Tools: allowedTools, + }, + }), + ) + + convey.So(opts.AgenticToolChoice, convey.ShouldNotBeNil) + convey.So(opts.AgenticToolChoice.Type, convey.ShouldEqual, toolChoice) + convey.So(opts.AgenticToolChoice.Forced, convey.ShouldNotBeNil) + convey.So(opts.AgenticToolChoice.Forced.Tools, convey.ShouldResemble, allowedTools) + }) +} + +type implOption struct { + userID int64 + name string +} + +func WithUserID(uid int64) Option { + return WrapImplSpecificOptFn[implOption](func(i *implOption) { + i.userID = uid + }) +} + +func WithName(n string) Option { + return WrapImplSpecificOptFn[implOption](func(i *implOption) { + i.name = n + }) +} + +func TestImplSpecificOption(t *testing.T) { + convey.Convey("impl_specific_option", t, func() { + opt := GetImplSpecificOptions(&implOption{}, WithUserID(101), WithName("Wang")) + + convey.So(opt, convey.ShouldEqual, &implOption{ + userID: 101, + name: "Wang", + }) + }) +} diff --git a/components/prompt/agentic_callback_extra.go b/components/prompt/agentic_callback_extra.go new file mode 100644 index 0000000..315d5a4 --- /dev/null +++ b/components/prompt/agentic_callback_extra.go @@ -0,0 +1,70 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// AgenticCallbackInput is the input for the callback. +type AgenticCallbackInput struct { + // Variables is the variables for the callback. + Variables map[string]any + // Templates is the agentic templates for the callback. + Templates []schema.AgenticMessagesTemplate + // Extra is the extra information for the callback. + Extra map[string]any +} + +// AgenticCallbackOutput is the output for the callback. +type AgenticCallbackOutput struct { + // Result is the agentic result for the callback. + Result []*schema.AgenticMessage + // Templates is the agentic templates for the callback. + Templates []schema.AgenticMessagesTemplate + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvAgenticCallbackInput converts the callback input to the agentic prompt callback input. +func ConvAgenticCallbackInput(src callbacks.CallbackInput) *AgenticCallbackInput { + switch t := src.(type) { + case *AgenticCallbackInput: + return t + case map[string]any: + return &AgenticCallbackInput{ + Variables: t, + } + default: + return nil + } +} + +// ConvAgenticCallbackOutput converts the callback output to the agentic prompt callback output. +func ConvAgenticCallbackOutput(src callbacks.CallbackOutput) *AgenticCallbackOutput { + switch t := src.(type) { + case *AgenticCallbackOutput: + return t + case []*schema.AgenticMessage: + return &AgenticCallbackOutput{ + Result: t, + } + default: + return nil + } +} diff --git a/components/prompt/agentic_callback_extra_test.go b/components/prompt/agentic_callback_extra_test.go new file mode 100644 index 0000000..67982be --- /dev/null +++ b/components/prompt/agentic_callback_extra_test.go @@ -0,0 +1,46 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestConvAgenticPrompt(t *testing.T) { + assert.NotNil(t, ConvAgenticCallbackInput(&AgenticCallbackInput{ + Variables: map[string]any{}, + Templates: []schema.AgenticMessagesTemplate{ + &schema.AgenticMessage{}, + }, + })) + assert.NotNil(t, ConvAgenticCallbackInput(map[string]any{})) + assert.Nil(t, ConvAgenticCallbackInput("asd")) + + assert.NotNil(t, ConvAgenticCallbackOutput(&AgenticCallbackOutput{ + Result: []*schema.AgenticMessage{ + {}, + }, + Templates: []schema.AgenticMessagesTemplate{ + &schema.AgenticMessage{}, + }, + })) + assert.NotNil(t, ConvAgenticCallbackOutput([]*schema.AgenticMessage{})) +} diff --git a/components/prompt/agentic_chat_template.go b/components/prompt/agentic_chat_template.go new file mode 100644 index 0000000..41d2910 --- /dev/null +++ b/components/prompt/agentic_chat_template.go @@ -0,0 +1,84 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "context" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/schema" +) + +// FromAgenticMessages creates a new DefaultAgenticChatTemplate from the given templates and format type. +// eg. +// +// template := prompt.FromAgenticMessages(schema.FString, &schema.AgenticMessage{}) +// // in chain, or graph +// chain := compose.NewChain[map[string]any, []*schema.AgenticMessage]() +// chain.AppendAgenticChatTemplate(template) +func FromAgenticMessages(formatType schema.FormatType, templates ...schema.AgenticMessagesTemplate) *DefaultAgenticChatTemplate { + return &DefaultAgenticChatTemplate{ + templates: templates, + formatType: formatType, + } +} + +type DefaultAgenticChatTemplate struct { + templates []schema.AgenticMessagesTemplate + formatType schema.FormatType +} + +func (t *DefaultAgenticChatTemplate) Format(ctx context.Context, vs map[string]any, opts ...Option) (result []*schema.AgenticMessage, err error) { + ctx = callbacks.EnsureRunInfo(ctx, t.GetType(), components.ComponentOfAgenticPrompt) + ctx = callbacks.OnStart(ctx, &AgenticCallbackInput{ + Variables: vs, + Templates: t.templates, + }) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + } + }() + + result = make([]*schema.AgenticMessage, 0, len(t.templates)) + for _, template := range t.templates { + msgs, err := template.Format(ctx, vs, t.formatType) + if err != nil { + return nil, err + } + + result = append(result, msgs...) + } + + _ = callbacks.OnEnd(ctx, &AgenticCallbackOutput{ + Result: result, + Templates: t.templates, + }) + + return result, nil +} + +// GetType returns the type of the agentic template (DefaultAgentic). +func (t *DefaultAgenticChatTemplate) GetType() string { + return "Default" +} + +// IsCallbacksEnabled checks if the callbacks are enabled for the chat template. +func (t *DefaultAgenticChatTemplate) IsCallbacksEnabled() bool { + return true +} diff --git a/components/prompt/agentic_chat_template_test.go b/components/prompt/agentic_chat_template_test.go new file mode 100644 index 0000000..f47020a --- /dev/null +++ b/components/prompt/agentic_chat_template_test.go @@ -0,0 +1,125 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +type mockAgenticTemplate struct { + err error +} + +func (m *mockAgenticTemplate) Format(ctx context.Context, vs map[string]any, formatType schema.FormatType) ([]*schema.AgenticMessage, error) { + if m.err != nil { + return nil, m.err + } + return []*schema.AgenticMessage{schema.UserAgenticMessage("mocked")}, nil +} + +func TestFromAgenticMessages(t *testing.T) { + t.Run("create template", func(t *testing.T) { + tpl := schema.UserAgenticMessage("hello") + ft := schema.FString + at := FromAgenticMessages(ft, tpl) + + assert.NotNil(t, at) + assert.Equal(t, ft, at.formatType) + assert.Len(t, at.templates, 1) + assert.Same(t, tpl, at.templates[0]) + }) +} + +func TestDefaultAgenticTemplate_GetType(t *testing.T) { + t.Run("get type", func(t *testing.T) { + at := &DefaultAgenticChatTemplate{} + assert.Equal(t, "Default", at.GetType()) + }) +} + +func TestDefaultAgenticTemplate_IsCallbacksEnabled(t *testing.T) { + t.Run("callbacks enabled", func(t *testing.T) { + at := &DefaultAgenticChatTemplate{} + assert.True(t, at.IsCallbacksEnabled()) + }) +} + +func TestDefaultAgenticTemplate_Format(t *testing.T) { + t.Run("success", func(t *testing.T) { + // Mock callback handler + cb := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + assert.Equal(t, "Default", info.Type) + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + assert.Equal(t, "Default", info.Type) + return ctx + }). + OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + assert.Fail(t, "unexpected error callback") + return ctx + }). + Build() + + tpl := schema.UserAgenticMessage("hello {val}") + at := FromAgenticMessages(schema.FString, tpl) + + ctx := context.Background() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{ + Type: "Default", + Component: "agentic_prompt", + }, cb) + + res, err := at.Format(ctx, map[string]any{"val": "world"}) + assert.NoError(t, err) + assert.Len(t, res, 1) + assert.Equal(t, "hello world", res[0].ContentBlocks[0].UserInputText.Text) + }) + + t.Run("template format error", func(t *testing.T) { + mockErr := errors.New("mock error") + mockTpl := &mockAgenticTemplate{err: mockErr} + at := FromAgenticMessages(schema.FString, mockTpl) + + // Mock callback handler to verify OnError + cb := callbacks.NewHandlerBuilder(). + OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + assert.Equal(t, mockErr, err) + return ctx + }). + Build() + + ctx := context.Background() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{ + Type: "Default", + Component: "agentic_prompt", + }, cb) + + res, err := at.Format(ctx, map[string]any{}) + assert.Error(t, err) + assert.Nil(t, res) + assert.Equal(t, mockErr, err) + }) +} diff --git a/components/prompt/callback_extra.go b/components/prompt/callback_extra.go new file mode 100644 index 0000000..324a418 --- /dev/null +++ b/components/prompt/callback_extra.go @@ -0,0 +1,70 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// CallbackInput is the input for the callback. +type CallbackInput struct { + // Variables is the variables for the callback. + Variables map[string]any + // Templates is the templates for the callback. + Templates []schema.MessagesTemplate + // Extra is the extra information for the callback. + Extra map[string]any +} + +// CallbackOutput is the output for the callback. +type CallbackOutput struct { + // Result is the result for the callback. + Result []*schema.Message + // Templates is the templates for the callback. + Templates []schema.MessagesTemplate + // Extra is the extra information for the callback. + Extra map[string]any +} + +// ConvCallbackInput converts the callback input to the prompt callback input. +func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput { + switch t := src.(type) { + case *CallbackInput: + return t + case map[string]any: + return &CallbackInput{ + Variables: t, + } + default: + return nil + } +} + +// ConvCallbackOutput converts the callback output to the prompt callback output. +func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput { + switch t := src.(type) { + case *CallbackOutput: + return t + case []*schema.Message: + return &CallbackOutput{ + Result: t, + } + default: + return nil + } +} diff --git a/components/prompt/callback_extra_test.go b/components/prompt/callback_extra_test.go new file mode 100644 index 0000000..ad8a3c0 --- /dev/null +++ b/components/prompt/callback_extra_test.go @@ -0,0 +1,45 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestConvPrompt(t *testing.T) { + assert.NotNil(t, ConvCallbackInput(&CallbackInput{ + Templates: []schema.MessagesTemplate{ + &schema.Message{}, + }, + })) + assert.NotNil(t, ConvCallbackInput(map[string]any{})) + assert.Nil(t, ConvCallbackInput("asd")) + + assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{ + Result: []*schema.Message{ + {}, + }, + Templates: []schema.MessagesTemplate{ + &schema.Message{}, + }, + })) + assert.NotNil(t, ConvCallbackOutput([]*schema.Message{})) +} diff --git a/components/prompt/chat_template.go b/components/prompt/chat_template.go new file mode 100644 index 0000000..2648058 --- /dev/null +++ b/components/prompt/chat_template.go @@ -0,0 +1,89 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "context" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/schema" +) + +// DefaultChatTemplate is the default chat template implementation. +type DefaultChatTemplate struct { + // templates is the templates for the chat template. + templates []schema.MessagesTemplate + // formatType is the format type for the chat template. + formatType schema.FormatType +} + +// FromMessages creates a new DefaultChatTemplate from the given templates and format type. +// eg. +// +// template := prompt.FromMessages(schema.FString, &schema.Message{Content: "Hello, {name}!"}, &schema.Message{Content: "how are you?"}) +// // in chain, or graph +// chain := compose.NewChain[map[string]any, []*schema.Message]() +// chain.AppendChatTemplate(template) +func FromMessages(formatType schema.FormatType, templates ...schema.MessagesTemplate) *DefaultChatTemplate { + return &DefaultChatTemplate{ + templates: templates, + formatType: formatType, + } +} + +// Format formats the chat template with the given context and variables. +func (t *DefaultChatTemplate) Format(ctx context.Context, + vs map[string]any, _ ...Option) (result []*schema.Message, err error) { + ctx = callbacks.EnsureRunInfo(ctx, t.GetType(), components.ComponentOfPrompt) + ctx = callbacks.OnStart(ctx, &CallbackInput{ + Variables: vs, + Templates: t.templates, + }) + defer func() { + if err != nil { + _ = callbacks.OnError(ctx, err) + } + }() + + result = make([]*schema.Message, 0, len(t.templates)) + for _, template := range t.templates { + msgs, err := template.Format(ctx, vs, t.formatType) + if err != nil { + return nil, err + } + + result = append(result, msgs...) + } + + _ = callbacks.OnEnd(ctx, &CallbackOutput{ + Result: result, + Templates: t.templates, + }) + + return result, nil +} + +// GetType returns the type of the chat template (Default). +func (t *DefaultChatTemplate) GetType() string { + return "Default" +} + +// IsCallbacksEnabled checks if the callbacks are enabled for the chat template. +func (t *DefaultChatTemplate) IsCallbacksEnabled() bool { + return true +} diff --git a/components/prompt/chat_template_test.go b/components/prompt/chat_template_test.go new file mode 100644 index 0000000..25f5ce5 --- /dev/null +++ b/components/prompt/chat_template_test.go @@ -0,0 +1,201 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestFormat(t *testing.T) { + pyFmtTestTemplate := []schema.MessagesTemplate{ + schema.SystemMessage( + "you are a helpful assistant.\n" + + "here is the context: {context}"), + schema.MessagesPlaceholder("chat_history", true), + schema.UserMessage("question: {question}"), + } + jinja2TestTemplate := []schema.MessagesTemplate{ + schema.SystemMessage( + "you are a helpful assistant.\n" + + "here is the context: {{context}}"), + schema.MessagesPlaceholder("chat_history", true), + schema.UserMessage("question: {{question}}"), + } + goFmtTestTemplate := []schema.MessagesTemplate{ + schema.SystemMessage( + "you are a helpful assistant.\n" + + "here is the context: {{.context}}"), + schema.MessagesPlaceholder("chat_history", true), + schema.UserMessage("question: {{.question}}"), + } + testValues := map[string]any{ + "context": "it's beautiful day", + "question": "how is the day today", + "chat_history": []*schema.Message{ + schema.UserMessage("who are you"), + schema.AssistantMessage("I'm a helpful assistant", nil), + }, + } + expected := []*schema.Message{ + schema.SystemMessage( + "you are a helpful assistant.\n" + + "here is the context: it's beautiful day"), + schema.UserMessage("who are you"), + schema.AssistantMessage("I'm a helpful assistant", nil), + schema.UserMessage("question: how is the day today"), + } + + // FString + chatTemplate := FromMessages(schema.FString, pyFmtTestTemplate...) + msgs, err := chatTemplate.Format(context.Background(), testValues) + assert.Nil(t, err) + assert.Equal(t, expected, msgs) + + // Jinja2 + chatTemplate = FromMessages(schema.Jinja2, jinja2TestTemplate...) + msgs, err = chatTemplate.Format(context.Background(), testValues) + assert.Nil(t, err) + assert.Equal(t, expected, msgs) + + // GoTemplate + chatTemplate = FromMessages(schema.GoTemplate, goFmtTestTemplate...) + msgs, err = chatTemplate.Format(context.Background(), testValues) + assert.Nil(t, err) + assert.Equal(t, expected, msgs) +} + +func TestDocumentFormat(t *testing.T) { + docs := []*schema.Document{ + { + ID: "1", + Content: "qwe", + MetaData: map[string]any{ + "hello": 888, + }, + }, + { + ID: "2", + Content: "asd", + MetaData: map[string]any{ + "bye": 111, + }, + }, + } + + template := FromMessages(schema.FString, + schema.SystemMessage("all:{all_docs}\nsingle:{single_doc}"), + ) + + msgs, err := template.Format(context.Background(), map[string]any{ + "all_docs": docs, + "single_doc": docs[0], + }) + + assert.Nil(t, err) + t.Log(msgs) +} + +func TestMultiContentFormat(t *testing.T) { + mtpl := []schema.MessagesTemplate{ + &schema.Message{ + Content: "{a}", + MultiContent: []schema.ChatMessagePart{ + { + Type: schema.ChatMessagePartTypeText, + Text: "{b}", + }, + { + Type: schema.ChatMessagePartTypeImageURL, + ImageURL: &schema.ChatMessageImageURL{ + URL: "{c}", + }, + }, + { + Type: schema.ChatMessagePartTypeAudioURL, + AudioURL: &schema.ChatMessageAudioURL{ + URL: "{d}", + }, + }, + { + Type: schema.ChatMessagePartTypeVideoURL, + VideoURL: &schema.ChatMessageVideoURL{ + URL: "{e}", + }, + }, + { + Type: schema.ChatMessagePartTypeFileURL, + FileURL: &schema.ChatMessageFileURL{ + URL: "{f}", + }, + }, + }, + }, + } + input := map[string]any{ + "a": "content", + "b": "text", + "c": "image url", + "d": "audio url", + "e": "video url", + "f": "file url", + } + expected := []*schema.Message{ + { + Content: "content", + MultiContent: []schema.ChatMessagePart{ + { + Type: schema.ChatMessagePartTypeText, + Text: "text", + }, + { + Type: schema.ChatMessagePartTypeImageURL, + ImageURL: &schema.ChatMessageImageURL{ + URL: "image url", + }, + }, + { + Type: schema.ChatMessagePartTypeAudioURL, + AudioURL: &schema.ChatMessageAudioURL{ + URL: "audio url", + }, + }, + { + Type: schema.ChatMessagePartTypeVideoURL, + VideoURL: &schema.ChatMessageVideoURL{ + URL: "video url", + }, + }, + { + Type: schema.ChatMessagePartTypeFileURL, + FileURL: &schema.ChatMessageFileURL{ + URL: "file url", + }, + }, + }, + }, + } + + tpl := FromMessages(schema.FString, mtpl...) + result, err := tpl.Format(context.Background(), input) + assert.Nil(t, err) + assert.Equal(t, expected, result) +} diff --git a/components/prompt/doc.go b/components/prompt/doc.go new file mode 100644 index 0000000..898aa5c --- /dev/null +++ b/components/prompt/doc.go @@ -0,0 +1,56 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package prompt defines the ChatTemplate component interface for building +// structured message lists from templates and runtime variables. +// +// # Overview +// +// A ChatTemplate takes a variables map and produces a []*schema.Message slice +// ready to pass to a [model.BaseChatModel]. It is typically the first node in +// a pipeline, sitting before the ChatModel. +// +// The built-in [DefaultChatTemplate] supports three template syntaxes: +// - FString: {variable} substitution +// - GoTemplate: Go's text/template with conditionals and loops +// - Jinja2: Jinja2 template syntax +// +// # Construction +// +// Use [FromMessages] to build a template from a list of message templates: +// +// tmpl := prompt.FromMessages(schema.FString, +// schema.SystemMessage("You are a helpful assistant."), +// schema.UserMessage("Answer this: {question}"), +// ) +// msgs, err := tmpl.Format(ctx, map[string]any{"question": "What is eino?"}) +// +// Use [schema.MessagesPlaceholder] to insert a dynamic list of messages +// (e.g. conversation history) at a fixed position in the template: +// +// tmpl := prompt.FromMessages(schema.FString, +// schema.SystemMessage("You are a helpful assistant."), +// schema.MessagesPlaceholder("history", true), +// schema.UserMessage("{question}"), +// ) +// +// # Common Pitfall +// +// Variable mismatches (a key present in the template but missing from the +// variables map) produce a runtime error — there is no compile-time check. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/chat_template_guide/ +package prompt diff --git a/components/prompt/interface.go b/components/prompt/interface.go new file mode 100644 index 0000000..2d5a2cb --- /dev/null +++ b/components/prompt/interface.go @@ -0,0 +1,50 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +var _ ChatTemplate = &DefaultChatTemplate{} +var _ AgenticChatTemplate = &DefaultAgenticChatTemplate{} + +// ChatTemplate formats a variables map into a list of messages for a ChatModel. +// +// Format substitutes the values from vs into the template's message list and +// returns the resulting []*schema.Message. The exact substitution syntax +// (FString, GoTemplate, Jinja2) is determined at construction time. +// +// Variable keys present in the template but absent from vs produce a runtime +// error — there is no compile-time safety. Prefer consistent variable naming +// across templates and callers. +// +// In a Graph or Chain, ChatTemplate typically precedes ChatModel. Use +// compose.WithOutputKey to convert the prior node's output into the map[string]any +// that Format expects. +// +// See [FromMessages] and [schema.MessagesPlaceholder] for construction helpers. +type ChatTemplate interface { + Format(ctx context.Context, vs map[string]any, opts ...Option) ([]*schema.Message, error) +} + +// AgenticChatTemplate formats variables into a list of agentic messages according to a prompt schema. +type AgenticChatTemplate interface { + Format(ctx context.Context, vs map[string]any, opts ...Option) ([]*schema.AgenticMessage, error) +} diff --git a/components/prompt/option.go b/components/prompt/option.go new file mode 100644 index 0000000..3b93c42 --- /dev/null +++ b/components/prompt/option.go @@ -0,0 +1,52 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +// Option is a call-time option for a ChatTemplate. The built-in +// [DefaultChatTemplate] has no common options — this type exists primarily for +// custom ChatTemplate implementations that need per-call configuration. +type Option struct { + implSpecificOptFn any +} + +// WrapImplSpecificOptFn wraps an implementation-specific option function so it +// can be passed alongside any future standard options. For use by custom +// ChatTemplate implementors. +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions extracts the implementation specific options from Option list, optionally providing a base options with default values. +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + s, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + s(base) + } + } + } + + return base +} diff --git a/components/prompt/option_test.go b/components/prompt/option_test.go new file mode 100644 index 0000000..851a013 --- /dev/null +++ b/components/prompt/option_test.go @@ -0,0 +1,51 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package prompt + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" +) + +type implOption struct { + userID int64 + name string +} + +func WithUserID(uid int64) Option { + return WrapImplSpecificOptFn[implOption](func(i *implOption) { + i.userID = uid + }) +} + +func WithName(n string) Option { + return WrapImplSpecificOptFn[implOption](func(i *implOption) { + i.name = n + }) +} + +func TestImplSpecificOption(t *testing.T) { + convey.Convey("impl_specific_option", t, func() { + opt := GetImplSpecificOptions(&implOption{}, WithUserID(101), WithName("Wang")) + + convey.So(opt, convey.ShouldEqual, &implOption{ + userID: 101, + name: "Wang", + }) + }) +} diff --git a/components/retriever/callback_extra.go b/components/retriever/callback_extra.go new file mode 100644 index 0000000..49ea421 --- /dev/null +++ b/components/retriever/callback_extra.go @@ -0,0 +1,74 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package retriever + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// CallbackInput is the input for the retriever callback. +type CallbackInput struct { + // Query is the query for the retriever. + Query string + + // TopK is the top k for the retriever, which means the top number of documents to retrieve. + TopK int + // Filter is the filter for the retriever. + Filter string + // ScoreThreshold is the score threshold for the retriever, eg 0.5 means the score of the document must be greater than 0.5. + ScoreThreshold *float64 + + // Extra is the extra information for the retriever. + Extra map[string]any +} + +// CallbackOutput is the output for the retriever callback. +type CallbackOutput struct { + // Docs is the documents for the retriever. + Docs []*schema.Document + // Extra is the extra information for the retriever. + Extra map[string]any +} + +// ConvCallbackInput converts the callback input to the retriever callback input. +func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput { + switch t := src.(type) { + case *CallbackInput: + return t + case string: + return &CallbackInput{ + Query: t, + } + default: + return nil + } +} + +// ConvCallbackOutput converts the callback output to the retriever callback output. +func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput { + switch t := src.(type) { + case *CallbackOutput: + return t + case []*schema.Document: + return &CallbackOutput{ + Docs: t, + } + default: + return nil + } +} diff --git a/components/retriever/callback_extra_test.go b/components/retriever/callback_extra_test.go new file mode 100644 index 0000000..83cb9b5 --- /dev/null +++ b/components/retriever/callback_extra_test.go @@ -0,0 +1,35 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package retriever + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestConvRetriever(t *testing.T) { + assert.NotNil(t, ConvCallbackInput(&CallbackInput{})) + assert.NotNil(t, ConvCallbackInput("asd")) + assert.Nil(t, ConvCallbackInput([]string{})) + + assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{})) + assert.NotNil(t, ConvCallbackOutput([]*schema.Document{})) + assert.Nil(t, ConvCallbackOutput("asd")) +} diff --git a/components/retriever/doc.go b/components/retriever/doc.go new file mode 100644 index 0000000..243bef3 --- /dev/null +++ b/components/retriever/doc.go @@ -0,0 +1,46 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package retriever defines the Retriever component interface for fetching +// relevant documents from a document store given a query. +// +// # Overview +// +// A Retriever is the read path of a RAG (Retrieval-Augmented Generation) +// pipeline. Given a query string it returns the most relevant [schema.Document] +// values from an underlying store (vector DB, keyword index, etc.). +// +// Concrete implementations (VikingDB, Milvus, Elasticsearch, …) live in +// eino-ext: +// +// github.com/cloudwego/eino-ext/components/retriever/ +// +// # Relationship to Indexer +// +// [Indexer] and Retriever are complementary: +// - Indexer writes documents (and their vectors) to the store +// - Retriever reads them back +// +// When both use an [embedding.Embedder], it must be the same model — vector +// dimensions must match or similarity scores will be meaningless. +// +// # Result Ordering +// +// Results are ordered by relevance score (descending). Scores and other +// backend metadata are available via [schema.Document].MetaData. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/retriever_guide/ +package retriever diff --git a/components/retriever/interface.go b/components/retriever/interface.go new file mode 100644 index 0000000..7b27e48 --- /dev/null +++ b/components/retriever/interface.go @@ -0,0 +1,50 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package retriever + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +//go:generate mockgen -destination ../../internal/mock/components/retriever/retriever_mock.go --package retriever -source interface.go + +// Retriever fetches the most relevant documents from a store for a given query. +// +// Retrieve accepts a natural-language query string and returns matching +// [schema.Document] values ordered by relevance (most relevant first). +// Relevance scores and backend-specific metadata are available in +// [schema.Document].MetaData. +// +// When [Options.Embedding] is set, the implementation converts the query to a +// vector before searching. The embedder must be the same model used at index +// time — see [indexer.Options.Embedding]. +// +// [Options.ScoreThreshold] is a filter, not a sort: documents scoring below +// the threshold are excluded entirely. [Options.TopK] caps the number of +// results returned. +// +// Retrieve can be used standalone or added to a Graph via AddRetrieverNode: +// +// retriever, _ := redis.NewRetriever(ctx, cfg) +// docs, _ := retriever.Retrieve(ctx, "what is eino?", retriever.WithTopK(5)) +// +// graph.AddRetrieverNode("retriever", retriever) +type Retriever interface { + Retrieve(ctx context.Context, query string, opts ...Option) ([]*schema.Document, error) +} diff --git a/components/retriever/option.go b/components/retriever/option.go new file mode 100644 index 0000000..a6bbef4 --- /dev/null +++ b/components/retriever/option.go @@ -0,0 +1,147 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package retriever + +import "github.com/cloudwego/eino/components/embedding" + +// Options is the options for the retriever. +type Options struct { + // Index is the index for the retriever, index in different retriever may be different. + Index *string + // SubIndex is the sub index for the retriever, sub index in different retriever may be different. + SubIndex *string + // TopK is the top k for the retriever, which means the top number of documents to retrieve. + TopK *int + // ScoreThreshold is the score threshold for the retriever, eg 0.5 means the score of the document must be greater than 0.5. + ScoreThreshold *float64 + // Embedding is the embedder for the retriever, which is used to embed the query for retrieval . + Embedding embedding.Embedder + + // DSLInfo carries backend-specific filter/query expressions. The structure and + // semantics are defined by the underlying store implementation. + DSLInfo map[string]any +} + +// WithIndex wraps the index option. +func WithIndex(index string) Option { + return Option{ + apply: func(opts *Options) { + opts.Index = &index + }, + } +} + +// WithSubIndex wraps the sub index option. +func WithSubIndex(subIndex string) Option { + return Option{ + apply: func(opts *Options) { + opts.SubIndex = &subIndex + }, + } +} + +// WithTopK wraps the top k option. +func WithTopK(topK int) Option { + return Option{ + apply: func(opts *Options) { + opts.TopK = &topK + }, + } +} + +// WithScoreThreshold wraps the score threshold option. +func WithScoreThreshold(threshold float64) Option { + return Option{ + apply: func(opts *Options) { + opts.ScoreThreshold = &threshold + }, + } +} + +// WithEmbedding wraps the embedder option. +func WithEmbedding(emb embedding.Embedder) Option { + return Option{ + apply: func(opts *Options) { + opts.Embedding = emb + }, + } +} + +// WithDSLInfo wraps the dsl info option. +func WithDSLInfo(dsl map[string]any) Option { + return Option{ + apply: func(opts *Options) { + opts.DSLInfo = dsl + }, + } +} + +// Option is a call-time option for a Retriever. +type Option struct { + apply func(opts *Options) + + implSpecificOptFn any +} + +// GetCommonOptions extracts standard [Options] from opts, merging onto base. +// Implementors must call this to honour caller-provided options: +// +// func (r *MyRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { +// options := retriever.GetCommonOptions(&retriever.Options{TopK: &r.defaultTopK}, opts...) +// // use options.TopK, options.ScoreThreshold, options.Embedding, etc. +// } +func GetCommonOptions(base *Options, opts ...Option) *Options { + if base == nil { + base = &Options{} + } + + for i := range opts { + if opts[i].apply != nil { + opts[i].apply(base) + } + } + + return base +} + +// WrapImplSpecificOptFn wraps an implementation-specific option function so it +// can be passed alongside standard options. For use by Retriever implementors. +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions extracts implementation-specific options from opts, +// merging onto base. Call alongside [GetCommonOptions] inside Retrieve. +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} diff --git a/components/retriever/option_test.go b/components/retriever/option_test.go new file mode 100644 index 0000000..b332dc7 --- /dev/null +++ b/components/retriever/option_test.go @@ -0,0 +1,60 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package retriever + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" + + "github.com/cloudwego/eino/internal/mock/components/embedding" +) + +func TestOptions(t *testing.T) { + convey.Convey("test options", t, func() { + var ( + index = "index" + topK = 2 + scoreThreshold = 4.0 + subIndex = "sub_index" + dslInfo = map[string]any{"dsl": "dsl"} + e = &embedding.MockEmbedder{} + defaultTopK = 1 + ) + + opts := GetCommonOptions( + &Options{ + TopK: &defaultTopK, + }, + WithIndex(index), + WithTopK(topK), + WithScoreThreshold(scoreThreshold), + WithSubIndex(subIndex), + WithDSLInfo(dslInfo), + WithEmbedding(e), + ) + + convey.So(opts, convey.ShouldResemble, &Options{ + Index: &index, + TopK: &topK, + ScoreThreshold: &scoreThreshold, + SubIndex: &subIndex, + DSLInfo: dslInfo, + Embedding: e, + }) + }) +} diff --git a/components/tool/callback_extra.go b/components/tool/callback_extra.go new file mode 100644 index 0000000..407f043 --- /dev/null +++ b/components/tool/callback_extra.go @@ -0,0 +1,68 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/schema" +) + +// CallbackInput is the input for the tool callback. +type CallbackInput struct { + // ArgumentsInJSON is the arguments in json format for the tool. + ArgumentsInJSON string + // Extra is the extra information for the tool. + Extra map[string]any +} + +// CallbackOutput is the output for the tool callback. +type CallbackOutput struct { + // Response is the response for the tool. + Response string + // ToolOutput is the multimodal output for the tool. Used when the tool returns structured data. + ToolOutput *schema.ToolResult + // Extra is the extra information for the tool. + Extra map[string]any +} + +// ConvCallbackInput converts the callback input to the tool callback input. +func ConvCallbackInput(src callbacks.CallbackInput) *CallbackInput { + switch t := src.(type) { + case *CallbackInput: + return t + case string: + return &CallbackInput{ArgumentsInJSON: t} + case *schema.ToolArgument: + return &CallbackInput{ArgumentsInJSON: t.Text} + default: + return nil + } +} + +// ConvCallbackOutput converts the callback output to the tool callback output. +func ConvCallbackOutput(src callbacks.CallbackOutput) *CallbackOutput { + switch t := src.(type) { + case *CallbackOutput: + return t + case string: + return &CallbackOutput{Response: t} + case *schema.ToolResult: + return &CallbackOutput{ToolOutput: t} + default: + return nil + } +} diff --git a/components/tool/callback_extra_test.go b/components/tool/callback_extra_test.go new file mode 100644 index 0000000..dc51edf --- /dev/null +++ b/components/tool/callback_extra_test.go @@ -0,0 +1,37 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestConvCallbackInput(t *testing.T) { + assert.NotNil(t, ConvCallbackInput(&CallbackInput{})) + assert.NotNil(t, ConvCallbackInput("asd")) + assert.Nil(t, ConvCallbackInput(123)) + assert.Nil(t, ConvCallbackInput(nil)) +} + +func TestConvCallbackOutput(t *testing.T) { + assert.NotNil(t, ConvCallbackOutput(&CallbackOutput{})) + assert.NotNil(t, ConvCallbackOutput("asd")) + assert.Nil(t, ConvCallbackOutput(123)) + assert.Nil(t, ConvCallbackOutput(nil)) +} diff --git a/components/tool/doc.go b/components/tool/doc.go new file mode 100644 index 0000000..a9753d2 --- /dev/null +++ b/components/tool/doc.go @@ -0,0 +1,53 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package tool defines the tool component interfaces that allow language models +// to invoke external capabilities, and helpers for interrupt/resume within tools. +// +// # Interface Hierarchy +// +// BaseTool — Info() only; for passing tool metadata to a ChatModel +// ├── InvokableTool — standard: args as JSON string, returns string +// ├── StreamableTool — standard streaming: args as JSON string, returns StreamReader[string] +// ├── EnhancedInvokableTool — multimodal: args as *schema.ToolArgument, returns *schema.ToolResult +// └── EnhancedStreamableTool— multimodal streaming +// +// # Choosing an Interface +// +// Implement [InvokableTool] for most tools — arguments arrive as a JSON string +// automatically decoded from the model's tool call, and the result is a string +// sent back to the model. +// +// Implement [EnhancedInvokableTool] when the tool needs to return structured +// multimodal content (images, audio, files) rather than plain text. When a +// tool implements both a standard and an enhanced interface, ToolsNode +// prioritises the enhanced interface. +// +// # Creating Tools +// +// The [utils] sub-package provides constructors that eliminate boilerplate: +// - [utils.InferTool] / [utils.InferStreamTool] — infer parameter schema from Go struct tags +// - [utils.NewTool] / [utils.NewStreamTool] — manual ToolInfo + typed function +// +// # Interrupt / Resume +// +// Tools can pause execution and wait for external input using [Interrupt], +// [StatefulInterrupt], and [CompositeInterrupt]. Use [GetInterruptState] and +// [GetResumeContext] inside the tool to distinguish first-run from resumed-run. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/ +// See https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/how_to_create_a_tool/ +package tool diff --git a/components/tool/interface.go b/components/tool/interface.go new file mode 100644 index 0000000..257c1fc --- /dev/null +++ b/components/tool/interface.go @@ -0,0 +1,79 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// BaseTool provides the metadata that a ChatModel uses to decide whether and +// how to call a tool. Info returns a [schema.ToolInfo] containing the tool +// name, description, and parameter JSON schema. +// +// BaseTool alone is sufficient when passing tool definitions to a ChatModel +// via WithTools — the model only needs the schema to generate tool calls. +// To also execute the tool, implement [InvokableTool] or [StreamableTool]. +type BaseTool interface { + Info(ctx context.Context) (*schema.ToolInfo, error) +} + +// InvokableTool is a tool that can be executed by ToolsNode. +// +// InvokableRun receives the model's tool call arguments as a JSON-encoded +// string and returns a plain string result that is sent back to the model as +// a tool message. The framework handles JSON decoding automatically when using +// the [utils.InferTool] or [utils.NewTool] constructors. +type InvokableTool interface { + BaseTool + + // InvokableRun executes the tool with arguments encoded as a JSON string. + InvokableRun(ctx context.Context, argumentsInJSON string, opts ...Option) (string, error) +} + +// StreamableTool is a streaming variant of [InvokableTool]. +// +// StreamableRun returns a [schema.StreamReader] that yields string chunks +// incrementally. The caller (ToolsNode) is responsible for closing the reader. +type StreamableTool interface { + BaseTool + + StreamableRun(ctx context.Context, argumentsInJSON string, opts ...Option) (*schema.StreamReader[string], error) +} + +// EnhancedInvokableTool is a tool that returns structured multimodal results. +// +// Unlike [InvokableTool], arguments arrive as a [schema.ToolArgument] (not a +// raw JSON string) and the result is a [schema.ToolResult] which can carry +// text, images, audio, video, and file content. +// +// When a tool implements both a standard and an enhanced interface, ToolsNode +// prioritises the enhanced interface. +type EnhancedInvokableTool interface { + BaseTool + InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...Option) (*schema.ToolResult, error) +} + +// EnhancedStreamableTool is the streaming variant of [EnhancedInvokableTool]. +// +// It streams [schema.ToolResult] chunks, enabling incremental multimodal +// output. The caller is responsible for closing the returned [schema.StreamReader]. +type EnhancedStreamableTool interface { + BaseTool + StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...Option) (*schema.StreamReader[*schema.ToolResult], error) +} diff --git a/components/tool/interrupt.go b/components/tool/interrupt.go new file mode 100644 index 0000000..1640097 --- /dev/null +++ b/components/tool/interrupt.go @@ -0,0 +1,185 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +import ( + "context" + "errors" + "fmt" + + "github.com/cloudwego/eino/internal/core" +) + +// Interrupt pauses tool execution and signals the orchestration layer to checkpoint. +// The tool can be resumed later with optional data. +// +// Parameters: +// - ctx: The context passed to InvokableRun/StreamableRun +// - info: User-facing information about why the tool is interrupting (e.g., "needs user confirmation") +// +// Returns an error that should be returned from InvokableRun/StreamableRun. +// +// Example: +// +// func (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) { +// if needsConfirmation(args) { +// return "", tool.Interrupt(ctx, "Please confirm this action") +// } +// return doWork(args), nil +// } +func Interrupt(ctx context.Context, info any) error { + is, err := core.Interrupt(ctx, info, nil, nil) + if err != nil { + return err + } + return is +} + +// StatefulInterrupt pauses tool execution with state preservation. +// Use this when the tool has internal state that must be restored on resume. +// +// Parameters: +// - ctx: The context passed to InvokableRun/StreamableRun +// - info: User-facing information about the interrupt +// - state: Internal state to persist (must be gob-serializable) +// +// Example: +// +// func (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) { +// wasInterrupted, hasState, state := tool.GetInterruptState[MyState](ctx) +// if !wasInterrupted { +// // First run - interrupt with state +// return "", tool.StatefulInterrupt(ctx, "processing", MyState{Step: 1}) +// } +// // Resumed - continue from saved state +// return continueFrom(state), nil +// } +func StatefulInterrupt(ctx context.Context, info any, state any) error { + is, err := core.Interrupt(ctx, info, state, nil) + if err != nil { + return err + } + return is +} + +// CompositeInterrupt creates an interrupt that aggregates multiple sub-interrupts. +// Use this when a tool internally executes a graph or other interruptible components. +// +// Parameters: +// - ctx: The context passed to InvokableRun/StreamableRun +// - info: User-facing information for this tool's interrupt +// - state: Internal state to persist for this tool +// - errs: Interrupt errors from sub-components (graphs, other tools, etc.) +// +// Example: +// +// func (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) { +// result, err := t.internalGraph.Invoke(ctx, input) +// if err != nil { +// if _, ok := tool.IsInterruptError(err); ok { +// return "", tool.CompositeInterrupt(ctx, "graph interrupted", myState, err) +// } +// return "", err +// } +// return result, nil +// } +func CompositeInterrupt(ctx context.Context, info any, state any, errs ...error) error { + if len(errs) == 0 { + return StatefulInterrupt(ctx, info, state) + } + + var cErrs []*core.InterruptSignal + for _, err := range errs { + ire := &core.InterruptSignal{} + if errors.As(err, &ire) { + cErrs = append(cErrs, ire) + continue + } + + var provider core.InterruptContextsProvider + if errors.As(err, &provider) { + is := core.FromInterruptContexts(provider.GetInterruptContexts()) + if is != nil { + cErrs = append(cErrs, is) + } + continue + } + + return fmt.Errorf("composite interrupt but one of the sub error is not interrupt error: %w", err) + } + + is, err := core.Interrupt(ctx, info, state, cErrs) + if err != nil { + return err + } + return is +} + +// GetInterruptState checks if the tool was previously interrupted and retrieves saved state. +// +// Returns: +// - wasInterrupted: true if this tool was part of a previous interruption +// - hasState: true if state was saved and successfully cast to type T +// - state: the saved state (zero value if hasState is false) +// +// Example: +// +// func (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) { +// wasInterrupted, hasState, state := tool.GetInterruptState[MyState](ctx) +// if wasInterrupted && hasState { +// // Continue from saved state +// return continueFrom(state), nil +// } +// // First run +// return "", tool.StatefulInterrupt(ctx, "need input", MyState{Step: 1}) +// } +func GetInterruptState[T any](ctx context.Context) (wasInterrupted bool, hasState bool, state T) { + return core.GetInterruptState[T](ctx) +} + +// GetResumeContext checks if this tool is the explicit target of a resume operation. +// +// Returns: +// - isResumeTarget: true if this tool was explicitly targeted for resume +// - hasData: true if resume data was provided +// - data: the resume data (zero value if hasData is false) +// +// Use this to differentiate between: +// - Being resumed as the target (should proceed with work) +// - Being re-executed because a sibling was resumed (should re-interrupt) +// +// Example: +// +// func (t *MyTool) InvokableRun(ctx context.Context, args string, opts ...Option) (string, error) { +// wasInterrupted, _, _ := tool.GetInterruptState[any](ctx) +// if !wasInterrupted { +// return "", tool.Interrupt(ctx, "need confirmation") +// } +// +// isTarget, hasData, data := tool.GetResumeContext[string](ctx) +// if !isTarget { +// // Not our turn - re-interrupt +// return "", tool.Interrupt(ctx, nil) +// } +// if hasData { +// return data, nil +// } +// return "default result", nil +// } +func GetResumeContext[T any](ctx context.Context) (isResumeTarget bool, hasData bool, data T) { + return core.GetResumeContext[T](ctx) +} diff --git a/components/tool/interrupt_test.go b/components/tool/interrupt_test.go new file mode 100644 index 0000000..2f20807 --- /dev/null +++ b/components/tool/interrupt_test.go @@ -0,0 +1,134 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/core" +) + +func TestInterrupt(t *testing.T) { + ctx := context.Background() + + t.Run("basic interrupt", func(t *testing.T) { + err := Interrupt(ctx, "test info") + assert.Error(t, err) + + var signal *core.InterruptSignal + assert.True(t, errors.As(err, &signal)) + assert.Equal(t, "test info", signal.Info) + assert.True(t, signal.IsRootCause) + }) +} + +func TestStatefulInterrupt(t *testing.T) { + ctx := context.Background() + + t.Run("stateful interrupt", func(t *testing.T) { + type myState struct { + Value int + } + state := &myState{Value: 42} + + err := StatefulInterrupt(ctx, "test info", state) + assert.Error(t, err) + + var signal *core.InterruptSignal + assert.True(t, errors.As(err, &signal)) + assert.Equal(t, "test info", signal.Info) + assert.Equal(t, state, signal.State) + assert.True(t, signal.IsRootCause) + }) +} + +func TestCompositeInterrupt(t *testing.T) { + ctx := context.Background() + + t.Run("no sub errors falls back to StatefulInterrupt", func(t *testing.T) { + err := CompositeInterrupt(ctx, "composite info", "my state") + assert.Error(t, err) + + var signal *core.InterruptSignal + assert.True(t, errors.As(err, &signal)) + assert.Equal(t, "composite info", signal.Info) + assert.Equal(t, "my state", signal.State) + assert.True(t, signal.IsRootCause) + assert.Empty(t, signal.Subs) + }) + + t.Run("with InterruptSignal sub error", func(t *testing.T) { + subSignal, _ := core.Interrupt(ctx, "sub info", "sub state", nil) + + err := CompositeInterrupt(ctx, "composite info", "my state", subSignal) + assert.Error(t, err) + + var signal *core.InterruptSignal + assert.True(t, errors.As(err, &signal)) + assert.Equal(t, "composite info", signal.Info) + assert.Equal(t, "my state", signal.State) + assert.Len(t, signal.Subs, 1) + assert.Equal(t, "sub info", signal.Subs[0].Info) + }) + + t.Run("with non-interrupt error returns error", func(t *testing.T) { + nonInterruptErr := errors.New("regular error") + + err := CompositeInterrupt(ctx, "composite info", "my state", nonInterruptErr) + assert.Error(t, err) + assert.Contains(t, err.Error(), "composite interrupt but one of the sub error is not interrupt error") + + var signal *core.InterruptSignal + assert.False(t, errors.As(err, &signal)) + }) + + t.Run("with multiple sub errors", func(t *testing.T) { + subSignal1, _ := core.Interrupt(ctx, "sub1 info", nil, nil) + subSignal2, _ := core.Interrupt(ctx, "sub2 info", nil, nil) + + err := CompositeInterrupt(ctx, "composite info", nil, subSignal1, subSignal2) + assert.Error(t, err) + + var signal *core.InterruptSignal + assert.True(t, errors.As(err, &signal)) + assert.Len(t, signal.Subs, 2) + }) +} + +func TestGetInterruptState(t *testing.T) { + t.Run("not interrupted returns false", func(t *testing.T) { + ctx := context.Background() + wasInterrupted, hasState, state := GetInterruptState[string](ctx) + assert.False(t, wasInterrupted) + assert.False(t, hasState) + assert.Empty(t, state) + }) +} + +func TestGetResumeContext(t *testing.T) { + t.Run("not resume target returns false", func(t *testing.T) { + ctx := context.Background() + isResumeTarget, hasData, data := GetResumeContext[string](ctx) + assert.False(t, isResumeTarget) + assert.False(t, hasData) + assert.Empty(t, data) + }) +} diff --git a/components/tool/option.go b/components/tool/option.go new file mode 100644 index 0000000..1458448 --- /dev/null +++ b/components/tool/option.go @@ -0,0 +1,78 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +// Option defines call option for InvokableTool or StreamableTool component, which is part of component interface signature. +// Each tool implementation could define its own options struct and option funcs within its own package, +// then wrap the impl specific option funcs into this type, before passing to InvokableRun or StreamableRun. +type Option struct { + implSpecificOptFn any +} + +// WrapImplSpecificOptFn wraps the impl specific option functions into Option type. +// T: the type of the impl specific options struct. +// Tool implementations are required to use this function to convert its own option functions into the unified Option type. +// For example, if the tool defines its own options struct: +// +// type customOptions struct { +// conf string +// } +// +// Then the tool needs to provide an option function as such: +// +// func WithConf(conf string) Option { +// return WrapImplSpecificOptFn(func(o *customOptions) { +// o.conf = conf +// } +// } +// +// . +func WrapImplSpecificOptFn[T any](optFn func(*T)) Option { + return Option{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions provides tool author the ability to extract their own custom options from the unified Option type. +// T: the type of the impl specific options struct. +// This function should be used within the tool implementation's InvokableRun or StreamableRun functions. +// It is recommended to provide a base T as the first argument, within which the tool author can provide default values for the impl specific options. +// eg. +// +// type customOptions struct { +// conf string +// } +// defaultOptions := &customOptions{} +// +// customOptions := tool.GetImplSpecificOptions(defaultOptions, opts...) +func GetImplSpecificOptions[T any](base *T, opts ...Option) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} diff --git a/components/tool/option_test.go b/components/tool/option_test.go new file mode 100644 index 0000000..2a27604 --- /dev/null +++ b/components/tool/option_test.go @@ -0,0 +1,54 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package tool + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" +) + +func TestImplSpecificOpts(t *testing.T) { + convey.Convey("TestImplSpecificOpts", t, func() { + type implSpecificOptions struct { + conf string + index int + } + + withConf := func(conf string) func(o *implSpecificOptions) { + return func(o *implSpecificOptions) { + o.conf = conf + } + } + + withIndex := func(index int) func(o *implSpecificOptions) { + return func(o *implSpecificOptions) { + o.index = index + } + } + + toolOption1 := WrapImplSpecificOptFn(withConf("test_conf")) + toolOption2 := WrapImplSpecificOptFn(withIndex(1)) + + implSpecificOpts := GetImplSpecificOptions(&implSpecificOptions{}, toolOption1, toolOption2) + + convey.So(implSpecificOpts, convey.ShouldResemble, &implSpecificOptions{ + conf: "test_conf", + index: 1, + }) + }) +} diff --git a/components/tool/utils/common.go b/components/tool/utils/common.go new file mode 100644 index 0000000..6228708 --- /dev/null +++ b/components/tool/utils/common.go @@ -0,0 +1,28 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "github.com/bytedance/sonic" +) + +func marshalString(resp any) (string, error) { + if rs, ok := resp.(string); ok { + return rs, nil + } + return sonic.MarshalString(resp) +} diff --git a/components/tool/utils/common_test.go b/components/tool/utils/common_test.go new file mode 100644 index 0000000..e015dd8 --- /dev/null +++ b/components/tool/utils/common_test.go @@ -0,0 +1,187 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestMarshalString(t *testing.T) { + tests := []struct { + name string + input interface{} + expected string + hasError bool + }{ + { + name: "string input should return as-is", + input: "hello world", + expected: "hello world", + hasError: false, + }, + { + name: "empty string should return empty", + input: "", + expected: "", + hasError: false, + }, + { + name: "string with special characters", + input: "hello\nworld\t\"test\"", + expected: "hello\nworld\t\"test\"", + hasError: false, + }, + { + name: "string with unicode", + input: "你好世界", + expected: "你好世界", + hasError: false, + }, + { + name: "integer should be marshaled to JSON", + input: 42, + expected: "42", + hasError: false, + }, + { + name: "float should be marshaled to JSON", + input: 3.14, + expected: "3.14", + hasError: false, + }, + { + name: "boolean true should be marshaled to JSON", + input: true, + expected: "true", + hasError: false, + }, + { + name: "boolean false should be marshaled to JSON", + input: false, + expected: "false", + hasError: false, + }, + { + name: "nil should be marshaled to JSON null", + input: nil, + expected: "null", + hasError: false, + }, + { + name: "slice should be marshaled to JSON array", + input: []int{1, 2, 3}, + expected: "[1,2,3]", + hasError: false, + }, + { + name: "empty slice should be marshaled to JSON empty array", + input: []int{}, + expected: "[]", + hasError: false, + }, + { + name: "empty map should be marshaled to JSON empty object", + input: map[string]int{}, + expected: "{}", + hasError: false, + }, + { + name: "struct should be marshaled to JSON", + input: struct{ Name string }{Name: "test"}, + expected: `{"Name":"test"}`, + hasError: false, + }, + { + name: "pointer to string should be handled as non-string", + input: func() *string { s := "test"; return &s }(), + expected: `"test"`, + hasError: false, + }, + { + name: "interface{} containing string should return as-is", + input: interface{}("test string"), + expected: "test string", + hasError: false, + }, + { + name: "interface{} containing int should be marshaled", + input: interface{}(123), + expected: "123", + hasError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := marshalString(tt.input) + + if tt.hasError { + assert.Error(t, err) + } else { + assert.NoError(t, err) + assert.Equal(t, tt.expected, result) + } + }) + } +} + +func TestMarshalStringEdgeCases(t *testing.T) { + t.Run("complex nested structure", func(t *testing.T) { + complex := map[string]interface{}{ + "string": "value", + "number": 42, + "nested": map[string]interface{}{ + "array": []string{"a", "b", "c"}, + "bool": true, + }, + } + + result, err := marshalString(complex) + assert.NoError(t, err) + assert.Contains(t, result, `"string":"value"`) + assert.Contains(t, result, `"number":42`) + assert.Contains(t, result, `"nested"`) + }) + + t.Run("string type assertion priority", func(t *testing.T) { + // Test that string type assertion has priority over JSON marshaling + var input interface{} = "direct string" + result, err := marshalString(input) + assert.NoError(t, err) + assert.Equal(t, "direct string", result) + + // Verify it's not JSON encoded + assert.NotEqual(t, `"direct string"`, result) + }) +} + +func TestMarshalStringConsistency(t *testing.T) { + t.Run("string vs JSON marshaling difference", func(t *testing.T) { + input := `{"key": "value"}` + + // Direct string should return as-is + result, err := marshalString(input) + assert.NoError(t, err) + assert.Equal(t, input, result) + + // Should not be double-encoded + assert.NotEqual(t, fmt.Sprintf(`"%s"`, input), result) + }) +} diff --git a/components/tool/utils/create_options.go b/components/tool/utils/create_options.go new file mode 100644 index 0000000..ea09e33 --- /dev/null +++ b/components/tool/utils/create_options.go @@ -0,0 +1,82 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + "reflect" + + "github.com/eino-contrib/jsonschema" +) + +// UnmarshalArguments is the function type for unmarshalling the arguments. +type UnmarshalArguments func(ctx context.Context, arguments string) (any, error) + +// MarshalOutput is the function type for marshalling the output. +type MarshalOutput func(ctx context.Context, output any) (string, error) + +type toolOptions struct { + um UnmarshalArguments + m MarshalOutput + scModifier SchemaModifierFn +} + +// Option is the option func for the tool. +type Option func(o *toolOptions) + +// WithUnmarshalArguments wraps the unmarshal arguments option. +// when you want to unmarshal the arguments by yourself, you can use this option. +func WithUnmarshalArguments(um UnmarshalArguments) Option { + return func(o *toolOptions) { + o.um = um + } +} + +// WithMarshalOutput wraps the marshal output option. +// when you want to marshal the output by yourself, you can use this option. +func WithMarshalOutput(m MarshalOutput) Option { + return func(o *toolOptions) { + o.m = m + } +} + +// SchemaModifierFn is the schema modifier function for inferring tool parameter from tagged go struct. +// Within this function, end-user can parse custom go struct tags into corresponding json schema field. +// Parameters: +// 1. jsonTagName: the name defined in the json tag. Specifically, the last 'jsonTagName' visited is fixed to be '_root', which represents the entire go struct. Also, for array field, both the field itself and the element within the array will trigger this function. +// 2. t: the type of current schema, usually the field type of the go struct. +// 3. tag: the struct tag of current schema, usually the field tag of the go struct. Note that the element within an array field will use the same go struct tag as the array field itself. +// 4. schema: the current json schema object to be modified. +type SchemaModifierFn func(jsonTagName string, t reflect.Type, tag reflect.StructTag, schema *jsonschema.Schema) + +// WithSchemaModifier sets a user-defined schema modifier for inferring tool parameter from tagged go struct. +func WithSchemaModifier(modifier SchemaModifierFn) Option { + return func(o *toolOptions) { + o.scModifier = modifier + } +} + +func getToolOptions(opt ...Option) *toolOptions { + opts := &toolOptions{ + um: nil, + m: nil, + } + for _, o := range opt { + o(opts) + } + return opts +} diff --git a/components/tool/utils/doc.go b/components/tool/utils/doc.go new file mode 100644 index 0000000..654e7ea --- /dev/null +++ b/components/tool/utils/doc.go @@ -0,0 +1,57 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package utils provides constructors for building tool implementations without +// writing boilerplate JSON serialization code. +// +// # Choosing a Constructor +// +// There are two main strategies: +// +// 1. Infer from struct tags (recommended): [InferTool], [InferStreamTool], +// [InferEnhancedTool], [InferEnhancedStreamTool]. +// The parameter JSON schema is derived automatically from the input struct's +// field names and tags. Requires a typed input struct. +// +// 2. Manual ToolInfo: [NewTool], [NewStreamTool], [NewEnhancedTool], +// [NewEnhancedStreamTool]. +// You supply a [schema.ToolInfo] directly. Useful when the schema cannot +// be expressed as a Go struct, or must be dynamically constructed. +// +// # Struct Tag Convention +// +// InferTool and friends use the following tags on the input struct fields: +// +// type Input struct { +// Query string `json:"query" jsonschema:"required" jsonschema_description:"The search query"` +// MaxItems int `json:"max_items" jsonschema_description:"Maximum results to return"` +// } +// +// Key rules: +// - Use a separate jsonschema_description tag for field descriptions — +// embedding descriptions inside the jsonschema tag causes comma-parsing +// issues. +// - Use jsonschema:"required" to mark mandatory parameters. +// - The json tag controls the parameter name visible to the model. +// +// # Schema Utilities +// +// [GoStruct2ToolInfo] and [GoStruct2ParamsOneOf] convert a Go struct to schema +// types without creating a tool — useful for ChatModel structured output via +// ResponseFormat or BindTools. +// +// See https://www.cloudwego.io/docs/eino/core_modules/components/tools_node_guide/how_to_create_a_tool/ +package utils diff --git a/components/tool/utils/error_handler.go b/components/tool/utils/error_handler.go new file mode 100644 index 0000000..a53d463 --- /dev/null +++ b/components/tool/utils/error_handler.go @@ -0,0 +1,166 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// ErrorHandler converts a tool error into a string response. +type ErrorHandler func(context.Context, error) string + +// WrapToolWithErrorHandler wraps any BaseTool with custom error handling. +// This function detects the tool type (InvokableTool, StreamableTool, or both) +// and applies the appropriate error handling wrapper. +// When the wrapped tool returns an error, the error handler function 'h' will be called +// to convert the error into a string result, and no error will be returned from the wrapper. +// +// Parameters: +// - t: The original BaseTool to be wrapped +// - h: A function that converts an error to a string +// +// Returns: +// - A wrapped BaseTool that handles errors internally based on its capabilities +func WrapToolWithErrorHandler(t tool.BaseTool, h ErrorHandler) tool.BaseTool { + ih := &infoHelper{info: t.Info} + var s tool.StreamableTool + if st, ok := t.(tool.StreamableTool); ok { + s = st + } + if it, ok := t.(tool.InvokableTool); ok { + if s == nil { + return WrapInvokableToolWithErrorHandler(it, h) + } else { + return &combinedErrorWrapper{ + infoHelper: ih, + errorHelper: &errorHelper{ + i: it.InvokableRun, + h: h, + }, + streamErrorHelper: &streamErrorHelper{ + s: s.StreamableRun, + h: h, + }, + } + } + } + if s != nil { + return WrapStreamableToolWithErrorHandler(s, h) + } + return t +} + +// WrapInvokableToolWithErrorHandler wraps an InvokableTool with custom error handling. +// When the wrapped tool returns an error, the error handler function 'h' will be called +// to convert the error into a string result, and no error will be returned from the wrapper. +// +// Parameters: +// - tool: The original InvokableTool to be wrapped +// - h: A function that converts an error to a string +// +// Returns: +// - A wrapped InvokableTool that handles errors internally +func WrapInvokableToolWithErrorHandler(t tool.InvokableTool, h ErrorHandler) tool.InvokableTool { + return &errorWrapper{ + infoHelper: &infoHelper{info: t.Info}, + errorHelper: &errorHelper{ + i: t.InvokableRun, + h: h, + }, + } +} + +// WrapStreamableToolWithErrorHandler wraps a StreamableTool with custom error handling. +// When the wrapped tool returns an error, the error handler function 'h' will be called +// to convert the error into a string result, which will be returned as a single-item stream, +// and no error will be returned from the wrapper. +// +// Parameters: +// - tool: The original StreamableTool to be wrapped +// - h: A function that converts an error to a string +// +// Returns: +// - A wrapped StreamableTool that handles errors internally +func WrapStreamableToolWithErrorHandler(t tool.StreamableTool, h ErrorHandler) tool.StreamableTool { + return &streamErrorWrapper{ + infoHelper: &infoHelper{info: t.Info}, + streamErrorHelper: &streamErrorHelper{ + s: t.StreamableRun, + h: h, + }, + } +} + +type errorWrapper struct { + *infoHelper + *errorHelper +} + +type streamErrorWrapper struct { + *infoHelper + *streamErrorHelper +} + +type combinedErrorWrapper struct { + *infoHelper + *errorHelper + *streamErrorHelper +} + +type infoHelper struct { + info func(ctx context.Context) (*schema.ToolInfo, error) +} + +func (i *infoHelper) Info(ctx context.Context) (*schema.ToolInfo, error) { + return i.info(ctx) +} + +type errorHelper struct { + i func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) + h ErrorHandler +} + +func (s *errorHelper) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + result, err := s.i(ctx, argumentsInJSON, opts...) + if _, ok := compose.IsInterruptRerunError(err); ok { + return result, err + } + if err != nil { + return s.h(ctx, err), nil + } + return result, nil +} + +type streamErrorHelper struct { + s func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) + h ErrorHandler +} + +func (s *streamErrorHelper) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + result, err := s.s(ctx, argumentsInJSON, opts...) + if _, ok := compose.IsInterruptRerunError(err); ok { + return result, err + } + if err != nil { + return schema.StreamReaderFromArray([]string{s.h(ctx, err)}), nil + } + return result, nil +} diff --git a/components/tool/utils/error_handler_test.go b/components/tool/utils/error_handler_test.go new file mode 100644 index 0000000..06648de --- /dev/null +++ b/components/tool/utils/error_handler_test.go @@ -0,0 +1,81 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +type testErrorTool struct{} + +func (t *testErrorTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return nil, nil +} + +func (t *testErrorTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + return "", errors.New("test error") +} + +func (t *testErrorTool) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + return nil, errors.New("test stream error") +} + +func TestErrorWrapper(t *testing.T) { + ctx := context.Background() + + nt := WrapToolWithErrorHandler(&testErrorTool{}, func(_ context.Context, err error) string { + return err.Error() + }) + result, err := nt.(tool.InvokableTool).InvokableRun(ctx, "") + assert.NoError(t, err) + assert.Equal(t, "test error", result) + streamResult, err := nt.(tool.StreamableTool).StreamableRun(ctx, "") + assert.NoError(t, err) + chunk, err := streamResult.Recv() + assert.NoError(t, err) + assert.Equal(t, "test stream error", chunk) + _, err = streamResult.Recv() + assert.True(t, errors.Is(err, io.EOF)) + + wrappedTool := WrapInvokableToolWithErrorHandler(&testErrorTool{}, func(_ context.Context, err error) string { + return err.Error() + }) + result, err = wrappedTool.InvokableRun(ctx, "") + assert.NoError(t, err) + assert.Equal(t, "test error", result) + + wrappedStreamTool := WrapStreamableToolWithErrorHandler(&testErrorTool{}, func(_ context.Context, err error) string { + return err.Error() + }) + + streamResult, err = wrappedStreamTool.StreamableRun(ctx, "") + assert.NoError(t, err) + + chunk, err = streamResult.Recv() + assert.NoError(t, err) + assert.Equal(t, "test stream error", chunk) + _, err = streamResult.Recv() + assert.True(t, errors.Is(err, io.EOF)) +} diff --git a/components/tool/utils/invokable_func.go b/components/tool/utils/invokable_func.go new file mode 100644 index 0000000..d5c9b3a --- /dev/null +++ b/components/tool/utils/invokable_func.go @@ -0,0 +1,318 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + "fmt" + "strings" + + "github.com/bytedance/sonic" + "github.com/eino-contrib/jsonschema" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +// InvokeFunc is the function type for the tool. +type InvokeFunc[T, D any] func(ctx context.Context, input T) (output D, err error) + +// OptionableInvokeFunc is the function type for the tool with tool option. +type OptionableInvokeFunc[T, D any] func(ctx context.Context, input T, opts ...tool.Option) (output D, err error) + +// InferTool creates an [tool.InvokableTool] by inferring the parameter JSON +// schema from the fields and tags of the input type T. +// +// The tool automatically JSON-decodes the model's argument string into T before +// calling fn, and JSON-encodes the D return value into the result string. +// +// Use [WithSchemaModifier] in opts to customise how struct tags are mapped to +// JSON schema fields. +func InferTool[T, D any](toolName, toolDesc string, i InvokeFunc[T, D], opts ...Option) (tool.InvokableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return NewTool(ti, i, opts...), nil +} + +// InferOptionableTool is like [InferTool] but the function also receives +// [tool.Option] values passed by ToolsNode at call time. +func InferOptionableTool[T, D any](toolName, toolDesc string, i OptionableInvokeFunc[T, D], opts ...Option) (tool.InvokableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return newOptionableTool(ti, i, opts...), nil +} + +// EnhancedInvokeFunc is the function type for the enhanced tool. +type EnhancedInvokeFunc[T any] func(ctx context.Context, input T) (output *schema.ToolResult, err error) + +// OptionableEnhancedInvokeFunc is the function type for the enhanced tool with tool option. +type OptionableEnhancedInvokeFunc[T any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.ToolResult, err error) + +// InferEnhancedTool creates an [tool.EnhancedInvokableTool] by inferring the +// parameter JSON schema from type T. The function returns a [schema.ToolResult] +// for multimodal output (text, images, audio, video, files). +func InferEnhancedTool[T any](toolName, toolDesc string, i EnhancedInvokeFunc[T], opts ...Option) (tool.EnhancedInvokableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return NewEnhancedTool(ti, i, opts...), nil +} + +// InferOptionableEnhancedTool creates an EnhancedInvokableTool from a given function by inferring the ToolInfo from the function's request parameters, with tool option. +func InferOptionableEnhancedTool[T any](toolName, toolDesc string, i OptionableEnhancedInvokeFunc[T], opts ...Option) (tool.EnhancedInvokableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return newOptionableEnhancedTool(ti, i, opts...), nil +} + +// GoStruct2ParamsOneOf converts a Go struct's fields and tags into a +// [schema.ParamsOneOf] (JSON Schema 2020-12). Useful for ChatModel structured +// output via ResponseFormat without creating a full tool. +func GoStruct2ParamsOneOf[T any](opts ...Option) (*schema.ParamsOneOf, error) { + return goStruct2ParamsOneOf[T](opts...) +} + +// GoStruct2ToolInfo converts a Go struct into a [schema.ToolInfo]. Useful for +// binding a typed schema to a ChatModel via BindTools for structured output, +// when you do not need a full executable tool. +func GoStruct2ToolInfo[T any](toolName, toolDesc string, opts ...Option) (*schema.ToolInfo, error) { + return goStruct2ToolInfo[T](toolName, toolDesc, opts...) +} + +func goStruct2ToolInfo[T any](toolName, toolDesc string, opts ...Option) (*schema.ToolInfo, error) { + paramsOneOf, err := goStruct2ParamsOneOf[T](opts...) + if err != nil { + return nil, err + } + return &schema.ToolInfo{ + Name: toolName, + Desc: toolDesc, + ParamsOneOf: paramsOneOf, + }, nil +} + +func goStruct2ParamsOneOf[T any](opts ...Option) (*schema.ParamsOneOf, error) { + options := getToolOptions(opts...) + + r := &jsonschema.Reflector{ + Anonymous: true, + DoNotReference: true, + SchemaModifier: jsonschema.SchemaModifierFn(options.scModifier), + } + + js := r.Reflect(generic.NewInstance[T]()) + js.Version = "" + + paramsOneOf := schema.NewParamsOneOfByJSONSchema(js) + + return paramsOneOf, nil +} + +// NewTool creates an [tool.InvokableTool] from an explicit [schema.ToolInfo] +// and a typed function. Use this when the schema cannot be inferred from struct +// tags (e.g. dynamic or complex parameter schemas). +// +// Note: you are responsible for keeping desc.ParamsOneOf consistent with the +// actual fields of T — there is no compile-time check. +func NewTool[T, D any](desc *schema.ToolInfo, i InvokeFunc[T, D], opts ...Option) tool.InvokableTool { + return newOptionableTool(desc, func(ctx context.Context, input T, _ ...tool.Option) (D, error) { + return i(ctx, input) + }, opts...) +} + +func newOptionableTool[T, D any](desc *schema.ToolInfo, i OptionableInvokeFunc[T, D], opts ...Option) tool.InvokableTool { + to := getToolOptions(opts...) + + return &invokableTool[T, D]{ + info: desc, + um: to.um, + m: to.m, + Fn: i, + } +} + +type invokableTool[T, D any] struct { + info *schema.ToolInfo + + um UnmarshalArguments + m MarshalOutput + + Fn OptionableInvokeFunc[T, D] +} + +func (i *invokableTool[T, D]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return i.info, nil +} + +// InvokableRun invokes the tool with the given arguments. +func (i *invokableTool[T, D]) InvokableRun(ctx context.Context, arguments string, opts ...tool.Option) (output string, err error) { + + var inst T + if i.um != nil { + var val any + val, err = i.um(ctx, arguments) + if err != nil { + return "", fmt.Errorf("[LocalFunc] failed to unmarshal arguments, toolName=%s, err=%w", i.getToolName(), err) + } + gt, ok := val.(T) + if !ok { + return "", fmt.Errorf("[LocalFunc] invalid type, toolName=%s, expected=%T, given=%T", i.getToolName(), inst, val) + } + inst = gt + } else { + inst = generic.NewInstance[T]() + + err = sonic.UnmarshalString(arguments, &inst) + if err != nil { + return "", fmt.Errorf("[LocalFunc] failed to unmarshal arguments in json, toolName=%s, err=%w", i.getToolName(), err) + } + } + + resp, err := i.Fn(ctx, inst, opts...) + if err != nil { + return "", fmt.Errorf("[LocalFunc] failed to invoke tool, toolName=%s, err=%w", i.getToolName(), err) + } + + if i.m != nil { + output, err = i.m(ctx, resp) + if err != nil { + return "", fmt.Errorf("[LocalFunc] failed to marshal output, toolName=%s, err=%w", i.getToolName(), err) + } + } else { + output, err = marshalString(resp) + if err != nil { + return "", fmt.Errorf("[LocalFunc] failed to marshal output in json, toolName=%s, err=%w", i.getToolName(), err) + } + } + + return output, nil +} + +func (i *invokableTool[T, D]) GetType() string { + return snakeToCamel(i.getToolName()) +} + +func (i *invokableTool[T, D]) getToolName() string { + if i.info == nil { + return "" + } + + return i.info.Name +} + +// snakeToCamel converts a snake_case string to CamelCase. +func snakeToCamel(s string) string { + if s == "" { + return "" + } + + parts := strings.Split(s, "_") + + for i := 0; i < len(parts); i++ { + if len(parts[i]) > 0 { + parts[i] = strings.ToUpper(string(parts[i][0])) + strings.ToLower(parts[i][1:]) + } + } + + return strings.Join(parts, "") +} + +// NewEnhancedTool creates an [tool.EnhancedInvokableTool] from an explicit +// [schema.ToolInfo] and a function that returns [schema.ToolResult]. +func NewEnhancedTool[T any](desc *schema.ToolInfo, i EnhancedInvokeFunc[T], opts ...Option) tool.EnhancedInvokableTool { + return newOptionableEnhancedTool(desc, func(ctx context.Context, input T, _ ...tool.Option) (*schema.ToolResult, error) { + return i(ctx, input) + }, opts...) +} + +func newOptionableEnhancedTool[T any](desc *schema.ToolInfo, i OptionableEnhancedInvokeFunc[T], opts ...Option) tool.EnhancedInvokableTool { + to := getToolOptions(opts...) + + return &enhancedInvokableTool[T]{ + info: desc, + um: to.um, + Fn: i, + } +} + +type enhancedInvokableTool[T any] struct { + info *schema.ToolInfo + + um UnmarshalArguments + + Fn OptionableEnhancedInvokeFunc[T] +} + +func (e *enhancedInvokableTool[T]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return e.info, nil +} + +func (e *enhancedInvokableTool[T]) InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + var inst T + var err error + + if e.um != nil { + var val any + val, err = e.um(ctx, toolArgument.Text) + if err != nil { + return nil, fmt.Errorf("[EnhancedLocalFunc] failed to unmarshal arguments, toolName=%s, err=%w", e.getToolName(), err) + } + gt, ok := val.(T) + if !ok { + return nil, fmt.Errorf("[EnhancedLocalFunc] invalid type, toolName=%s, expected=%T, given=%T", e.getToolName(), inst, val) + } + inst = gt + } else { + inst = generic.NewInstance[T]() + + err = sonic.UnmarshalString(toolArgument.Text, &inst) + if err != nil { + return nil, fmt.Errorf("[EnhancedLocalFunc] failed to unmarshal arguments in json, toolName=%s, err=%w", e.getToolName(), err) + } + } + + resp, err := e.Fn(ctx, inst, opts...) + if err != nil { + return nil, fmt.Errorf("[EnhancedLocalFunc] failed to invoke tool, toolName=%s, err=%w", e.getToolName(), err) + } + + return resp, nil +} + +func (e *enhancedInvokableTool[T]) GetType() string { + return snakeToCamel(e.getToolName()) +} + +func (e *enhancedInvokableTool[T]) getToolName() string { + if e.info == nil { + return "" + } + + return e.info.Name +} diff --git a/components/tool/utils/invokable_func_test.go b/components/tool/utils/invokable_func_test.go new file mode 100644 index 0000000..1330f17 --- /dev/null +++ b/components/tool/utils/invokable_func_test.go @@ -0,0 +1,405 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "github.com/eino-contrib/jsonschema" + "github.com/stretchr/testify/assert" + orderedmap "github.com/wk8/go-ordered-map/v2" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +type Job struct { + Company string `json:"company" jsonschema:"description=the company where the user works"` + Position string `json:"position,omitempty" jsonschema:"description=the position of the user's job"` + ServiceLength float32 `json:"service_length,omitempty" jsonschema:"description=the year of user's service"` // 司龄，年 +} + +type Income struct { + Source string `json:"source" jsonschema:"description=the source of income"` + Amount int `json:"amount" jsonschema:"description=the amount of income"` + HasPayTax bool `json:"has_pay_tax" jsonschema:"description=whether the user has paid tax"` + Job *Job `json:"job,omitempty" jsonschema:"description=the job of the user when earning this income"` +} + +type User struct { + Name string `json:"name" jsonschema:"required,description=the name of the user"` + Age int `json:"age" jsonschema:"required,description=the age of the user"` + + Job *Job `json:"job,omitempty" jsonschema:"description=the job of the user"` + + Incomes []*Income `json:"incomes" jsonschema:"description=the incomes of the user"` +} + +type UserResult struct { + Code int `json:"code"` + Msg string `json:"msg"` +} + +var toolInfo = &schema.ToolInfo{ + Name: "update_user_info", + Desc: "full update user info", + ParamsOneOf: schema.NewParamsOneOfByJSONSchema( + &jsonschema.Schema{ + Type: "object", + Required: []string{"name", "age", "incomes"}, + AdditionalProperties: jsonschema.FalseSchema, + Properties: orderedmap.New[string, *jsonschema.Schema]( + orderedmap.WithInitialData( + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "name", + Value: &jsonschema.Schema{ + Type: "string", + Description: "the name of the user", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "age", + Value: &jsonschema.Schema{ + Type: "integer", + Description: "the age of the user", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "job", + Value: &jsonschema.Schema{ + Type: "object", + Required: []string{"company"}, + AdditionalProperties: jsonschema.FalseSchema, + Description: "the job of the user", + Properties: orderedmap.New[string, *jsonschema.Schema]( + orderedmap.WithInitialData( + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "company", + Value: &jsonschema.Schema{ + Type: "string", + Description: "the company where the user works", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "position", + Value: &jsonschema.Schema{ + Type: "string", + Description: "the position of the user's job", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "service_length", + Value: &jsonschema.Schema{ + Type: "number", + Description: "the year of user's service", + }, + }, + ), + ), + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "incomes", + Value: &jsonschema.Schema{ + Type: "array", + Description: "the incomes of the user", + Items: &jsonschema.Schema{ + Type: "object", + AdditionalProperties: jsonschema.FalseSchema, + Required: []string{"source", "amount", "has_pay_tax"}, + Properties: orderedmap.New[string, *jsonschema.Schema]( + orderedmap.WithInitialData( + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "source", + Value: &jsonschema.Schema{ + Type: "string", + Description: "the source of income", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "amount", + Value: &jsonschema.Schema{ + Type: "integer", + Description: "the amount of income", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "has_pay_tax", + Value: &jsonschema.Schema{ + Type: "boolean", + Description: "whether the user has paid tax", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "job", + Value: &jsonschema.Schema{ + Type: "object", + AdditionalProperties: jsonschema.FalseSchema, + Required: []string{"company"}, + Description: "the job of the user when earning this income", + Properties: orderedmap.New[string, *jsonschema.Schema]( + orderedmap.WithInitialData( + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "company", + Value: &jsonschema.Schema{ + Type: "string", + Description: "the company where the user works", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "position", + Value: &jsonschema.Schema{ + Type: "string", + Description: "the position of the user's job", + }, + }, + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "service_length", + Value: &jsonschema.Schema{ + Type: "number", + Description: "the year of user's service", + }, + }, + ), + ), + }, + }, + ), + ), + }, + }, + }, + ), + ), + }), +} + +func updateUserInfo(ctx context.Context, input *User) (output *UserResult, err error) { + return &UserResult{ + Code: 200, + Msg: fmt.Sprintf("update %v success", input.Name), + }, nil +} + +type UserInfoOption struct { + Field1 string +} + +func WithUserInfoOption(s string) tool.Option { + return tool.WrapImplSpecificOptFn(func(t *UserInfoOption) { + t.Field1 = s + }) +} + +func updateUserInfoWithOption(_ context.Context, input *User, opts ...tool.Option) (output *UserResult, err error) { + baseOption := &UserInfoOption{ + Field1: "test_origin", + } + + option := tool.GetImplSpecificOptions(baseOption, opts...) + return &UserResult{ + Code: 200, + Msg: option.Field1, + }, nil +} + +func TestInferTool(t *testing.T) { + t.Run("invoke_infer_tool", func(t *testing.T) { + ctx := context.Background() + + tl, err := InferTool("update_user_info", "full update user info", updateUserInfo) + assert.NoError(t, err) + + info, err := tl.Info(context.Background()) + assert.NoError(t, err) + + actual, err := info.ToJSONSchema() + assert.NoError(t, err) + actualStr, err := json.Marshal(actual) + assert.NoError(t, err) + + expect, err := toolInfo.ToJSONSchema() + assert.NoError(t, err) + expectStr, err := json.Marshal(expect) + assert.NoError(t, err) + + assert.Equal(t, string(expectStr), string(actualStr)) + + content, err := tl.InvokableRun(ctx, `{"name": "bruce lee"}`) + assert.NoError(t, err) + assert.JSONEq(t, `{"code":200,"msg":"update bruce lee success"}`, content) + }) +} + +func TestInferOptionableTool(t *testing.T) { + ctx := context.Background() + + t.Run("invoke_infer_optionable_tool", func(t *testing.T) { + + tl, err := InferOptionableTool("invoke_infer_optionable_tool", "full update user info", updateUserInfoWithOption) + assert.NoError(t, err) + + content, err := tl.InvokableRun(ctx, `{"name": "bruce lee"}`, WithUserInfoOption("hello world")) + assert.NoError(t, err) + assert.JSONEq(t, `{"code":200,"msg":"hello world"}`, content) + }) +} + +func TestNewTool(t *testing.T) { + ctx := context.Background() + type Input struct { + Name string `json:"name"` + } + type Output struct { + Name string `json:"name"` + } + + t.Run("struct_input_struct_output", func(t *testing.T) { + + tl := NewTool[Input, Output](nil, func(ctx context.Context, input Input) (output Output, err error) { + return Output{ + Name: input.Name, + }, nil + }) + + _, err := tl.InvokableRun(ctx, `{"name":"test"}`) + assert.Nil(t, err) + }) + + t.Run("pointer_input_pointer_output", func(t *testing.T) { + tl := NewTool[*Input, *Output](nil, func(ctx context.Context, input *Input) (output *Output, err error) { + return &Output{ + Name: input.Name, + }, nil + }) + + content, err := tl.InvokableRun(ctx, `{"name":"test"}`) + assert.NoError(t, err) + assert.Equal(t, `{"name":"test"}`, content) + }) + + t.Run("string_input_int64_output", func(t *testing.T) { + tl := NewTool(nil, func(ctx context.Context, input string) (output int64, err error) { + return 10, nil + }) + + content, err := tl.InvokableRun(ctx, `100`) // json unmarshal must contains double quote if is not json string. + assert.Error(t, err) + assert.Equal(t, "", content) + }) + + t.Run("string_pointer_input_int64_pointer_output", func(t *testing.T) { + tl := NewTool[*string, *int64](nil, func(ctx context.Context, input *string) (output *int64, err error) { + n := int64(10) + return &n, nil + }) + + content, err := tl.InvokableRun(ctx, `"100"`) + assert.NoError(t, err) + assert.Equal(t, `10`, content) + }) +} + +func TestSnakeToCamel(t *testing.T) { + t.Run("normal_case", func(t *testing.T) { + assert.Equal(t, "GoogleSearch3", snakeToCamel("google_search_3")) + }) + + t.Run("empty_case", func(t *testing.T) { + assert.Equal(t, "", snakeToCamel("")) + }) + + t.Run("single_word_case", func(t *testing.T) { + assert.Equal(t, "Google", snakeToCamel("google")) + }) + + t.Run("upper_case", func(t *testing.T) { + assert.Equal(t, "HttpHost", snakeToCamel("_HTTP_HOST_")) + }) + + t.Run("underscore_case", func(t *testing.T) { + assert.Equal(t, "", snakeToCamel("_")) + }) +} + +type stringAlias string +type integerAlias uint32 +type floatAlias float64 +type boolAlias bool + +type testEnumStruct struct { + Field1 string `json:"field1" jsonschema:"enum=a,enum=b"` + Field2 int `json:"field2" jsonschema:"enum=1,enum=2"` + Field3 float32 `json:"field3" jsonschema:"enum=1.1,enum=2.2"` + Field4 bool `json:"field4" jsonschema:"default=true"` + Field5 stringAlias `json:"field5" jsonschema:"enum=a,enum=c"` + Field6 integerAlias `json:"field6" jsonschema:"enum=3,enum=4"` + Field7 floatAlias `json:"field7" jsonschema:"enum=3.3,enum=4.4"` + Field8 boolAlias `json:"field8" jsonschema:"enum=false"` +} + +type testEnumStruct2 struct { + Field1 int8 `json:"field1" jsonschema:"enum=1.1"` +} + +type testEnumStruct3 struct { + Field1 float64 `json:"field1" jsonschema:"enum=a"` +} + +func TestEnumTag(t *testing.T) { + info, err := goStruct2ParamsOneOf[testEnumStruct]() + assert.NoError(t, err) + s, err := info.ToJSONSchema() + assert.NoError(t, err) + + enum, ok := s.Properties.Get("field1") + assert.True(t, ok) + assert.Equal(t, []any{"a", "b"}, enum.Enum) + + enum, ok = s.Properties.Get("field2") + assert.True(t, ok) + assert.Equal(t, []any{json.Number("1"), json.Number("2")}, enum.Enum) + + enum, ok = s.Properties.Get("field3") + assert.True(t, ok) + assert.Equal(t, []any{json.Number("1.1"), json.Number("2.2")}, enum.Enum) + + enum, ok = s.Properties.Get("field4") + assert.True(t, ok) + assert.Equal(t, true, enum.Default) + + enum, ok = s.Properties.Get("field5") + assert.True(t, ok) + assert.Equal(t, []any{"a", "c"}, enum.Enum) + + enum, ok = s.Properties.Get("field6") + assert.True(t, ok) + assert.Equal(t, []any{json.Number("3"), json.Number("4")}, enum.Enum) + + enum, ok = s.Properties.Get("field7") + assert.True(t, ok) + assert.Equal(t, []any{json.Number("3.3"), json.Number("4.4")}, enum.Enum) + + _, err = goStruct2ParamsOneOf[testEnumStruct2]() + assert.NoError(t, err) + + _, err = goStruct2ParamsOneOf[testEnumStruct3]() + assert.NoError(t, err) +} diff --git a/components/tool/utils/streamable_func.go b/components/tool/utils/streamable_func.go new file mode 100644 index 0000000..c6601e8 --- /dev/null +++ b/components/tool/utils/streamable_func.go @@ -0,0 +1,259 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + "fmt" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +// StreamFunc is the function type for the streamable tool. +type StreamFunc[T, D any] func(ctx context.Context, input T) (output *schema.StreamReader[D], err error) + +// OptionableStreamFunc is the function type for the streamable tool with tool option. +type OptionableStreamFunc[T, D any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.StreamReader[D], err error) + +// InferStreamTool creates a [tool.StreamableTool] by inferring the parameter +// JSON schema from type T. The function returns a [schema.StreamReader] of D +// values which the framework serialises to a string stream. +func InferStreamTool[T, D any](toolName, toolDesc string, s StreamFunc[T, D], opts ...Option) (tool.StreamableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return NewStreamTool(ti, s, opts...), nil +} + +// InferOptionableStreamTool is like [InferStreamTool] but the function also +// receives [tool.Option] values passed by ToolsNode at call time. +func InferOptionableStreamTool[T, D any](toolName, toolDesc string, s OptionableStreamFunc[T, D], opts ...Option) (tool.StreamableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return newOptionableStreamTool(ti, s, opts...), nil +} + +// NewStreamTool creates a [tool.StreamableTool] from an explicit [schema.ToolInfo] +// and a typed streaming function. +func NewStreamTool[T, D any](desc *schema.ToolInfo, s StreamFunc[T, D], opts ...Option) tool.StreamableTool { + return newOptionableStreamTool(desc, + func(ctx context.Context, input T, _ ...tool.Option) (output *schema.StreamReader[D], err error) { + return s(ctx, input) + }, + opts...) +} + +func newOptionableStreamTool[T, D any](desc *schema.ToolInfo, s OptionableStreamFunc[T, D], opts ...Option) tool.StreamableTool { + + to := getToolOptions(opts...) + + return &streamableTool[T, D]{ + info: desc, + + um: to.um, + m: to.m, + Fn: s, + } +} + +type streamableTool[T, D any] struct { + info *schema.ToolInfo + + um UnmarshalArguments + m MarshalOutput + + Fn OptionableStreamFunc[T, D] +} + +// Info returns the tool info, implement the BaseTool interface. +func (s *streamableTool[T, D]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return s.info, nil +} + +// StreamableRun invokes the tool with the given arguments, implement the StreamableTool interface. +func (s *streamableTool[T, D]) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) ( + outStream *schema.StreamReader[string], err error) { + + var inst T + if s.um != nil { + var val any + val, err = s.um(ctx, argumentsInJSON) + if err != nil { + return nil, fmt.Errorf("[LocalStreamFunc] failed to unmarshal arguments, toolName=%s, err=%w", s.getToolName(), err) + } + + gt, ok := val.(T) + if !ok { + return nil, fmt.Errorf("[LocalStreamFunc] type err, toolName=%s, expected=%T, given=%T", s.getToolName(), inst, val) + } + inst = gt + } else { + + inst = generic.NewInstance[T]() + + err = sonic.UnmarshalString(argumentsInJSON, &inst) + if err != nil { + return nil, fmt.Errorf("[LocalStreamFunc] failed to unmarshal arguments in json, toolName=%s, err=%w", s.getToolName(), err) + } + } + + streamD, err := s.Fn(ctx, inst, opts...) + if err != nil { + return nil, err + } + + outStream = schema.StreamReaderWithConvert(streamD, func(d D) (string, error) { + var out string + var e error + if s.m != nil { + out, e = s.m(ctx, d) + if e != nil { + return "", fmt.Errorf("[LocalStreamFunc] failed to marshal output, toolName=%s, err=%w", s.getToolName(), e) + } + } else { + out, e = marshalString(d) + if e != nil { + return "", fmt.Errorf("[LocalStreamFunc] failed to marshal output in json, toolName=%s, err=%w", s.getToolName(), e) + } + } + + return out, nil + }) + + return outStream, nil +} + +func (s *streamableTool[T, D]) GetType() string { + return snakeToCamel(s.getToolName()) +} + +func (s *streamableTool[T, D]) getToolName() string { + if s.info == nil { + return "" + } + + return s.info.Name +} + +// EnhancedStreamFunc is the function type for the enhanced streamable tool. +type EnhancedStreamFunc[T any] func(ctx context.Context, input T) (output *schema.StreamReader[*schema.ToolResult], err error) + +// OptionableEnhancedStreamFunc is the function type for the enhanced streamable tool with tool option. +type OptionableEnhancedStreamFunc[T any] func(ctx context.Context, input T, opts ...tool.Option) (output *schema.StreamReader[*schema.ToolResult], err error) + +// InferEnhancedStreamTool creates an [tool.EnhancedStreamableTool] by inferring +// the parameter JSON schema from type T. The function streams [schema.ToolResult] +// values for multimodal output. +func InferEnhancedStreamTool[T any](toolName, toolDesc string, s EnhancedStreamFunc[T], opts ...Option) (tool.EnhancedStreamableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return NewEnhancedStreamTool(ti, s, opts...), nil +} + +// InferOptionableEnhancedStreamTool creates an EnhancedStreamableTool from a given function by inferring the ToolInfo from the function's request parameters, with tool option. +func InferOptionableEnhancedStreamTool[T any](toolName, toolDesc string, s OptionableEnhancedStreamFunc[T], opts ...Option) (tool.EnhancedStreamableTool, error) { + ti, err := goStruct2ToolInfo[T](toolName, toolDesc, opts...) + if err != nil { + return nil, err + } + + return newOptionableEnhancedStreamTool(ti, s, opts...), nil +} + +// NewEnhancedStreamTool creates an [tool.EnhancedStreamableTool] from an +// explicit [schema.ToolInfo] and a typed streaming function. +func NewEnhancedStreamTool[T any](desc *schema.ToolInfo, s EnhancedStreamFunc[T], opts ...Option) tool.EnhancedStreamableTool { + return newOptionableEnhancedStreamTool(desc, + func(ctx context.Context, input T, _ ...tool.Option) (output *schema.StreamReader[*schema.ToolResult], err error) { + return s(ctx, input) + }, + opts...) +} + +func newOptionableEnhancedStreamTool[T any](desc *schema.ToolInfo, s OptionableEnhancedStreamFunc[T], opts ...Option) tool.EnhancedStreamableTool { + to := getToolOptions(opts...) + + return &enhancedStreamableTool[T]{ + info: desc, + um: to.um, + Fn: s, + } +} + +type enhancedStreamableTool[T any] struct { + info *schema.ToolInfo + + um UnmarshalArguments + + Fn OptionableEnhancedStreamFunc[T] +} + +func (s *enhancedStreamableTool[T]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return s.info, nil +} + +func (s *enhancedStreamableTool[T]) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) ( + outStream *schema.StreamReader[*schema.ToolResult], err error) { + + var inst T + if s.um != nil { + var val any + val, err = s.um(ctx, toolArgument.Text) + if err != nil { + return nil, fmt.Errorf("[EnhancedLocalStreamFunc] failed to unmarshal arguments, toolName=%s, err=%w", s.getToolName(), err) + } + + gt, ok := val.(T) + if !ok { + return nil, fmt.Errorf("[EnhancedLocalStreamFunc] type err, toolName=%s, expected=%T, given=%T", s.getToolName(), inst, val) + } + inst = gt + } else { + inst = generic.NewInstance[T]() + + err = sonic.UnmarshalString(toolArgument.Text, &inst) + if err != nil { + return nil, fmt.Errorf("[EnhancedLocalStreamFunc] failed to unmarshal arguments in json, toolName=%s, err=%w", s.getToolName(), err) + } + } + + return s.Fn(ctx, inst, opts...) +} + +func (s *enhancedStreamableTool[T]) GetType() string { + return snakeToCamel(s.getToolName()) +} + +func (s *enhancedStreamableTool[T]) getToolName() string { + if s.info == nil { + return "" + } + + return s.info.Name +} diff --git a/components/tool/utils/streamable_func_test.go b/components/tool/utils/streamable_func_test.go new file mode 100644 index 0000000..56632a0 --- /dev/null +++ b/components/tool/utils/streamable_func_test.go @@ -0,0 +1,329 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/eino-contrib/jsonschema" + "github.com/stretchr/testify/assert" + orderedmap "github.com/wk8/go-ordered-map/v2" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +func TestNewStreamableTool(t *testing.T) { + ctx := context.Background() + type Input struct { + Name string `json:"name"` + } + type Output struct { + Name string `json:"name"` + } + + t.Run("simple_case", func(t *testing.T) { + tl := NewStreamTool[*Input, *Output]( + &schema.ToolInfo{ + Name: "search_user", + Desc: "search user info", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Type: "string", + Desc: "user name", + }, + }), + }, + func(ctx context.Context, input *Input) (output *schema.StreamReader[*Output], err error) { + sr, sw := schema.Pipe[*Output](2) + sw.Send(&Output{ + Name: input.Name, + }, nil) + sw.Send(&Output{ + Name: "lee", + }, nil) + sw.Close() + + return sr, nil + }, + ) + + info, err := tl.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "search_user", info.Name) + + js, err := info.ToJSONSchema() + assert.NoError(t, err) + + assert.Equal(t, &jsonschema.Schema{ + Type: "object", + Properties: orderedmap.New[string, *jsonschema.Schema]( + orderedmap.WithInitialData[string, *jsonschema.Schema]( + orderedmap.Pair[string, *jsonschema.Schema]{ + Key: "name", + Value: &jsonschema.Schema{ + Type: "string", + Description: "user name", + }, + }, + ), + ), + Required: make([]string, 0), + }, js) + + sr, err := tl.StreamableRun(ctx, `{"name":"xxx"}`) + assert.NoError(t, err) + + defer sr.Close() + + idx := 0 + for { + m, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + assert.NoError(t, err) + + if idx == 0 { + assert.Equal(t, `{"name":"xxx"}`, m) + } else { + assert.Equal(t, `{"name":"lee"}`, m) + } + idx++ + } + + assert.Equal(t, 2, idx) + }) +} + +type FakeStreamOption struct { + Field string +} + +type FakeStreamInferToolInput struct { + Field string `json:"field"` +} + +type FakeStreamInferToolOutput struct { + Field string `json:"field"` +} + +func FakeWithToolOption(s string) tool.Option { + return tool.WrapImplSpecificOptFn(func(t *FakeStreamOption) { + t.Field = s + }) +} + +func fakeStreamFunc(ctx context.Context, input FakeStreamInferToolInput, opts ...tool.Option) (output *schema.StreamReader[*FakeStreamInferToolOutput], err error) { + baseOpt := &FakeStreamOption{ + Field: "default_field_value", + } + option := tool.GetImplSpecificOptions(baseOpt, opts...) + + return schema.StreamReaderFromArray([]*FakeStreamInferToolOutput{ + { + Field: option.Field, + }, + }), nil +} + +func TestInferStreamTool(t *testing.T) { + st, err := InferOptionableStreamTool("infer_optionable_stream_tool", "test infer stream tool with option", fakeStreamFunc) + assert.Nil(t, err) + + sr, err := st.StreamableRun(context.Background(), `{"field": "value"}`, FakeWithToolOption("hello world")) + assert.Nil(t, err) + + defer sr.Close() + + idx := 0 + for { + m, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + assert.NoError(t, err) + + if idx == 0 { + assert.JSONEq(t, `{"field":"hello world"}`, m) + } + } +} + +type EnhancedStreamInput struct { + Query string `json:"query" jsonschema:"description=the search query"` +} + +func TestNewEnhancedStreamTool(t *testing.T) { + ctx := context.Background() + + t.Run("simple_case", func(t *testing.T) { + tl := NewEnhancedStreamTool[*EnhancedStreamInput]( + &schema.ToolInfo{ + Name: "enhanced_stream_search", + Desc: "search with enhanced stream output", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "query": { + Type: "string", + Desc: "the search query", + }, + }), + }, + func(ctx context.Context, input *EnhancedStreamInput) (*schema.StreamReader[*schema.ToolResult], error) { + sr, sw := schema.Pipe[*schema.ToolResult](2) + sw.Send(&schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "result for: " + input.Query}, + }, + }, nil) + sw.Send(&schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "more results"}, + }, + }, nil) + sw.Close() + return sr, nil + }, + ) + + info, err := tl.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "enhanced_stream_search", info.Name) + + sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{"query":"test"}`}) + assert.NoError(t, err) + defer sr.Close() + + idx := 0 + for { + m, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + assert.NoError(t, err) + + if idx == 0 { + assert.Len(t, m.Parts, 1) + assert.Equal(t, schema.ToolPartTypeText, m.Parts[0].Type) + assert.Equal(t, "result for: test", m.Parts[0].Text) + } else { + assert.Len(t, m.Parts, 1) + assert.Equal(t, "more results", m.Parts[0].Text) + } + idx++ + } + assert.Equal(t, 2, idx) + }) +} + +type FakeEnhancedStreamOption struct { + Prefix string +} + +func FakeWithEnhancedStreamOption(prefix string) tool.Option { + return tool.WrapImplSpecificOptFn(func(t *FakeEnhancedStreamOption) { + t.Prefix = prefix + }) +} + +func fakeEnhancedStreamFunc(ctx context.Context, input EnhancedStreamInput) (*schema.StreamReader[*schema.ToolResult], error) { + return schema.StreamReaderFromArray([]*schema.ToolResult{ + { + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "result: " + input.Query}, + }, + }, + }), nil +} + +func fakeOptionableEnhancedStreamFunc(ctx context.Context, input EnhancedStreamInput, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + baseOpt := &FakeEnhancedStreamOption{ + Prefix: "default", + } + option := tool.GetImplSpecificOptions(baseOpt, opts...) + + return schema.StreamReaderFromArray([]*schema.ToolResult{ + { + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: option.Prefix + ": " + input.Query}, + }, + }, + }), nil +} + +func TestInferEnhancedStreamTool(t *testing.T) { + ctx := context.Background() + + t.Run("infer_enhanced_stream_tool", func(t *testing.T) { + tl, err := InferEnhancedStreamTool("infer_enhanced_stream", "test infer enhanced stream tool", fakeEnhancedStreamFunc) + assert.NoError(t, err) + + info, err := tl.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "infer_enhanced_stream", info.Name) + + sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{"query":"hello"}`}) + assert.NoError(t, err) + defer sr.Close() + + m, err := sr.Recv() + assert.NoError(t, err) + assert.Len(t, m.Parts, 1) + assert.Equal(t, "result: hello", m.Parts[0].Text) + }) +} + +func TestInferOptionableEnhancedStreamTool(t *testing.T) { + ctx := context.Background() + + t.Run("infer_optionable_enhanced_stream_tool", func(t *testing.T) { + tl, err := InferOptionableEnhancedStreamTool("infer_optionable_enhanced_stream", "test infer optionable enhanced stream tool", fakeOptionableEnhancedStreamFunc) + assert.NoError(t, err) + + info, err := tl.Info(ctx) + assert.NoError(t, err) + assert.Equal(t, "infer_optionable_enhanced_stream", info.Name) + + sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{"query":"world"}`}, FakeWithEnhancedStreamOption("custom")) + assert.NoError(t, err) + defer sr.Close() + + m, err := sr.Recv() + assert.NoError(t, err) + assert.Len(t, m.Parts, 1) + assert.Equal(t, "custom: world", m.Parts[0].Text) + }) + + t.Run("infer_optionable_enhanced_stream_tool_default_option", func(t *testing.T) { + tl, err := InferOptionableEnhancedStreamTool("infer_optionable_enhanced_stream", "test infer optionable enhanced stream tool", fakeOptionableEnhancedStreamFunc) + assert.NoError(t, err) + + sr, err := tl.StreamableRun(ctx, &schema.ToolArgument{Text: `{"query":"test"}`}) + assert.NoError(t, err) + defer sr.Close() + + m, err := sr.Recv() + assert.NoError(t, err) + assert.Len(t, m.Parts, 1) + assert.Equal(t, "default: test", m.Parts[0].Text) + }) +} diff --git a/components/types.go b/components/types.go new file mode 100644 index 0000000..2b0ad8f --- /dev/null +++ b/components/types.go @@ -0,0 +1,87 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package components defines common interfaces that describe component +// types and callback capabilities used across Eino. +package components + +// Typer provides a human-readable type name for a component implementation. +// +// When implemented, the component's full display name in DevOps tooling +// (visual debugger, IDE plugin, dashboards) becomes "{GetType()}{ComponentKind}" +// — e.g. "OpenAIChatModel". Use CamelCase naming. +// +// Also used by [utils.InferTool] and similar constructors to set the display +// name of tool instances. +type Typer interface { + GetType() string +} + +// GetType returns the type name for a component that implements Typer. +func GetType(component any) (string, bool) { + if typer, ok := component.(Typer); ok { + return typer.GetType(), true + } + + return "", false +} + +// Checker controls whether the framework's automatic callback instrumentation +// is active for a component. +// +// When IsCallbacksEnabled returns true, the framework skips its default +// OnStart/OnEnd wrapping and trusts the component to invoke callbacks itself +// at the correct points. Implement this when your component needs precise +// control over callback timing or content — for example, when streaming +// requires callbacks to fire mid-stream rather than only at completion. +type Checker interface { + IsCallbacksEnabled() bool +} + +// IsCallbacksEnabled reports whether a component implements Checker and enables callbacks. +func IsCallbacksEnabled(i any) bool { + if checker, ok := i.(Checker); ok { + return checker.IsCallbacksEnabled() + } + + return false +} + +// Component names representing the different categories of components. +type Component string + +const ( + // ComponentOfPrompt identifies chat template components. + ComponentOfPrompt Component = "ChatTemplate" + // ComponentOfAgenticPrompt identifies agentic template components. + ComponentOfAgenticPrompt Component = "AgenticChatTemplate" + // ComponentOfChatModel identifies chat model components. + ComponentOfChatModel Component = "ChatModel" + // ComponentOfAgenticModel identifies agentic model components. + ComponentOfAgenticModel Component = "AgenticModel" + // ComponentOfEmbedding identifies embedding components. + ComponentOfEmbedding Component = "Embedding" + // ComponentOfIndexer identifies indexer components. + ComponentOfIndexer Component = "Indexer" + // ComponentOfRetriever identifies retriever components. + ComponentOfRetriever Component = "Retriever" + // ComponentOfLoader identifies loader components. + ComponentOfLoader Component = "Loader" + // ComponentOfTransformer identifies document transformer components. + ComponentOfTransformer Component = "DocumentTransformer" + // ComponentOfTool identifies tool components. + ComponentOfTool Component = "Tool" +) diff --git a/compose/agentic_tools_node.go b/compose/agentic_tools_node.go new file mode 100644 index 0000000..10e430a --- /dev/null +++ b/compose/agentic_tools_node.go @@ -0,0 +1,260 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + + "github.com/cloudwego/eino/schema" +) + +// NewAgenticToolsNode creates a new AgenticToolsNode. +// e.g. +// +// conf := &ToolsNodeConfig{ +// Tools: []tool.BaseTool{invokableTool1, streamableTool2}, +// } +// toolsNode, err := NewAgenticToolsNode(ctx, conf) +func NewAgenticToolsNode(ctx context.Context, conf *ToolsNodeConfig) (*AgenticToolsNode, error) { + tn, err := NewToolNode(ctx, conf) + if err != nil { + return nil, err + } + return &AgenticToolsNode{inner: tn}, nil +} + +type AgenticToolsNode struct { + inner *ToolsNode +} + +func (a *AgenticToolsNode) Invoke(ctx context.Context, input *schema.AgenticMessage, opts ...ToolsNodeOption) ([]*schema.AgenticMessage, error) { + result, err := a.inner.Invoke(ctx, agenticMessageToToolCallMessage(input), opts...) + if err != nil { + return nil, err + } + return toolMessageToAgenticMessage(result), nil +} + +func (a *AgenticToolsNode) Stream(ctx context.Context, input *schema.AgenticMessage, + opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.AgenticMessage], error) { + result, err := a.inner.Stream(ctx, agenticMessageToToolCallMessage(input), opts...) + if err != nil { + return nil, err + } + return streamToolMessageToAgenticMessage(result), nil +} + +func agenticMessageToToolCallMessage(input *schema.AgenticMessage) *schema.Message { + var tc []schema.ToolCall + for _, block := range input.ContentBlocks { + if block.Type != schema.ContentBlockTypeFunctionToolCall || block.FunctionToolCall == nil { + continue + } + tc = append(tc, schema.ToolCall{ + ID: block.FunctionToolCall.CallID, + Function: schema.FunctionCall{ + Name: block.FunctionToolCall.Name, + Arguments: block.FunctionToolCall.Arguments, + }, + Extra: block.Extra, + }) + } + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: tc, + } +} + +func toolMessageToAgenticMessage(input []*schema.Message) []*schema.AgenticMessage { + results := make([]*schema.AgenticMessage, len(input)) + for i, m := range input { + if msg, ok := toolSearchResultMessageToAgenticMessage(m, nil); ok { + results[i] = msg + continue + } + + ftr := &schema.FunctionToolResult{ + CallID: m.ToolCallID, + Name: m.ToolName, + } + if len(m.UserInputMultiContent) > 0 { + ftr.Content = messageInputPartsToFunctionToolBlocks(m.UserInputMultiContent) + } else if m.Content != "" { + ftr.Content = []*schema.FunctionToolResultContentBlock{ + newFuncToolResultContentBlock(&schema.UserInputText{Text: m.Content}), + } + } + results[i] = &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{{ + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: ftr, + Extra: m.Extra, + }}, + Extra: m.Extra, + } + } + return results +} + +func streamToolMessageToAgenticMessage(input *schema.StreamReader[[]*schema.Message]) *schema.StreamReader[[]*schema.AgenticMessage] { + return schema.StreamReaderWithConvert(input, func(t []*schema.Message) ([]*schema.AgenticMessage, error) { + results := make([]*schema.AgenticMessage, len(t)) + for i, m := range t { + if m == nil { + continue + } + if msg, ok := toolSearchResultMessageToAgenticMessage(m, &schema.StreamingMeta{Index: i}); ok { + results[i] = msg + continue + } + + ftr := &schema.FunctionToolResult{ + CallID: m.ToolCallID, + Name: m.ToolName, + } + if len(m.UserInputMultiContent) > 0 { + ftr.Content = messageInputPartsToFunctionToolBlocks(m.UserInputMultiContent) + } else if m.Content != "" { + ftr.Content = []*schema.FunctionToolResultContentBlock{ + newFuncToolResultContentBlock(&schema.UserInputText{Text: m.Content}), + } + } + results[i] = &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{{ + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: ftr, + StreamingMeta: &schema.StreamingMeta{Index: i}, + Extra: m.Extra, + }}, + Extra: m.Extra, + } + } + return results, nil + }) +} + +func toolSearchResultMessageToAgenticMessage(m *schema.Message, meta *schema.StreamingMeta) (*schema.AgenticMessage, bool) { + if m == nil || len(m.UserInputMultiContent) != 1 { + return nil, false + } + + part := m.UserInputMultiContent[0] + if part.Type != schema.ChatMessagePartTypeToolSearchResult || part.ToolSearchResult == nil { + return nil, false + } + + block := schema.NewContentBlock(&schema.ToolSearchFunctionToolResult{ + CallID: m.ToolCallID, + Name: m.ToolName, + Result: part.ToolSearchResult, + }) + block.StreamingMeta = meta + block.Extra = m.Extra + + return &schema.AgenticMessage{ + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{block}, + Extra: m.Extra, + }, true +} + +func messageInputPartsToFunctionToolBlocks(parts []schema.MessageInputPart) []*schema.FunctionToolResultContentBlock { + blocks := make([]*schema.FunctionToolResultContentBlock, 0, len(parts)) + for _, p := range parts { + var block *schema.FunctionToolResultContentBlock + switch p.Type { + case schema.ChatMessagePartTypeText: + block = newFuncToolResultContentBlock(&schema.UserInputText{Text: p.Text}) + block.Extra = p.Extra + case schema.ChatMessagePartTypeImageURL: + if p.Image != nil { + block = newFuncToolResultContentBlock(&schema.UserInputImage{ + URL: derefString(p.Image.URL), + Base64Data: derefString(p.Image.Base64Data), + MIMEType: p.Image.MIMEType, + Detail: p.Image.Detail, + }) + block.Extra = p.Extra + } + case schema.ChatMessagePartTypeAudioURL: + if p.Audio != nil { + block = newFuncToolResultContentBlock(&schema.UserInputAudio{ + URL: derefString(p.Audio.URL), + Base64Data: derefString(p.Audio.Base64Data), + MIMEType: p.Audio.MIMEType, + }) + block.Extra = p.Extra + } + case schema.ChatMessagePartTypeVideoURL: + if p.Video != nil { + block = newFuncToolResultContentBlock(&schema.UserInputVideo{ + URL: derefString(p.Video.URL), + Base64Data: derefString(p.Video.Base64Data), + MIMEType: p.Video.MIMEType, + }) + block.Extra = p.Extra + } + case schema.ChatMessagePartTypeFileURL: + if p.File != nil { + block = newFuncToolResultContentBlock(&schema.UserInputFile{ + URL: derefString(p.File.URL), + Base64Data: derefString(p.File.Base64Data), + Name: p.File.Name, + MIMEType: p.File.MIMEType, + }) + block.Extra = p.Extra + } + } + if block != nil { + blocks = append(blocks, block) + } + } + return blocks +} + +type userInputVariant interface { + schema.UserInputText | schema.UserInputImage | schema.UserInputAudio | schema.UserInputVideo | schema.UserInputFile +} + +// newFuncToolResultContentBlock creates a FunctionToolResultContentBlock from a typed content pointer. +func newFuncToolResultContentBlock[T userInputVariant](content *T) *schema.FunctionToolResultContentBlock { + switch c := any(content).(type) { + case *schema.UserInputText: + return &schema.FunctionToolResultContentBlock{Type: schema.FunctionToolResultContentBlockTypeText, Text: c} + case *schema.UserInputImage: + return &schema.FunctionToolResultContentBlock{Type: schema.FunctionToolResultContentBlockTypeImage, Image: c} + case *schema.UserInputAudio: + return &schema.FunctionToolResultContentBlock{Type: schema.FunctionToolResultContentBlockTypeAudio, Audio: c} + case *schema.UserInputVideo: + return &schema.FunctionToolResultContentBlock{Type: schema.FunctionToolResultContentBlockTypeVideo, Video: c} + case *schema.UserInputFile: + return &schema.FunctionToolResultContentBlock{Type: schema.FunctionToolResultContentBlockTypeFile, File: c} + default: + return nil + } +} + +func derefString(s *string) string { + if s == nil { + return "" + } + return *s +} + +func (a *AgenticToolsNode) GetType() string { return "" } diff --git a/compose/agentic_tools_node_test.go b/compose/agentic_tools_node_test.go new file mode 100644 index 0000000..1f57963 --- /dev/null +++ b/compose/agentic_tools_node_test.go @@ -0,0 +1,467 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "io" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +func TestAgenticMessageToToolCallMessage(t *testing.T) { + input := &schema.AgenticMessage{ + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{ + CallID: "1", + Name: "name1", + Arguments: "arg1", + }, + }, + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{ + CallID: "2", + Name: "name2", + Arguments: "arg2", + }, + }, + { + Type: schema.ContentBlockTypeFunctionToolCall, + FunctionToolCall: &schema.FunctionToolCall{ + CallID: "3", + Name: "name3", + Arguments: "arg3", + }, + }, + }, + } + ret := agenticMessageToToolCallMessage(input) + assert.Equal(t, schema.Assistant, ret.Role) + assert.Equal(t, []schema.ToolCall{ + { + ID: "1", + Function: schema.FunctionCall{ + Name: "name1", + Arguments: "arg1", + }, + }, + { + ID: "2", + Function: schema.FunctionCall{ + Name: "name2", + Arguments: "arg2", + }, + }, + { + ID: "3", + Function: schema.FunctionCall{ + Name: "name3", + Arguments: "arg3", + }, + }, + }, ret.ToolCalls) +} + +func TestToolMessageToAgenticMessage(t *testing.T) { + t.Run("text only", func(t *testing.T) { + input := []*schema.Message{ + { + Role: schema.Tool, + Content: "content1", + ToolCallID: "1", + ToolName: "name1", + }, + { + Role: schema.Tool, + Content: "content2", + ToolCallID: "2", + ToolName: "name2", + }, + { + Role: schema.Tool, + Content: "content3", + ToolCallID: "3", + ToolName: "name3", + }, + } + ret := toolMessageToAgenticMessage(input) + assert.Equal(t, 3, len(ret)) + for i, msg := range ret { + assert.Equal(t, schema.AgenticRoleTypeUser, msg.Role) + assert.Equal(t, 1, len(msg.ContentBlocks)) + assert.Equal(t, schema.ContentBlockTypeFunctionToolResult, msg.ContentBlocks[0].Type) + ftr := msg.ContentBlocks[0].FunctionToolResult + assert.Equal(t, input[i].ToolCallID, ftr.CallID) + assert.Equal(t, input[i].ToolName, ftr.Name) + assert.Equal(t, 1, len(ftr.Content)) + assert.Equal(t, input[i].Content, ftr.Content[0].Text.Text) + } + }) + + t.Run("with multimodal content", func(t *testing.T) { + imageURL := "https://example.com/image.png" + audioBase64 := "YXVkaW9kYXRh" + videoURL := "https://example.com/video.mp4" + fileURL := "https://example.com/file.pdf" + + input := []*schema.Message{ + { + Role: schema.Tool, + Content: "text result", + ToolCallID: "1", + ToolName: "tool1", + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "hello"}, + {Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageInputImage{ + MessagePartCommon: schema.MessagePartCommon{URL: &imageURL, MIMEType: "image/png"}, + Detail: schema.ImageURLDetailHigh, + }}, + {Type: schema.ChatMessagePartTypeAudioURL, Audio: &schema.MessageInputAudio{ + MessagePartCommon: schema.MessagePartCommon{Base64Data: &audioBase64, MIMEType: "audio/wav"}, + }}, + {Type: schema.ChatMessagePartTypeVideoURL, Video: &schema.MessageInputVideo{ + MessagePartCommon: schema.MessagePartCommon{URL: &videoURL, MIMEType: "video/mp4"}, + }}, + {Type: schema.ChatMessagePartTypeFileURL, File: &schema.MessageInputFile{ + MessagePartCommon: schema.MessagePartCommon{URL: &fileURL, MIMEType: "application/pdf"}, + }}, + }, + }, + { + Role: schema.Tool, + Content: "plain result", + ToolCallID: "2", + ToolName: "tool2", + }, + } + + ret := toolMessageToAgenticMessage(input) + assert.Equal(t, 2, len(ret)) + + // first message: multimodal tool result + assert.Equal(t, schema.AgenticRoleTypeUser, ret[0].Role) + assert.Equal(t, 1, len(ret[0].ContentBlocks)) + ftr1 := ret[0].ContentBlocks[0].FunctionToolResult + assert.Equal(t, "1", ftr1.CallID) + assert.Equal(t, 5, len(ftr1.Content)) + + assert.Equal(t, "hello", ftr1.Content[0].Text.Text) + + assert.Equal(t, imageURL, ftr1.Content[1].Image.URL) + assert.Equal(t, schema.ImageURLDetailHigh, ftr1.Content[1].Image.Detail) + + assert.Equal(t, audioBase64, ftr1.Content[2].Audio.Base64Data) + + assert.Equal(t, videoURL, ftr1.Content[3].Video.URL) + + assert.Equal(t, fileURL, ftr1.Content[4].File.URL) + + // second message: text-only tool result + assert.Equal(t, schema.AgenticRoleTypeUser, ret[1].Role) + assert.Equal(t, 1, len(ret[1].ContentBlocks)) + ftr2 := ret[1].ContentBlocks[0].FunctionToolResult + assert.Equal(t, "2", ftr2.CallID) + assert.Equal(t, 1, len(ftr2.Content)) + assert.Equal(t, "plain result", ftr2.Content[0].Text.Text) + }) + + t.Run("nil media fields are skipped", func(t *testing.T) { + input := []*schema.Message{ + { + Role: schema.Tool, + Content: "result", + ToolCallID: "1", + ToolName: "tool1", + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeImageURL, Image: nil}, + {Type: schema.ChatMessagePartTypeAudioURL, Audio: nil}, + {Type: schema.ChatMessagePartTypeVideoURL, Video: nil}, + {Type: schema.ChatMessagePartTypeFileURL, File: nil}, + {Type: schema.ChatMessagePartTypeText, Text: "only text"}, + }, + }, + } + ret := toolMessageToAgenticMessage(input) + assert.Equal(t, 1, len(ret)) + ftr := ret[0].ContentBlocks[0].FunctionToolResult + assert.Equal(t, 1, len(ftr.Content)) + assert.Equal(t, "only text", ftr.Content[0].Text.Text) + }) + + t.Run("tool search result", func(t *testing.T) { + input := []*schema.Message{ + { + Role: schema.Tool, + ToolCallID: "call_1", + ToolName: "tool_search", + UserInputMultiContent: []schema.MessageInputPart{ + { + Type: schema.ChatMessagePartTypeToolSearchResult, + ToolSearchResult: &schema.ToolSearchResult{Tools: []*schema.ToolInfo{ + {Name: "dynamic_tool", Desc: "dynamic tool"}, + }}, + }, + }, + }, + } + + ret := toolMessageToAgenticMessage(input) + require.Len(t, ret, 1) + require.Len(t, ret[0].ContentBlocks, 1) + block := ret[0].ContentBlocks[0] + assert.Equal(t, schema.ContentBlockTypeToolSearchResult, block.Type) + require.NotNil(t, block.ToolSearchFunctionToolResult) + assert.Equal(t, "call_1", block.ToolSearchFunctionToolResult.CallID) + assert.Equal(t, "tool_search", block.ToolSearchFunctionToolResult.Name) + require.NotNil(t, block.ToolSearchFunctionToolResult.Result) + require.Len(t, block.ToolSearchFunctionToolResult.Result.Tools, 1) + assert.Equal(t, "dynamic_tool", block.ToolSearchFunctionToolResult.Result.Tools[0].Name) + }) +} + +func TestStreamToolMessageToAgenticMessage(t *testing.T) { + t.Run("text only", func(t *testing.T) { + testStreamToolMessageTextOnly(t) + }) + + t.Run("with multimodal content", func(t *testing.T) { + imageURL := "https://example.com/image.png" + input := schema.StreamReaderFromArray([][]*schema.Message{ + { + { + Role: schema.Tool, + Content: "result1", + ToolName: "tool1", + ToolCallID: "1", + UserInputMultiContent: []schema.MessageInputPart{ + {Type: schema.ChatMessagePartTypeText, Text: "text part"}, + {Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageInputImage{ + MessagePartCommon: schema.MessagePartCommon{URL: &imageURL}, + }}, + }, + }, + nil, + }, + { + nil, + { + Role: schema.Tool, + Content: "result2", + ToolName: "tool2", + ToolCallID: "2", + }, + }, + }) + ret := streamToolMessageToAgenticMessage(input) + var chunks [][]*schema.AgenticMessage + for { + chunk, err := ret.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + chunks = append(chunks, chunk) + } + result, err := schema.ConcatAgenticMessagesArray(chunks) + assert.NoError(t, err) + + assert.Equal(t, 2, len(result)) + + // first message: multimodal tool result (single chunk → StreamingMeta preserved) + assert.Equal(t, schema.AgenticRoleTypeUser, result[0].Role) + assert.Equal(t, 1, len(result[0].ContentBlocks)) + ftr1 := result[0].ContentBlocks[0].FunctionToolResult + assert.Equal(t, "1", ftr1.CallID) + assert.Equal(t, 2, len(ftr1.Content)) + assert.NotNil(t, ftr1.Content[0].Text) + assert.NotNil(t, ftr1.Content[1].Image) + assert.Equal(t, imageURL, ftr1.Content[1].Image.URL) + + // second message: text-only tool result (single chunk → StreamingMeta preserved) + assert.Equal(t, schema.AgenticRoleTypeUser, result[1].Role) + assert.Equal(t, 1, len(result[1].ContentBlocks)) + ftr2 := result[1].ContentBlocks[0].FunctionToolResult + assert.Equal(t, "2", ftr2.CallID) + assert.Equal(t, 1, len(ftr2.Content)) + assert.Equal(t, "result2", ftr2.Content[0].Text.Text) + }) + + t.Run("tool search result", func(t *testing.T) { + input := schema.StreamReaderFromArray([][]*schema.Message{ + { + { + Role: schema.Tool, + ToolName: "tool_search", + ToolCallID: "call_1", + UserInputMultiContent: []schema.MessageInputPart{ + { + Type: schema.ChatMessagePartTypeToolSearchResult, + ToolSearchResult: &schema.ToolSearchResult{Tools: []*schema.ToolInfo{ + {Name: "dynamic_tool", Desc: "dynamic tool"}, + }}, + }, + }, + }, + }, + }) + + ret := streamToolMessageToAgenticMessage(input) + chunk, err := ret.Recv() + require.NoError(t, err) + require.Len(t, chunk, 1) + require.Len(t, chunk[0].ContentBlocks, 1) + block := chunk[0].ContentBlocks[0] + assert.Equal(t, schema.ContentBlockTypeToolSearchResult, block.Type) + assert.Equal(t, &schema.StreamingMeta{Index: 0}, block.StreamingMeta) + require.NotNil(t, block.ToolSearchFunctionToolResult) + assert.Equal(t, "call_1", block.ToolSearchFunctionToolResult.CallID) + assert.Equal(t, "tool_search", block.ToolSearchFunctionToolResult.Name) + require.NotNil(t, block.ToolSearchFunctionToolResult.Result) + require.Len(t, block.ToolSearchFunctionToolResult.Result.Tools, 1) + assert.Equal(t, "dynamic_tool", block.ToolSearchFunctionToolResult.Result.Tools[0].Name) + }) +} + +func testStreamToolMessageTextOnly(t *testing.T) { + input := schema.StreamReaderFromArray([][]*schema.Message{ + { + { + Role: schema.Tool, + Content: "content1-1", + ToolName: "name1", + ToolCallID: "1", + }, + nil, nil, + }, + { + nil, + { + Role: schema.Tool, + Content: "content2-1", + ToolName: "name2", + ToolCallID: "2", + }, + nil, + }, + { + nil, + { + Role: schema.Tool, + Content: "content2-2", + ToolName: "name2", + ToolCallID: "2", + }, + nil, + }, + { + nil, nil, + { + Role: schema.Tool, + Content: "content3-1", + ToolName: "name3", + ToolCallID: "3", + }, + }, + { + nil, nil, + { + Role: schema.Tool, + Content: "content3-2", + ToolName: "name3", + ToolCallID: "3", + }, + }, + }) + ret := streamToolMessageToAgenticMessage(input) + var chunks [][]*schema.AgenticMessage + for { + chunk, err := ret.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + chunks = append(chunks, chunk) + } + result, err := schema.ConcatAgenticMessagesArray(chunks) + assert.NoError(t, err) + + actualStr, err := sonic.MarshalString(result) + assert.NoError(t, err) + + expected := []*schema.AgenticMessage{ + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "1", + Name: "name1", + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content1-1"}}, + }, + }, + StreamingMeta: &schema.StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "2", + Name: "name2", + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content2-1"}}, + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content2-2"}}, + }, + }, + }, + }, + }, + { + Role: schema.AgenticRoleTypeUser, + ContentBlocks: []*schema.ContentBlock{ + { + Type: schema.ContentBlockTypeFunctionToolResult, + FunctionToolResult: &schema.FunctionToolResult{ + CallID: "3", + Name: "name3", + Content: []*schema.FunctionToolResultContentBlock{ + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content3-1"}}, + {Type: schema.FunctionToolResultContentBlockTypeText, Text: &schema.UserInputText{Text: "content3-2"}}, + }, + }, + }, + }, + }, + } + + expectedStr, err := sonic.MarshalString(expected) + assert.NoError(t, err) + + assert.Equal(t, expectedStr, actualStr) +} diff --git a/compose/branch.go b/compose/branch.go new file mode 100644 index 0000000..5507d10 --- /dev/null +++ b/compose/branch.go @@ -0,0 +1,176 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +// GraphBranchCondition is the condition type for the branch. +type GraphBranchCondition[T any] func(ctx context.Context, in T) (endNode string, err error) + +// StreamGraphBranchCondition is the condition type for the stream branch. +type StreamGraphBranchCondition[T any] func(ctx context.Context, in *schema.StreamReader[T]) (endNode string, err error) + +// GraphMultiBranchCondition is the condition type for the multi choice branch. +type GraphMultiBranchCondition[T any] func(ctx context.Context, in T) (endNode map[string]bool, err error) + +// StreamGraphMultiBranchCondition is the condition type for the stream multi choice branch. +type StreamGraphMultiBranchCondition[T any] func(ctx context.Context, in *schema.StreamReader[T]) (endNodes map[string]bool, err error) + +// GraphBranch is the branch type for the graph. +// It is used to determine the next node based on the condition. +type GraphBranch struct { + invoke func(ctx context.Context, input any) (output []string, err error) + collect func(ctx context.Context, input streamReader) (output []string, err error) + inputType reflect.Type + *genericHelper + endNodes map[string]bool + idx int // used to distinguish branches in parallel + noDataFlow bool +} + +// GetEndNode returns the all end nodes of the branch. +func (gb *GraphBranch) GetEndNode() map[string]bool { + return gb.endNodes +} + +func newGraphBranch[T any](r *runnablePacker[T, []string, any], endNodes map[string]bool) *GraphBranch { + return &GraphBranch{ + invoke: func(ctx context.Context, input any) (output []string, err error) { + in, ok := input.(T) + if !ok { + // When a nil is passed as an 'any' type, its original type information is lost, + // becoming an untyped nil. This would cause type assertions to fail. + // So if the input is nil and the target type T is an interface, we need to explicitly create a nil of type T. + if input == nil && generic.TypeOf[T]().Kind() == reflect.Interface { + var i T + in = i + } else { + panic(newUnexpectedInputTypeErr(generic.TypeOf[T](), reflect.TypeOf(input))) + } + } + return r.Invoke(ctx, in) + }, + collect: func(ctx context.Context, input streamReader) (output []string, err error) { + in, ok := unpackStreamReader[T](input) + if !ok { + panic(newUnexpectedInputTypeErr(generic.TypeOf[T](), input.getType())) + } + return r.Collect(ctx, in) + }, + inputType: generic.TypeOf[T](), + genericHelper: newGenericHelper[T, T](), + endNodes: endNodes, + } +} + +// NewGraphMultiBranch creates a branch for graphs where a condition selects +// multiple end nodes; only keys present in endNodes are allowed. +func NewGraphMultiBranch[T any](condition GraphMultiBranchCondition[T], endNodes map[string]bool) *GraphBranch { + condRun := func(ctx context.Context, in T, opts ...any) ([]string, error) { + ends, err := condition(ctx, in) + if err != nil { + return nil, err + } + ret := make([]string, 0, len(ends)) + for end := range ends { + if !endNodes[end] { + return nil, fmt.Errorf("branch invocation returns unintended end node: %s", end) + } + ret = append(ret, end) + } + + return ret, nil + } + + return newGraphBranch(newRunnablePacker(condRun, nil, nil, nil, false), endNodes) +} + +// NewStreamGraphMultiBranch creates a streaming branch where a condition on +// the input stream selects multiple end nodes. +func NewStreamGraphMultiBranch[T any](condition StreamGraphMultiBranchCondition[T], + endNodes map[string]bool) *GraphBranch { + + condRun := func(ctx context.Context, in *schema.StreamReader[T], opts ...any) ([]string, error) { + ends, err := condition(ctx, in) + if err != nil { + return nil, err + } + + ret := make([]string, 0, len(ends)) + for end := range ends { + if !endNodes[end] { + return nil, fmt.Errorf("branch invocation returns unintended end node: %s", end) + } + ret = append(ret, end) + } + return ret, nil + } + + return newGraphBranch(newRunnablePacker(nil, nil, condRun, nil, false), endNodes) +} + +// NewGraphBranch creates a new graph branch. +// It is used to determine the next node based on the condition. +// e.g. +// +// condition := func(ctx context.Context, in string) (string, error) { +// // logic to determine the next node +// return "next_node_key", nil +// } +// endNodes := map[string]bool{"path01": true, "path02": true} +// branch := compose.NewGraphBranch(condition, endNodes) +// +// graph.AddBranch("key_of_node_before_branch", branch) +func NewGraphBranch[T any](condition GraphBranchCondition[T], endNodes map[string]bool) *GraphBranch { + return NewGraphMultiBranch(func(ctx context.Context, in T) (endNode map[string]bool, err error) { + ret, err := condition(ctx, in) + if err != nil { + return nil, err + } + return map[string]bool{ret: true}, nil + }, endNodes) +} + +// NewStreamGraphBranch creates a new stream graph branch. +// It is used to determine the next node based on the condition of stream input. +// e.g. +// +// condition := func(ctx context.Context, in *schema.StreamReader[T]) (string, error) { +// // logic to determine the next node. +// // to use the feature of stream, you can use the first chunk to determine the next node. +// return "next_node_key", nil +// } +// endNodes := map[string]bool{"path01": true, "path02": true} +// branch := compose.NewStreamGraphBranch(condition, endNodes) +// +// graph.AddBranch("key_of_node_before_branch", branch) +func NewStreamGraphBranch[T any](condition StreamGraphBranchCondition[T], endNodes map[string]bool) *GraphBranch { + return NewStreamGraphMultiBranch(func(ctx context.Context, in *schema.StreamReader[T]) (endNode map[string]bool, err error) { + ret, err := condition(ctx, in) + if err != nil { + return nil, err + } + return map[string]bool{ret: true}, nil + }, endNodes) +} diff --git a/compose/branch_test.go b/compose/branch_test.go new file mode 100644 index 0000000..0c13059 --- /dev/null +++ b/compose/branch_test.go @@ -0,0 +1,132 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "io" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestMultiBranch(t *testing.T) { + g := NewGraph[string, map[string]any]() + emptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil }) + err := g.AddLambdaNode("1", emptyLambda, WithOutputKey("1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", emptyLambda, WithOutputKey("2")) + assert.NoError(t, err) + err = g.AddLambdaNode("3", emptyLambda, WithOutputKey("3")) + assert.NoError(t, err) + + err = g.AddBranch(START, NewGraphMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) { + return map[string]bool{"1": true, "2": true}, nil + }, map[string]bool{"1": true, "2": true, "3": true})) + assert.NoError(t, err) + + err = g.AddEdge("1", END) + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx) + assert.NoError(t, err) + + result, err := r.Invoke(ctx, "start") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) + + streamResult, err := r.Stream(ctx, "start") + assert.NoError(t, err) + result = map[string]any{} + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + for k, v := range chunk { + result[k] = v + } + } + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) +} + +func TestStreamMultiBranch(t *testing.T) { + g := NewGraph[string, map[string]any]() + emptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil }) + err := g.AddLambdaNode("1", emptyLambda, WithOutputKey("1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", emptyLambda, WithOutputKey("2")) + assert.NoError(t, err) + err = g.AddLambdaNode("3", emptyLambda, WithOutputKey("3")) + assert.NoError(t, err) + + err = g.AddBranch(START, NewStreamGraphMultiBranch(func(ctx context.Context, in *schema.StreamReader[string]) (endNode map[string]bool, err error) { + in.Close() + return map[string]bool{"1": true, "2": true}, nil + }, map[string]bool{"1": true, "2": true, "3": true})) + assert.NoError(t, err) + + err = g.AddEdge("1", END) + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx) + assert.NoError(t, err) + + result, err := r.Invoke(ctx, "start") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) + + streamResult, err := r.Stream(ctx, "start") + assert.NoError(t, err) + result = map[string]any{} + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + for k, v := range chunk { + result[k] = v + } + } + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) +} diff --git a/compose/chain.go b/compose/chain.go new file mode 100644 index 0000000..abfa6bf --- /dev/null +++ b/compose/chain.go @@ -0,0 +1,600 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + "reflect" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/internal/gmap" + "github.com/cloudwego/eino/internal/gslice" +) + +// NewChain create a chain with input/output type. +func NewChain[I, O any](opts ...NewGraphOption) *Chain[I, O] { + ch := &Chain[I, O]{ + gg: NewGraph[I, O](opts...), + } + + ch.gg.cmp = ComponentOfChain + + return ch +} + +// Chain is a chain of components. +// Chain nodes can be parallel / branch / sequence components. +// Chain is designed to be used in a builder pattern (should Compile() before use). +// And the interface is `Chain style`, you can use it like: `chain.AppendXX(...).AppendXX(...)` +// +// Normal usage: +// 1. create a chain with input/output type: `chain := NewChain[inputType, outputType]()` +// 2. add components to chainable list: +// 2.1 add components: `chain.AppendChatTemplate(...).AppendChatModel(...).AppendToolsNode(...)` +// 2.2 add parallel or branch node if needed: `chain.AppendParallel()`, `chain.AppendBranch()` +// 3. compile: `r, err := c.Compile()` +// 4. run: +// 4.1 `one input & one output` use `r.Invoke(ctx, input)` +// 4.2 `one input & multi output chunk` use `r.Stream(ctx, input)` +// 4.3 `multi input chunk & one output` use `r.Collect(ctx, inputReader)` +// 4.4 `multi input chunk & multi output chunk` use `r.Transform(ctx, inputReader)` +// +// Using in graph or other chain: +// chain1 := NewChain[inputType, outputType]() +// graph := NewGraph[](runTypePregel) +// graph.AddGraph("key", chain1) // chain is an AnyGraph implementation +// +// // or in another chain: +// chain2 := NewChain[inputType, outputType]() +// chain2.AppendGraph(chain1) +type Chain[I, O any] struct { + err error + + gg *Graph[I, O] + + nodeIdx int + + preNodeKeys []string + + hasEnd bool +} + +// ErrChainCompiled is returned when attempting to modify a chain after it has been compiled +var ErrChainCompiled = errors.New("chain has been compiled, cannot be modified") + +// implements AnyGraph. +func (c *Chain[I, O]) compile(ctx context.Context, option *graphCompileOptions) (*composableRunnable, error) { + if err := c.addEndIfNeeded(); err != nil { + return nil, err + } + + return c.gg.compile(ctx, option) +} + +// addEndIfNeeded add END edge of the chain/graph. +// only run once when compiling. +func (c *Chain[I, O]) addEndIfNeeded() error { + if c.hasEnd { + return nil + } + + if c.err != nil { + return c.err + } + + if len(c.preNodeKeys) == 0 { + return fmt.Errorf("pre node keys not set, number of nodes in chain= %d", len(c.gg.nodes)) + } + + for _, nodeKey := range c.preNodeKeys { + err := c.gg.AddEdge(nodeKey, END) + if err != nil { + return err + } + } + + c.hasEnd = true + + return nil +} + +func (c *Chain[I, O]) getGenericHelper() *genericHelper { + return newGenericHelper[I, O]() +} + +// inputType returns the input type of the chain. +// implements AnyGraph. +func (c *Chain[I, O]) inputType() reflect.Type { + return generic.TypeOf[I]() +} + +// outputType returns the output type of the chain. +// implements AnyGraph. +func (c *Chain[I, O]) outputType() reflect.Type { + return generic.TypeOf[O]() +} + +// compositeType returns the composite type of the chain. +// implements AnyGraph. +func (c *Chain[I, O]) component() component { + return c.gg.component() +} + +// Compile to a Runnable. +// Runnable can be used directly. +// e.g. +// +// chain := NewChain[string, string]() +// r, err := chain.Compile() +// if err != nil {} +// +// r.Invoke(ctx, input) // ping => pong +// r.Stream(ctx, input) // ping => stream out +// r.Collect(ctx, inputReader) // stream in => pong +// r.Transform(ctx, inputReader) // stream in => stream out +func (c *Chain[I, O]) Compile(ctx context.Context, opts ...GraphCompileOption) (Runnable[I, O], error) { + if err := c.addEndIfNeeded(); err != nil { + return nil, err + } + + return c.gg.Compile(ctx, opts...) +} + +// AppendChatModel add a ChatModel node to the chain. +// e.g. +// +// model, err := openai.NewChatModel(ctx, config) +// if err != nil {...} +// chain.AppendChatModel(model) +func (c *Chain[I, O]) AppendChatModel(node model.BaseChatModel, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toChatModelNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendAgenticModel add a agentic.Model node to the chain. +// e.g. +// +// model, err := openai.NewAgenticModel(ctx, config) +// if err != nil {...} +// chain.AppendAgenticModel(model) +func (c *Chain[I, O]) AppendAgenticModel(node model.AgenticModel, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toAgenticModelNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendChatTemplate add a ChatTemplate node to the chain. +// eg. +// +// chatTemplate, err := prompt.FromMessages(schema.FString, &schema.Message{ +// Role: schema.System, +// Content: "You are acting as a {role}.", +// }) +// +// chain.AppendChatTemplate(chatTemplate) +func (c *Chain[I, O]) AppendChatTemplate(node prompt.ChatTemplate, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toChatTemplateNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendAgenticChatTemplate add a prompt.AgenticChatTemplate node to the chain. +// eg. +// +// chatTemplate, err := prompt.FromAgenticMessages(schema.FString, &schema.AgenticMessage{}) +// +// chain.AppendAgenticChatTemplate(chatTemplate) +func (c *Chain[I, O]) AppendAgenticChatTemplate(node prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toAgenticChatTemplateNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendToolsNode add a ToolsNode node to the chain. +// e.g. +// +// toolsNode, err := compose.NewToolNode(ctx, &compose.ToolsNodeConfig{ +// Tools: []tools.BaseTool{...}, +// }) +// +// chain.AppendToolsNode(toolsNode) +func (c *Chain[I, O]) AppendToolsNode(node *ToolsNode, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toToolsNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendAgenticToolsNode add a AgenticToolsNode node to the chain. +// e.g. +// +// toolsNode, err := compose.NewAgenticToolsNode(ctx, &compose.ToolsNodeConfig{ +// Tools: []tools.BaseTool{...}, +// }) +// +// chain.AppendAgenticToolsNode(toolsNode) +func (c *Chain[I, O]) AppendAgenticToolsNode(node *AgenticToolsNode, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toAgenticToolsNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendDocumentTransformer add a DocumentTransformer node to the chain. +// e.g. +// +// markdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{}) +// +// chain.AppendDocumentTransformer(markdownSplitter) +func (c *Chain[I, O]) AppendDocumentTransformer(node document.Transformer, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toDocumentTransformerNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendLambda add a Lambda node to the chain. +// Lambda is a node that can be used to implement custom logic. +// e.g. +// +// lambdaNode := compose.InvokableLambda(func(ctx context.Context, docs []*schema.Document) (string, error) {...}) +// chain.AppendLambda(lambdaNode) +// +// Note: +// to create a Lambda node, you need to use `compose.AnyLambda` or `compose.InvokableLambda` or `compose.StreamableLambda` or `compose.TransformableLambda`. +// if you want this node has real stream output, you need to use `compose.StreamableLambda` or `compose.TransformableLambda`, for example. +func (c *Chain[I, O]) AppendLambda(node *Lambda, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toLambdaNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendEmbedding add a Embedding node to the chain. +// e.g. +// +// embedder, err := openai.NewEmbedder(ctx, config) +// if err != nil {...} +// chain.AppendEmbedding(embedder) +func (c *Chain[I, O]) AppendEmbedding(node embedding.Embedder, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toEmbeddingNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendRetriever add a Retriever node to the chain. +// e.g. +// +// retriever, err := vectorstore.NewRetriever(ctx, config) +// if err != nil {...} +// chain.AppendRetriever(retriever) +// +// or using fornax knowledge as retriever: +// +// config := fornaxknowledge.Config{...} +// retriever, err := fornaxknowledge.NewKnowledgeRetriever(ctx, config) +// if err != nil {...} +// chain.AppendRetriever(retriever) +func (c *Chain[I, O]) AppendRetriever(node retriever.Retriever, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toRetrieverNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendLoader adds a Loader node to the chain. +// e.g. +// +// loader, err := file.NewFileLoader(ctx, &file.FileLoaderConfig{}) +// if err != nil {...} +// chain.AppendLoader(loader) +func (c *Chain[I, O]) AppendLoader(node document.Loader, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toLoaderNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendIndexer add an Indexer node to the chain. +// Indexer is a node that can store documents. +// e.g. +// +// vectorStoreImpl, err := vikingdb.NewVectorStorer(ctx, vikingdbConfig) // in components/vectorstore/vikingdb/vectorstore.go +// if err != nil {...} +// +// config := vectorstore.IndexerConfig{VectorStore: vectorStoreImpl} +// indexer, err := vectorstore.NewIndexer(ctx, config) +// if err != nil {...} +// +// chain.AppendIndexer(indexer) +func (c *Chain[I, O]) AppendIndexer(node indexer.Indexer, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toIndexerNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendBranch add a conditional branch to chain. +// Each branch within the ChainBranch can be an AnyGraph. +// All branches should either lead to END, or converge to another node within the Chain. +// e.g. +// +// cb := compose.NewChainBranch(conditionFunc) +// cb.AddChatTemplate("chat_template_key_01", chatTemplate) +// cb.AddChatTemplate("chat_template_key_02", chatTemplate2) +// chain.AppendBranch(cb) +func (c *Chain[I, O]) AppendBranch(b *ChainBranch) *Chain[I, O] { + if b == nil { + c.reportError(fmt.Errorf("append branch invalid, branch is nil")) + return c + } + + if b.err != nil { + c.reportError(fmt.Errorf("append branch error: %w", b.err)) + return c + } + + if len(b.key2BranchNode) == 0 { + c.reportError(fmt.Errorf("append branch invalid, nodeList is empty")) + return c + } + + if len(b.key2BranchNode) == 1 { + c.reportError(fmt.Errorf("append branch invalid, nodeList length = 1")) + return c + } + + var startNode string + if len(c.preNodeKeys) == 0 { // branch appended directly to START + startNode = START + } else if len(c.preNodeKeys) == 1 { + startNode = c.preNodeKeys[0] + } else { + c.reportError(fmt.Errorf("append branch invalid, multiple previous nodes: %v ", c.preNodeKeys)) + return c + } + + prefix := c.nextNodeKey() + key2NodeKey := make(map[string]string, len(b.key2BranchNode)) + + for key := range b.key2BranchNode { + node := b.key2BranchNode[key] + + var nodeKey string + + if node.Second != nil && node.Second.nodeOptions != nil && node.Second.nodeOptions.nodeKey != "" { + nodeKey = node.Second.nodeOptions.nodeKey + } else { + nodeKey = fmt.Sprintf("%s_branch_%s", prefix, key) + } + + if err := c.gg.addNode(nodeKey, node.First, node.Second); err != nil { + c.reportError(fmt.Errorf("add branch node[%s] to chain failed: %w", nodeKey, err)) + return c + } + + key2NodeKey[key] = nodeKey + } + + gBranch := *b.internalBranch + + invokeCon := func(ctx context.Context, in any) (endNode []string, err error) { + ends, err := b.internalBranch.invoke(ctx, in) + if err != nil { + return nil, err + } + + nodeKeyEnds := make([]string, 0, len(ends)) + for _, end := range ends { + if nodeKey, ok := key2NodeKey[end]; !ok { + return nil, fmt.Errorf("branch invocation returns unintended end node: %s", end) + } else { + nodeKeyEnds = append(nodeKeyEnds, nodeKey) + } + } + + return nodeKeyEnds, nil + } + gBranch.invoke = invokeCon + + collectCon := func(ctx context.Context, sr streamReader) ([]string, error) { + ends, err := b.internalBranch.collect(ctx, sr) + if err != nil { + return nil, err + } + + nodeKeyEnds := make([]string, 0, len(ends)) + for _, end := range ends { + if nodeKey, ok := key2NodeKey[end]; !ok { + return nil, fmt.Errorf("branch invocation returns unintended end node: %s", end) + } else { + nodeKeyEnds = append(nodeKeyEnds, nodeKey) + } + } + + return nodeKeyEnds, nil + } + gBranch.collect = collectCon + + gBranch.endNodes = gslice.ToMap(gmap.Values(key2NodeKey), func(k string) (string, bool) { + return k, true + }) + + if err := c.gg.AddBranch(startNode, &gBranch); err != nil { + c.reportError(fmt.Errorf("chain append branch failed: %w", err)) + return c + } + + c.preNodeKeys = gmap.Values(key2NodeKey) + + return c +} + +// AppendParallel add a Parallel structure (multiple concurrent nodes) to the chain. +// e.g. +// +// parallel := compose.NewParallel() +// parallel.AddChatModel("openai", model1) // => "openai": *schema.Message{} +// parallel.AddChatModel("maas", model2) // => "maas": *schema.Message{} +// +// chain.AppendParallel(parallel) // => multiple concurrent nodes are added to the Chain +// +// The next node in the chain is either an END, or a node which accepts a map[string]any, where keys are `openai` `maas` as specified above. +func (c *Chain[I, O]) AppendParallel(p *Parallel) *Chain[I, O] { + if p == nil { + c.reportError(fmt.Errorf("append parallel invalid, parallel is nil")) + return c + } + + if p.err != nil { + c.reportError(fmt.Errorf("append parallel invalid, parallel error: %w", p.err)) + return c + } + + if len(p.nodes) <= 1 { + c.reportError(fmt.Errorf("append parallel invalid, not enough nodes, count = %d", len(p.nodes))) + return c + } + + var startNode string + if len(c.preNodeKeys) == 0 { // parallel appended directly to START + startNode = START + } else if len(c.preNodeKeys) == 1 { + startNode = c.preNodeKeys[0] + } else { + c.reportError(fmt.Errorf("append parallel invalid, multiple previous nodes: %v ", c.preNodeKeys)) + return c + } + + prefix := c.nextNodeKey() + var nodeKeys []string + + for i := range p.nodes { + node := p.nodes[i] + + var nodeKey string + if node.Second != nil && node.Second.nodeOptions != nil && node.Second.nodeOptions.nodeKey != "" { + nodeKey = node.Second.nodeOptions.nodeKey + } else { + nodeKey = fmt.Sprintf("%s_parallel_%d", prefix, i) + } + + if err := c.gg.addNode(nodeKey, node.First, node.Second); err != nil { + c.reportError(fmt.Errorf("add parallel node to chain failed, key=%s, err: %w", nodeKey, err)) + return c + } + + if err := c.gg.AddEdge(startNode, nodeKey); err != nil { + c.reportError(fmt.Errorf("add parallel edge failed, from=%s, to=%s, err: %w", startNode, nodeKey, err)) + return c + } + + nodeKeys = append(nodeKeys, nodeKey) + } + + c.preNodeKeys = nodeKeys + + return c +} + +// AppendGraph add a AnyGraph node to the chain. +// AnyGraph can be a chain or a graph. +// e.g. +// +// graph := compose.NewGraph[string, string]() +// chain.AppendGraph(graph) +func (c *Chain[I, O]) AppendGraph(node AnyGraph, opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toAnyGraphNode(node, opts...) + c.addNode(gNode, options) + return c +} + +// AppendPassthrough add a Passthrough node to the chain. +// Could be used to connect multiple ChainBranch or Parallel. +// e.g. +// +// chain.AppendPassthrough() +func (c *Chain[I, O]) AppendPassthrough(opts ...GraphAddNodeOpt) *Chain[I, O] { + gNode, options := toPassthroughNode(opts...) + c.addNode(gNode, options) + return c +} + +// nextIdx. +// get the next idx for the chain. +// chain key is: node_idx => eg: node_0 => represent the first node of the chain (idx start from 0) +// if has parallel: node_idx_parallel_idx => eg: node_0_parallel_1 => represent the first node of the chain, and is a parallel node, and the second node of the parallel +// if has branch: node_idx_branch_key => eg: node_1_branch_customkey => represent the second node of the chain, and is a branch node, and the 'customkey' is the key of the branch +func (c *Chain[I, O]) nextNodeKey() string { + idx := c.nodeIdx + c.nodeIdx++ + return fmt.Sprintf("node_%d", idx) +} + +// reportError. +// save the first error in the chain. +func (c *Chain[I, O]) reportError(err error) { + if c.err == nil { + c.err = err + } +} + +// addNode. +// add a node to the chain. +func (c *Chain[I, O]) addNode(node *graphNode, options *graphAddNodeOpts) { + if c.err != nil { + return + } + + if c.gg.compiled { + c.reportError(ErrChainCompiled) + return + } + + if node == nil { + c.reportError(fmt.Errorf("chain add node invalid, node is nil")) + return + } + + nodeKey := options.nodeOptions.nodeKey + defaultNodeKey := c.nextNodeKey() + if nodeKey == "" { + nodeKey = defaultNodeKey + } + + err := c.gg.addNode(nodeKey, node, options) + if err != nil { + c.reportError(err) + return + } + + if len(c.preNodeKeys) == 0 { + c.preNodeKeys = append(c.preNodeKeys, START) + } + + for _, preNodeKey := range c.preNodeKeys { + e := c.gg.AddEdge(preNodeKey, nodeKey) + if e != nil { + c.reportError(e) + return + } + } + + c.preNodeKeys = []string{nodeKey} +} diff --git a/compose/chain_branch.go b/compose/chain_branch.go new file mode 100644 index 0000000..84fb110 --- /dev/null +++ b/compose/chain_branch.go @@ -0,0 +1,343 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +type nodeOptionsPair generic.Pair[*graphNode, *graphAddNodeOpts] + +// ChainBranch represents a conditional branch in a chain of operations. +// It allows for dynamic routing of execution based on a condition. +// All branches within ChainBranch are expected to either end the Chain, or converge to another node in the Chain. +type ChainBranch struct { + internalBranch *GraphBranch + key2BranchNode map[string]nodeOptionsPair + err error +} + +// NewChainMultiBranch creates a chain branch where a condition selects +// multiple end nodes to route execution. +func NewChainMultiBranch[T any](cond GraphMultiBranchCondition[T]) *ChainBranch { + invokeCond := func(ctx context.Context, in T, opts ...any) (endNodes []string, err error) { + ends, err := cond(ctx, in) + if err != nil { + return nil, err + } + endNodes = make([]string, 0, len(ends)) + for end := range ends { + endNodes = append(endNodes, end) + } + return endNodes, nil + } + + return &ChainBranch{ + key2BranchNode: make(map[string]nodeOptionsPair), + internalBranch: newGraphBranch(newRunnablePacker(invokeCond, nil, nil, nil, false), nil), + } +} + +// NewStreamChainMultiBranch creates a chain branch that selects multiple end +// nodes based on a condition evaluated on the input stream. +func NewStreamChainMultiBranch[T any](cond StreamGraphMultiBranchCondition[T]) *ChainBranch { + collectCon := func(ctx context.Context, in *schema.StreamReader[T], opts ...any) (endNodes []string, err error) { + ends, err := cond(ctx, in) + if err != nil { + return nil, err + } + endNodes = make([]string, 0, len(ends)) + for end := range ends { + endNodes = append(endNodes, end) + } + return endNodes, nil + } + + return &ChainBranch{ + key2BranchNode: make(map[string]nodeOptionsPair), + internalBranch: newGraphBranch(newRunnablePacker(nil, nil, collectCon, nil, false), nil), + } +} + +// NewChainBranch creates a new ChainBranch instance based on a given condition. +// It takes a generic type T and a GraphBranchCondition function for that type. +// The returned ChainBranch will have an empty key2BranchNode map and a condition function +// that wraps the provided cond to handle type assertions and error checking. +// eg. +// +// condition := func(ctx context.Context, in string, opts ...any) (endNode string, err error) { +// // logic to determine the next node +// return "some_next_node_key", nil +// } +// +// cb := NewChainBranch[string](condition) +// cb.AddPassthrough("next_node_key_01", xxx) // node in branch, represent one path of branch +// cb.AddPassthrough("next_node_key_02", xxx) // node in branch +func NewChainBranch[T any](cond GraphBranchCondition[T]) *ChainBranch { + return NewChainMultiBranch(func(ctx context.Context, in T) (endNode map[string]bool, err error) { + ret, err := cond(ctx, in) + if err != nil { + return nil, err + } + return map[string]bool{ret: true}, nil + }) +} + +// NewStreamChainBranch creates a new ChainBranch instance based on a given stream condition. +// It takes a generic type T and a StreamGraphBranchCondition function for that type. +// The returned ChainBranch will have an empty key2BranchNode map and a condition function +// that wraps the provided cond to handle type assertions and error checking. +// eg. +// +// condition := func(ctx context.Context, in *schema.StreamReader[string], opts ...any) (endNode string, err error) { +// // logic to determine the next node, you can read the stream and make a decision. +// // to save time, usually read the first chunk of stream, then make a decision which path to go. +// return "some_next_node_key", nil +// } +// +// cb := NewStreamChainBranch[string](condition) +func NewStreamChainBranch[T any](cond StreamGraphBranchCondition[T]) *ChainBranch { + return NewStreamChainMultiBranch(func(ctx context.Context, in *schema.StreamReader[T]) (endNodes map[string]bool, err error) { + ret, err := cond(ctx, in) + if err != nil { + return nil, err + } + return map[string]bool{ret: true}, nil + }) +} + +// AddChatModel adds a ChatModel node to the branch. +// eg. +// +// chatModel01, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{ +// Model: "gpt-4o", +// }) +// chatModel02, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{ +// Model: "gpt-4o-mini", +// }) +// cb.AddChatModel("chat_model_key_01", chatModel01) +// cb.AddChatModel("chat_model_key_02", chatModel02) +func (cb *ChainBranch) AddChatModel(key string, node model.BaseChatModel, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toChatModelNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddAgenticModel adds a agentic.Model node to the branch. +// eg. +// +// model1, err := openai.NewAgenticModel(ctx, &openai.AgenticModelConfig{ +// Model: "gpt-4o", +// }) +// model2, err := openai.NewAgenticModel(ctx, &openai.AgenticModelConfig{ +// Model: "gpt-4o-mini", +// }) +// cb.AddAgenticModel("agentic_model_key_1", model1) +// cb.AddAgenticModel("agentic_model_key_2", model2) +func (cb *ChainBranch) AddAgenticModel(key string, node model.AgenticModel, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toAgenticModelNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddChatTemplate adds a ChatTemplate node to the branch. +// eg. +// +// chatTemplate, err := prompt.FromMessages(schema.FString, &schema.Message{ +// Role: schema.System, +// Content: "You are acting as a {role}.", +// }) +// +// cb.AddChatTemplate("chat_template_key_01", chatTemplate) +// +// chatTemplate2, err := prompt.FromMessages(schema.FString, &schema.Message{ +// Role: schema.System, +// Content: "You are acting as a {role}, you are not allowed to chat in other topics.", +// }) +// +// cb.AddChatTemplate("chat_template_key_02", chatTemplate2) +func (cb *ChainBranch) AddChatTemplate(key string, node prompt.ChatTemplate, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toChatTemplateNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddAgenticChatTemplate adds a prompt.AgenticChatTemplate node to the branch. +// eg. +// +// chatTemplate, err := prompt.FromAgenticMessages(schema.FString, &schema.AgenticMessage{}) +// +// cb.AddAgenticChatTemplate("chat_template_key_01", chatTemplate) +// +// chatTemplate2, err := prompt.FromAgenticMessages(schema.FString, &schema.AgenticMessage{}) +// +// cb.AddAgenticChatTemplate("chat_template_key_02", chatTemplate2) +func (cb *ChainBranch) AddAgenticChatTemplate(key string, node prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toAgenticChatTemplateNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddToolsNode adds a ToolsNode to the branch. +// eg. +// +// toolsNode, err := compose.NewToolNode(ctx, &compose.ToolsNodeConfig{ +// Tools: []tools.BaseTool{...}, +// }) +// +// cb.AddToolsNode("tools_node_key", toolsNode) +func (cb *ChainBranch) AddToolsNode(key string, node *ToolsNode, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toToolsNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddAgenticToolsNode adds a AgenticToolsNode to the branch. +// eg. +// +// toolsNode, err := compose.NewAgenticToolsNode(ctx, &compose.ToolsNodeConfig{ +// Tools: []tools.BaseTool{...}, +// }) +// +// cb.AddAgenticToolsNode("tools_node_key", toolsNode) +func (cb *ChainBranch) AddAgenticToolsNode(key string, node *AgenticToolsNode, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toAgenticToolsNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddLambda adds a Lambda node to the branch. +// eg. +// +// lambdaFunc := func(ctx context.Context, in string, opts ...any) (out string, err error) { +// // logic to process the input +// return "processed_output", nil +// } +// +// cb.AddLambda("lambda_node_key", compose.InvokeLambda(lambdaFunc)) +func (cb *ChainBranch) AddLambda(key string, node *Lambda, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toLambdaNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddEmbedding adds an Embedding node to the branch. +// eg. +// +// embeddingNode, err := openai.NewEmbedder(ctx, &openai.EmbeddingConfig{ +// Model: "text-embedding-3-small", +// }) +// +// cb.AddEmbedding("embedding_node_key", embeddingNode) +func (cb *ChainBranch) AddEmbedding(key string, node embedding.Embedder, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toEmbeddingNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddRetriever adds a Retriever node to the branch. +// eg. +// +// retriever, err := volc_vikingdb.NewRetriever(ctx, &volc_vikingdb.RetrieverConfig{ +// Collection: "my_collection", +// }) +// +// cb.AddRetriever("retriever_node_key", retriever) +func (cb *ChainBranch) AddRetriever(key string, node retriever.Retriever, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toRetrieverNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddLoader adds a Loader node to the branch. +// eg. +// +// pdfParser, err := pdf.NewPDFParser() +// loader, err := file.NewFileLoader(ctx, &file.FileLoaderConfig{ +// Parser: pdfParser, +// }) +// +// cb.AddLoader("loader_node_key", loader) +func (cb *ChainBranch) AddLoader(key string, node document.Loader, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toLoaderNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddIndexer adds an Indexer node to the branch. +// eg. +// +// indexer, err := volc_vikingdb.NewIndexer(ctx, &volc_vikingdb.IndexerConfig{ +// Collection: "my_collection", +// }) +// +// cb.AddIndexer("indexer_node_key", indexer) +func (cb *ChainBranch) AddIndexer(key string, node indexer.Indexer, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toIndexerNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddDocumentTransformer adds an Document Transformer node to the branch. +// eg. +// +// markdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{}) +// +// cb.AddDocumentTransformer("document_transformer_node_key", markdownSplitter) +func (cb *ChainBranch) AddDocumentTransformer(key string, node document.Transformer, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toDocumentTransformerNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddGraph adds a generic Graph node to the branch. +// eg. +// +// graph, err := compose.NewGraph[string, string]() +// +// cb.AddGraph("graph_node_key", graph) +func (cb *ChainBranch) AddGraph(key string, node AnyGraph, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toAnyGraphNode(node, opts...) + return cb.addNode(key, gNode, options) +} + +// AddPassthrough adds a Passthrough node to the branch. +// eg. +// +// cb.AddPassthrough("passthrough_node_key") +func (cb *ChainBranch) AddPassthrough(key string, opts ...GraphAddNodeOpt) *ChainBranch { + gNode, options := toPassthroughNode(opts...) + return cb.addNode(key, gNode, options) +} + +func (cb *ChainBranch) addNode(key string, node *graphNode, options *graphAddNodeOpts) *ChainBranch { + if cb.err != nil { + return cb + } + + if cb.key2BranchNode == nil { + cb.key2BranchNode = make(map[string]nodeOptionsPair) + } + + _, ok := cb.key2BranchNode[key] + if ok { + cb.err = fmt.Errorf("chain branch add node, duplicate branch node key= %s", key) + return cb + } + + cb.key2BranchNode[key] = nodeOptionsPair{node, options} + + return cb +} diff --git a/compose/chain_branch_test.go b/compose/chain_branch_test.go new file mode 100644 index 0000000..6b3b038 --- /dev/null +++ b/compose/chain_branch_test.go @@ -0,0 +1,348 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/schema" +) + +func TestChainBranch(t *testing.T) { + cond := func(ctx context.Context, input string) (key string, err error) { + switch input { + case "one": + return "one_key", nil + case "two": + return "two_key", nil + case "three": + return "three_key", nil + default: + return "", fmt.Errorf("invalid input= %s", input) + } + } + + t.Run("nested chain", func(t *testing.T) { + inner := NewChain[string, string]() + inner.AppendBranch(NewChainBranch(cond). + AddLambda("one_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in, nil + })). + AddLambda("two_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in + in, nil + }))) + inner.AppendParallel(NewParallel(). + AddLambda("one_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in, nil + })). + AddLambda("two_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in + in, nil + }))) + + outer := NewChain[string, string]() + outer.AppendGraph(inner) + _, err := outer.Compile(context.Background()) + assert.Error(t, err) + }) + + t.Run("bad param", func(t *testing.T) { + c := NewChain[string, string]() + c.AppendBranch(nil) + assert.NotNil(t, c.err) + + c = NewChain[string, string]() + c.AppendBranch(NewChainBranch[string](nil)) + assert.NotNil(t, c.err) + + c = NewChain[string, string]() + c.AppendBranch(NewChainBranch(cond).AddChatTemplate("template", prompt.FromMessages(schema.FString, schema.SystemMessage("hello")))) + assert.NotNil(t, c.err) + + c = NewChain[string, string]() + c.AppendBranch(NewChainBranch(cond).AddChatTemplate("1", prompt.FromMessages(schema.FString)).AddChatTemplate("1", prompt.FromMessages(schema.FString))) + assert.NotNil(t, c.err) + }) + + t.Run("different Node types in branch", func(t *testing.T) { + c := NewChain[string, string]() + c.AppendBranch(NewChainBranch(cond). + AddChatTemplate("t", prompt.FromMessages(schema.FString)). + AddGraph("c", NewChain[string, string]())) + assert.NotNil(t, c.err) + }) + + t.Run("type mismatch", func(t *testing.T) { + c := NewChain[int, string]() + c.AppendBranch(NewChainBranch(cond). + AddLambda("one_key", InvokableLambda(func(ctx context.Context, in int) (output string, err error) { + return strconv.Itoa(in), nil + })). + AddLambda("two_key", InvokableLambda(func(ctx context.Context, in int) (output string, err error) { + return strconv.Itoa(in), nil + }))) + _, err := c.Compile(context.Background()) + assert.NotNil(t, err) + }) + + t.Run("invoke", func(t *testing.T) { + c := NewChain[string, string]() + c.AppendBranch(NewChainBranch(cond). + AddLambda("one_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in, nil + })). + AddLambda("two_key", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in + in, nil + }))) + c.AppendLambda(InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + in, nil + })) + assert.Nil(t, c.err) + compiledChain, err := c.Compile(context.Background()) + assert.Nil(t, err) + + out, err := compiledChain.Invoke(context.Background(), "two") + assert.Nil(t, err) + assert.Equal(t, "twotwotwotwotwotwo", out) + + _, err = compiledChain.Invoke(context.Background(), "three") + assert.NotNil(t, err) + + _, err = compiledChain.Invoke(context.Background(), "four") + assert.NotNil(t, err) + }) + + t.Run("fake stream", func(t *testing.T) { + c := NewChain[string, string]() + c.AppendLambda(StreamableLambda(func(ctx context.Context, in string) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](utf8.RuneCountInString(in)) + + go func() { + for _, field := range strings.Fields(in) { + sw.Send(field, nil) + } + sw.Close() + }() + + return sr, nil + })) + c.AppendBranch(NewChainBranch[string](cond).AddLambda("one_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) { + defer in.Close() + for { + v, err := in.Recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + return "", err + } + + output += v + } + + return output + output, nil + })). + AddLambda("two_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) { + defer in.Close() + for { + v, err := in.Recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + return "", err + } + + output += v + } + + return output + output + output, nil + }))) + + assert.Nil(t, c.err) + compiledChain, err := c.Compile(context.Background()) + assert.Nil(t, err) + + out, err := compiledChain.Invoke(context.Background(), "one") + assert.Nil(t, err) + assert.Equal(t, "oneone", out) + }) + + t.Run("real stream", func(t *testing.T) { + streamCon := func(ctx context.Context, sr *schema.StreamReader[string]) (key string, err error) { + msg, err := sr.Recv() + if err != nil { + return "", err + } + defer sr.Close() + + switch msg { + case "one": + return "one_key", nil + case "two": + return "two_key", nil + case "three": + return "three_key", nil + default: + return "", fmt.Errorf("invalid input= %s", msg) + } + } + + c := NewChain[string, string]() + c.AppendLambda(StreamableLambda(func(ctx context.Context, in string) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](utf8.RuneCountInString(in)) + + go func() { + for _, field := range strings.Fields(in) { + sw.Send(field, nil) + } + sw.Close() + }() + + return sr, nil + })) + c.AppendBranch(NewStreamChainBranch(streamCon).AddLambda("one_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) { + defer in.Close() + for { + v, err := in.Recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + return "", err + } + + output += v + } + + return output + output, nil + })). + AddLambda("two_key", CollectableLambda(func(ctx context.Context, in *schema.StreamReader[string]) (output string, err error) { + defer in.Close() + for { + v, err := in.Recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + return "", err + } + + output += v + } + + return output + output + output, nil + }))) + + assert.Nil(t, c.err) + compiledChain, err := c.Compile(context.Background()) + assert.Nil(t, err) + + out, err := compiledChain.Stream(context.Background(), "one size fit all") + assert.Nil(t, err) + concat, err := concatStreamReader(out) + assert.Nil(t, err) + assert.Equal(t, "onesizefitallonesizefitall", concat) + }) +} + +func TestChainMultiBranch(t *testing.T) { + emptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil }) + + ctx := context.Background() + r, err := NewChain[string, map[string]any](). + AppendBranch(NewChainMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) { + return map[string]bool{"1": true, "2": true}, nil + }).AddLambda("1", emptyLambda, WithOutputKey("1")).AddLambda("2", emptyLambda, WithOutputKey("2")).AddLambda("3", emptyLambda, WithOutputKey("3"))). + Compile(ctx) + assert.Nil(t, err) + + result, err := r.Invoke(ctx, "start") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) + + streamResult, err := r.Stream(ctx, "start") + assert.NoError(t, err) + result = map[string]any{} + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + for k, v := range chunk { + result[k] = v + } + } + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) +} + +func TestStreamChainMultiBranch(t *testing.T) { + emptyLambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input, nil }) + + ctx := context.Background() + r, err := NewChain[string, map[string]any](). + AppendBranch(NewStreamChainMultiBranch(func(ctx context.Context, in *schema.StreamReader[string]) (endNode map[string]bool, err error) { + return map[string]bool{"1": true, "2": true}, nil + }).AddLambda("1", emptyLambda, WithOutputKey("1")).AddLambda("2", emptyLambda, WithOutputKey("2")).AddLambda("3", emptyLambda, WithOutputKey("3"))). + Compile(ctx) + assert.Nil(t, err) + + result, err := r.Invoke(ctx, "start") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) + + streamResult, err := r.Stream(ctx, "start") + assert.NoError(t, err) + result = map[string]any{} + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + for k, v := range chunk { + result[k] = v + } + } + assert.Equal(t, map[string]any{ + "1": "start", + "2": "start", + }, result) +} diff --git a/compose/chain_parallel.go b/compose/chain_parallel.go new file mode 100644 index 0000000..463140b --- /dev/null +++ b/compose/chain_parallel.go @@ -0,0 +1,263 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "fmt" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" +) + +// NewParallel creates a new parallel type. +// it is useful when you want to run multiple nodes in parallel in a chain. +func NewParallel() *Parallel { + return &Parallel{ + outputKeys: make(map[string]bool), + } +} + +// Parallel run multiple nodes in parallel +// +// use `NewParallel()` to create a new parallel type +// Example: +// +// parallel := NewParallel() +// parallel.AddChatModel("output_key01", chat01) +// parallel.AddChatModel("output_key01", chat02) +// +// chain := NewChain[any,any]() +// chain.AppendParallel(parallel) +type Parallel struct { + nodes []nodeOptionsPair + outputKeys map[string]bool + err error +} + +// AddChatModel adds a chat model to the parallel. +// eg. +// +// chatModel01, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{ +// Model: "gpt-4o", +// }) +// +// chatModel02, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{ +// Model: "gpt-4o", +// }) +// +// p.AddChatModel("output_key01", chatModel01) +// p.AddChatModel("output_key02", chatModel02) +func (p *Parallel) AddChatModel(outputKey string, node model.BaseChatModel, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toChatModelNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddAgenticModel adds a agentic.Model to the parallel. +// eg. +// +// model1, err := openai.NewAgenticModel(ctx, &openai.AgenticModelConfig{ +// Model: "gpt-4o", +// }) +// +// model2, err := openai.NewAgenticModel(ctx, &openai.AgenticModelConfig{ +// Model: "gpt-4o", +// }) +// +// p.AddAgenticModel("output_key1", model1) +// p.AddAgenticModel("output_key2", model2) +func (p *Parallel) AddAgenticModel(outputKey string, node model.AgenticModel, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toAgenticModelNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddChatTemplate adds a chat template to the parallel. +// eg. +// +// chatTemplate01, err := prompt.FromMessages(schema.FString, &schema.Message{ +// Role: schema.System, +// Content: "You are acting as a {role}.", +// }) +// +// p.AddChatTemplate("output_key01", chatTemplate01) +func (p *Parallel) AddChatTemplate(outputKey string, node prompt.ChatTemplate, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toChatTemplateNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddAgenticChatTemplate adds a prompt.AgenticChatTemplate to the parallel. +// eg. +// +// chatTemplate01, err := prompt.FromAgenticMessages(schema.FString, &schema.AgenticMessage{}) +// +// p.AddAgenticChatTemplate("output_key01", chatTemplate01) +func (p *Parallel) AddAgenticChatTemplate(outputKey string, node prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toAgenticChatTemplateNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddToolsNode adds a tools node to the parallel. +// eg. +// +// toolsNode, err := compose.NewToolNode(ctx, &compose.ToolsNodeConfig{ +// Tools: []tool.BaseTool{...}, +// }) +// +// p.AddToolsNode("output_key01", toolsNode) +func (p *Parallel) AddToolsNode(outputKey string, node *ToolsNode, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toToolsNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddAgenticToolsNode adds a tools node to the parallel. +// eg. +// +// toolsNode, err := compose.NewAgenticToolsNode(ctx, &compose.ToolsNodeConfig{ +// Tools: []tool.BaseTool{...}, +// }) +// +// p.AddAgenticToolsNode("output_key01", toolsNode) +func (p *Parallel) AddAgenticToolsNode(outputKey string, node *AgenticToolsNode, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toAgenticToolsNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddLambda adds a lambda node to the parallel. +// eg. +// +// lambdaFunc := func(ctx context.Context, input *schema.Message) ([]*schema.Message, error) { +// return []*schema.Message{input}, nil +// } +// +// p.AddLambda("output_key01", compose.InvokeLambda(lambdaFunc)) +func (p *Parallel) AddLambda(outputKey string, node *Lambda, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toLambdaNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddEmbedding adds an embedding node to the parallel. +// eg. +// +// embeddingNode, err := openai.NewEmbedder(ctx, &openai.EmbeddingConfig{ +// Model: "text-embedding-3-small", +// }) +// +// p.AddEmbedding("output_key01", embeddingNode) +func (p *Parallel) AddEmbedding(outputKey string, node embedding.Embedder, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toEmbeddingNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddRetriever adds a retriever node to the parallel. +// eg. +// +// retriever, err := vikingdb.NewRetriever(ctx, &vikingdb.RetrieverConfig{}) +// +// p.AddRetriever("output_key01", retriever) +func (p *Parallel) AddRetriever(outputKey string, node retriever.Retriever, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toRetrieverNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddLoader adds a loader node to the parallel. +// eg. +// +// loader, err := file.NewLoader(ctx, &file.LoaderConfig{}) +// +// p.AddLoader("output_key01", loader) +func (p *Parallel) AddLoader(outputKey string, node document.Loader, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toLoaderNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddIndexer adds an indexer node to the parallel. +// eg. +// +// indexer, err := volc_vikingdb.NewIndexer(ctx, &volc_vikingdb.IndexerConfig{ +// Collection: "my_collection", +// }) +// +// p.AddIndexer("output_key01", indexer) +func (p *Parallel) AddIndexer(outputKey string, node indexer.Indexer, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toIndexerNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddDocumentTransformer adds an Document Transformer node to the parallel. +// eg. +// +// markdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{}) +// +// p.AddDocumentTransformer("output_key01", markdownSplitter) +func (p *Parallel) AddDocumentTransformer(outputKey string, node document.Transformer, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toDocumentTransformerNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddGraph adds a graph node to the parallel. +// It is useful when you want to use a graph or a chain as a node in the parallel. +// eg. +// +// graph, err := compose.NewChain[any,any]() +// +// p.AddGraph("output_key01", graph) +func (p *Parallel) AddGraph(outputKey string, node AnyGraph, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toAnyGraphNode(node, append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +// AddPassthrough adds a passthrough node to the parallel. +// eg. +// +// p.AddPassthrough("output_key01") +func (p *Parallel) AddPassthrough(outputKey string, opts ...GraphAddNodeOpt) *Parallel { + gNode, options := toPassthroughNode(append(opts, WithOutputKey(outputKey))...) + return p.addNode(outputKey, gNode, options) +} + +func (p *Parallel) addNode(outputKey string, node *graphNode, options *graphAddNodeOpts) *Parallel { + if p.err != nil { + return p + } + + if node == nil { + p.err = fmt.Errorf("chain parallel add node invalid, node is nil") + return p + } + + if p.outputKeys == nil { + p.outputKeys = make(map[string]bool) + } + + if _, ok := p.outputKeys[outputKey]; ok { + p.err = fmt.Errorf("parallel add node err, duplicate output key= %s", outputKey) + return p + } + + if node.nodeInfo == nil { + p.err = fmt.Errorf("chain parallel add node invalid, nodeInfo is nil") + return p + } + + node.nodeInfo.outputKey = outputKey + p.nodes = append(p.nodes, nodeOptionsPair{node, options}) + p.outputKeys[outputKey] = true + return p +} diff --git a/compose/chain_test.go b/compose/chain_test.go new file mode 100644 index 0000000..27f17a8 --- /dev/null +++ b/compose/chain_test.go @@ -0,0 +1,873 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/internal/mock/components/document" + "github.com/cloudwego/eino/internal/mock/components/embedding" + "github.com/cloudwego/eino/internal/mock/components/indexer" + "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/internal/mock/components/retriever" + "github.com/cloudwego/eino/schema" +) + +func TestChain(t *testing.T) { + + cm := &mockIntentChatModel{} + + // 构建 branch + branchCond := func(ctx context.Context, input map[string]any) (string, error) { + if rand.Intn(2) == 1 { + return "b1", nil + } + return "b2", nil + } + + b1 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("hello in branch lambda 01") + kvs["role"] = "cat" + return kvs, nil + }) + b2 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("hello in branch lambda 02") + kvs["role"] = "dog" + return kvs, nil + }) + + // 并发节点 + parallel := NewParallel() + parallel. + AddLambda("role", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + // may be change role to others by input kvs, for example (dentist/doctor...) + role := kvs["role"] + if role.(string) == "" { + role = "bird" + } + return role.(string), nil + })). + AddLambda("input", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + return "你的叫声是怎样的？", nil + })) + + // 顺序节点 + rolePlayChain := NewChain[map[string]any, *schema.Message]() + rolePlayChain. + AppendChatTemplate(prompt.FromMessages(schema.FString, schema.SystemMessage(`You are a {role}.`), schema.UserMessage(`{input}`))). + AppendChatModel(cm) + + // 构建 chain + + chain := NewChain[map[string]any, string]() + chain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + // do some logic to prepare kv as variables for next Node + // just pass through + t.Log("in view lambda: ", kvs) + return kvs, nil + })). + AppendBranch(NewChainBranch[map[string]any](branchCond).AddLambda("b1", b1).AddLambda("b2", b2)). + AppendPassthrough(). + AppendParallel(parallel). + AppendGraph(rolePlayChain). + AppendLambda(InvokableLambda(func(ctx context.Context, m *schema.Message) (string, error) { + // do some logic to check the output or something + t.Log("in view of messages: ", m.Content) + + return m.Content, nil + })) + + r, err := chain.Compile(context.Background()) + assert.Nil(t, err) + + out, err := r.Invoke(context.Background(), map[string]any{}) + assert.Nil(t, err) + t.Log(err) + + t.Log("out is : ", out) +} + +func TestChainWithException(t *testing.T) { + chain := NewChain[map[string]any, string]() + chain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + // do some logic to prepare kv as variables for next Node + // just pass through + t.Log("in view lambda: ", kvs) + return kvs, nil + })). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in view lambda 02: ", kvs) + return kvs, nil + }), WithNodeKey("xlam")) + + // items with parallels + parallel := NewParallel() + parallel. + AddLambda("hello", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + t.Log("in parallel item 01") + return "world", nil + })). + AddLambda("world", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + t.Log("in parallel item 02") + return "hello", nil + })) + + // sequence items + nchain := NewChain[map[string]any, map[string]any]() + nchain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence item 01") + return kvs, nil + })). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence item 02") + return kvs, nil + })) + + branchCond := func(ctx context.Context, input map[string]any) (string, error) { + if rand.Intn(2) == 1 { + return "b1", nil + } + return "b2", nil + } + + b1 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("hello in branch lambda 01") + kvs["role"] = "cat" + return kvs, nil + }) + b2 := InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }) + + // sequence with branch + chain.AppendBranch(NewChainBranch[map[string]any](branchCond).AddLambda("b1", b1).AddLambda("b2", b2)) + + // parallel with sequence + parallel.AddGraph("test_sequence", nchain) + + // parallel with parallel + npara := NewParallel(). + AddLambda("test_parallel1", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })). + AddLambda("test_parallel2", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + + // parallel with graph + ngraph := NewChain[map[string]any, map[string]any](). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in graph item 01") + return kvs, nil + })). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in graph item 02") + return kvs, nil + })) + nc := NewChain[map[string]any, map[string]any]() + nc.AppendGraph(ngraph) + parallel.AddGraph("test_graph", nc) + + chain.AppendPassthrough() + + // sequence with parallel + chain.AppendParallel(npara) + + // 构建 chain + chain. + AppendGraph(nchain). + AppendParallel(parallel). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + t.Log("in last view lambda: ", kvs) + return "hello last", nil + })) + + ctx := context.Background() + + r, err := chain.Compile(ctx) + assert.Nil(t, err) + + out, err := r.Invoke(ctx, map[string]any{"test": "test"}) + assert.Nil(t, err) + t.Log("out is : ", out) +} + +func TestEmptyList(t *testing.T) { + ctx := context.Background() + + // no nodes in chain + chain := NewChain[map[string]any, map[string]any]() + _, err := chain.Compile(ctx) + assert.Error(t, err) + + // no nodes in parallel + parallel := NewParallel() + chain = NewChain[map[string]any, map[string]any]() + chain.AppendParallel(parallel) + + _, err = chain.Compile(ctx) + assert.Error(t, err) + + // no nodes in sequence + emptyChain := NewChain[map[string]any, map[string]any]() + chain = NewChain[map[string]any, map[string]any]() + + chain. + AppendParallel(parallel). + AppendGraph(emptyChain). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + + _, err = chain.Compile(ctx) + assert.Error(t, err) +} + +func TestChainList(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in view lambda: ", kvs) + return kvs, nil + })) + + // parallel + parallel := NewParallel() + parallel. + AddLambda("test_parallel1", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in parallel item 01") + return kvs, nil + })) + + // seq in parallel + nchain := NewChain[map[string]any, map[string]any]() + nchain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence in parallel item 01") + return kvs, nil + })). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence in parallel item 02") + return kvs, nil + })) + + // seq in seq + nchainInChain := NewChain[map[string]any, map[string]any]() + nchainInChain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence in sequence item 01") + return kvs, nil + })). + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence in sequence item 02") + return kvs, nil + })) + + nchain.AppendGraph(nchainInChain) + + parallel.AddGraph("test_seq_in_parallel", nchain) + + chain.AppendParallel(parallel) + + r, err := chain.Compile(context.Background()) + assert.Nil(t, err) + out, err := r.Invoke(context.Background(), map[string]any{"test": "test"}) + assert.Nil(t, err) + t.Log("out is : ", out) +} + +func TestChainSingleNode(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in view lambda: ", kvs) + return kvs, nil + })) + + // single Node in chain (prepare for parallel) + singleNodeChain := NewChain[map[string]any, map[string]any]() + singleNodeChain. + AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in sequence item 01") + return kvs, nil + })) + + // add parallel + parallel := NewParallel() + parallel. + AddLambda("test_parallel1_lambda", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in parallel item 01") + return kvs, nil + })) + + parallel.AddGraph("test_parallel2_chain", singleNodeChain) + + ctx := context.Background() + + chain.AppendParallel(parallel) + r, err := chain.Compile(ctx) + assert.Nil(t, err) + + out, err := r.Invoke(ctx, map[string]any{"test": "test"}) + assert.Nil(t, err) + t.Log("out is : ", out) +} + +func TestParallelModels(t *testing.T) { + cm := &mockIntentChatModel{} + chain := NewChain[map[string]any, map[string]any]() + chatSuite := NewChain[map[string]any, string]() + chatSuite. + AppendChatTemplate(prompt.FromMessages(schema.FString, schema.SystemMessage(`You are a {role}.`), schema.UserMessage(`{input}`))). + AppendChatModel(cm). + AppendLambda(InvokableLambda(func(ctx context.Context, msg *schema.Message) (string, error) { + t.Log("in parallel item 01") + return msg.Content, nil + })) + + parallel := NewParallel() + parallel. + AddGraph("time001", chatSuite). + AddGraph("time002", chatSuite). + AddGraph("time003", chatSuite) + + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + t.Log("in view lambda: ", kvs) + return kvs, nil + })) + + chain.AppendParallel(parallel) + + ctx := context.Background() + + r, err := chain.Compile(ctx) + assert.Nil(t, err) + + out, err := r.Invoke(ctx, map[string]any{"role": "cat", "input": "你怎么叫的？"}) + assert.Nil(t, err) + + t.Log("out is : ", out) +} + +func TestChainMultiNodes(t *testing.T) { + ctx := context.Background() + + t.Run("test embedding Node", func(t *testing.T) { + chain := NewChain[[]string, [][]float64]() + + mockCtrl := gomock.NewController(t) + eb := embedding.NewMockEmbedder(mockCtrl) + chain.AppendEmbedding(eb) + + r, err := chain.Compile(ctx) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test retriever Node", func(t *testing.T) { + chain := NewChain[string, []*schema.Document]() + + chain.AppendRetriever(retriever.NewMockRetriever(gomock.NewController(t))) + + r, err := chain.Compile(ctx) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test chat model", func(t *testing.T) { + chain := NewChain[[]*schema.Message, *schema.Message]() + + cm := &mockIntentChatModel{} + chain.AppendChatModel(cm) + + r, err := chain.Compile(ctx) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test chat template", func(t *testing.T) { + chain := NewChain[map[string]any, []*schema.Message]() + + chatTemplate := prompt.FromMessages(schema.FString) + chain.AppendChatTemplate(chatTemplate) + + r, err := chain.Compile(ctx) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test lambda", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + + r, err := chain.Compile(ctx) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test indexer", func(t *testing.T) { + chain := NewChain[[]*schema.Document, []string]() + + chain.AppendIndexer(indexer.NewMockIndexer(gomock.NewController(t))) + + r, err := chain.Compile(ctx) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test parallel", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + parallel := NewParallel() + parallel.AddLambda("test_parallel", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + chain.AppendParallel(parallel) + _, err := chain.Compile(ctx) + assert.Error(t, err) + + chain = NewChain[map[string]any, map[string]any]() + parallel = NewParallel() + parallel.AddLambda("test_parallel", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + parallel.AddLambda("test_parallel", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + chain.AppendParallel(parallel) + _, err = chain.Compile(ctx) + assert.Error(t, err) + + chain = NewChain[map[string]any, map[string]any]() + parallel = NewParallel() + parallel.AddLambda("test_parallel", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + parallel.AddLambda("test_parallel1", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + chain.AppendParallel(parallel) + _, err = chain.Compile(ctx) + assert.NoError(t, err) + + chain = NewChain[map[string]any, map[string]any]() + parallel = NewParallel() + parallel.AddLambda("test_parallel", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + parallel.AddLambda("test_parallel1", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + chain.AppendParallel(parallel) + + parallel1 := NewParallel() + parallel1.AddLambda("test_parallel", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + parallel1.AddLambda("test_parallel1", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + chain.AppendParallel(parallel1) + + _, err = chain.Compile(ctx) + assert.Error(t, err) + }) + + t.Run("test tools Node", func(t *testing.T) { + ctx := context.Background() + chain := NewChain[map[string]any, map[string]any]() + toolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{}) + assert.NoError(t, err) + chain.AppendToolsNode(toolsNode) + }) + + t.Run("test chain with compile option", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + r, err := chain.Compile(ctx, WithMaxRunSteps(10)) + assert.NoError(t, err) + assert.NotNil(t, r) + }) + + t.Run("test chain return type", func(t *testing.T) { + t.Run("test chain any output type", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (any, error) { + return 1, nil + })) + _, err := chain.Compile(ctx) + assert.Nil(t, err) + }) + + t.Run("test chain error output type", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + return "123", nil + })) + _, err := chain.Compile(ctx) + assert.Error(t, err) + }) + + t.Run("test chain error input type", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, input string) (map[string]any, error) { + return nil, nil + })) + _, err := chain.Compile(ctx) + assert.Error(t, err) + }) + }) + +} + +func TestParallelMultiNodes(t *testing.T) { + ctx := context.Background() + p := NewParallel() + p.AddLambda("lambda", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + p.AddGraph("graph", NewChain[map[string]any, map[string]any]()) + p.AddIndexer("indexer", indexer.NewMockIndexer(gomock.NewController(t))) + p.AddLoader("loader", document.NewMockLoader(gomock.NewController(t))) + p.AddDocumentTransformer("document transformer", document.NewMockTransformer(gomock.NewController(t))) + p.AddRetriever("retriever", retriever.NewMockRetriever(gomock.NewController(t))) + p.AddChatModel("chatmodel", model.NewMockChatModel(gomock.NewController(t))) + p.AddChatTemplate("chatTemplate", prompt.FromMessages(schema.FString, schema.SystemMessage("hello"))) + p.AddEmbedding("embedding", embedding.NewMockEmbedder(gomock.NewController(t))) + p.AddPassthrough("passthrough") + toolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{}) + assert.NoError(t, err) + p.AddToolsNode("tools", toolsNode) + + assert.Greater(t, len(p.nodes), 6) + + ctrl := gomock.NewController(t) + p = NewParallel() + p.AddIndexer("key", indexer.NewMockIndexer(ctrl)) + p.AddLoader("key", document.NewMockLoader(ctrl)) + p.AddRetriever("r", retriever.NewMockRetriever(ctrl)) + assert.NotNil(t, p.err) + + p = NewParallel() + p.addNode("k", nil, nil) + assert.NotNil(t, p.err) + + p = &Parallel{ + outputKeys: nil, + } + p.addNode("k", &graphNode{}, nil) + assert.NotNil(t, p.err) +} + +type FakeLambdaOptions struct { + Info string +} + +type FakeLambdaOption func(opt *FakeLambdaOptions) + +func FakeWithLambdaInfo(info string) FakeLambdaOption { + return func(opt *FakeLambdaOptions) { + opt.Info = info + } +} + +func TestChainWithNodeKey(t *testing.T) { + ctx := context.Background() + + t.Run("test normal chain with node key option", func(t *testing.T) { + + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_01")) + + b := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) { + return "lambda_02", nil + }) + + b.AddLambda("lambda_02", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (map[string]any, error) { + opt := &FakeLambdaOptions{} + for _, optFn := range opts { + optFn(opt) + } + kvs["lambda_02"] = opt.Info + return kvs, nil + }), WithNodeKey("lambda_02")) + + b.AddLambda("lambda_03", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (map[string]any, error) { + opt := &FakeLambdaOptions{} + for _, optFn := range opts { + optFn(opt) + } + kvs["lambda_03"] = opt.Info + return kvs, nil + }), WithNodeKey("lambda_03")) + + chain.AppendBranch(b) + + chain.AppendPassthrough() + + p := NewParallel() + p.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + return kvs["lambda_02"].(string), nil + })) + p.AddLambda("lambda_04", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (string, error) { + opt := &FakeLambdaOptions{} + for _, optFn := range opts { + optFn(opt) + } + return opt.Info, nil + }), WithNodeKey("lambda_04")) + + p.AddLambda("lambda_05", InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (string, error) { + opt := &FakeLambdaOptions{} + for _, optFn := range opts { + optFn(opt) + } + return opt.Info, nil + }), WithNodeKey("lambda_05")) + chain.AppendParallel(p) + + chain.AppendLambda(InvokableLambdaWithOption(func(ctx context.Context, kvs map[string]any, opts ...FakeLambdaOption) (map[string]any, error) { + opt := &FakeLambdaOptions{} + for _, optFn := range opts { + optFn(opt) + } + kvs["lambda_06"] = opt.Info + return kvs, nil + }), WithNodeKey("lambda_06")) + + r, err := chain.Compile(ctx) + assert.Nil(t, err) + + res, err := r.Invoke(ctx, map[string]any{}, + WithLambdaOption(FakeWithLambdaInfo("normal")), + WithLambdaOption(FakeWithLambdaInfo("info_lambda_02")).DesignateNode("lambda_02"), // branch + WithLambdaOption(FakeWithLambdaInfo("info_lambda_03")).DesignateNode("lambda_03"), // branch (wont run) + WithLambdaOption(FakeWithLambdaInfo("info_lambda_05")).DesignateNode("lambda_05"), // parallel + ) + assert.Nil(t, err) + + assert.Equal(t, "info_lambda_02", res["lambda_02"]) // transmit option with DesigateNode + assert.Equal(t, "info_lambda_05", res["lambda_05"]) // transmit option with DesigateNode + assert.Equal(t, "normal", res["lambda_06"]) // without DesigateNode, using default option + }) + + t.Run("test chain with node key option and error with correct error info", func(t *testing.T) { + + t.Run("compile error of chain", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + return "123", nil + }), WithNodeKey("lambda_01")) + + c, err := chain.Compile(ctx) + assert.Nil(t, c) + fmt.Printf("%+v\n", err) + + assert.Contains(t, err.Error(), "edge[lambda_01]") + }) + + t.Run("compile error of branch", func(t *testing.T) { + t.Run("without node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_01")) + b := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) { + return "lambda_02", nil + }) + b.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + b.AddLambda("lambda_03", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + return "", nil + })) + chain.AppendBranch(b) + c, err := chain.Compile(ctx) + assert.Nil(t, c) + fmt.Printf("%+v\n", err) + assert.Contains(t, err.Error(), "edge[node_1_branch_lambda_03]") // with no node key option, will use default node key + }) + + t.Run("with node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_01")) + + b := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) { + return "lambda_02", nil + }) + b.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_02")) + + b.AddLambda("lambda_03", InvokableLambda(func(ctx context.Context, kvs map[string]any) (string, error) { + return "123", nil + }), WithNodeKey("key_of_lambda_03")) + + chain.AppendBranch(b) + c, err := chain.Compile(ctx) + assert.Nil(t, c) + fmt.Printf("%+v\n", err) + assert.Contains(t, err.Error(), "edge[key_of_lambda_03]") + }) + }) + + t.Run("compile error of parallel", func(t *testing.T) { + t.Run("without node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_01")) + p := NewParallel() + p.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + })) + p.AddLambda("lambda_03", InvokableLambda(func(ctx context.Context, v string) (string, error) { + return "", nil + })) + + chain.AppendParallel(p) + + c, err := chain.Compile(ctx) + assert.Nil(t, c) + fmt.Printf("%+v\n", err) + assert.Contains(t, err.Error(), "to=node_1_parallel_1") // with no node key option, will use default node key + }) + + t.Run("with node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + chain.AppendLambda(InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_01")) + p := NewParallel() + p.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return kvs, nil + }), WithNodeKey("lambda_02")) + p.AddLambda("lambda_03", InvokableLambda(func(ctx context.Context, v string) (string, error) { + return "", nil + }), WithNodeKey("key_of_lambda_03")) + chain.AppendParallel(p) + c, err := chain.Compile(ctx) + assert.Nil(t, c) + fmt.Printf("%+v\n", err) + assert.Contains(t, err.Error(), "to=key_of_lambda_03") + }) + }) + + t.Run("invoke error", func(t *testing.T) { + t.Run("branch with out node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + + b := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) { + return "lambda_01", nil + }) + + b.AddLambda("lambda_01", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, fmt.Errorf("fake error") + })) + b.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, nil + })) + + chain.AppendBranch(b) + c, err := chain.Compile(ctx) + assert.Nil(t, err) + + _, err = c.Invoke(ctx, map[string]any{}) + fmt.Printf("%+v\n", err) + assert.Error(t, err) + assert.Contains(t, err.Error(), "node_0_branch_lambda_01") // with no node key option, will use default node key + }) + + t.Run("branch with node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + b := NewChainBranch(func(ctx context.Context, input map[string]any) (string, error) { + return "lambda_01", nil + }) + b.AddLambda("lambda_01", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, fmt.Errorf("fake error") + }), WithNodeKey("key_of_lambda_01")) + b.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, nil + })) + + chain.AppendBranch(b) + c, err := chain.Compile(ctx) + assert.Nil(t, err) + _, err = c.Invoke(ctx, map[string]any{}) + fmt.Printf("%+v\n", err) + assert.Error(t, err) + assert.Contains(t, err.Error(), "key_of_lambda_01") + }) + + t.Run("parallel with out node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + p := NewParallel() + p.AddLambda("lambda_01", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, fmt.Errorf("fake error") + })) + p.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, nil + })) + chain.AppendParallel(p) + c, err := chain.Compile(ctx) + assert.Nil(t, err) + _, err = c.Invoke(ctx, map[string]any{}) + fmt.Printf("%+v\n", err) + assert.Error(t, err) + assert.Contains(t, err.Error(), "node_0_parallel_0") // with no node key option, will use default node key + }) + + t.Run("parallel with node key", func(t *testing.T) { + chain := NewChain[map[string]any, map[string]any]() + p := NewParallel() + p.AddLambda("lambda_01", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, fmt.Errorf("fake error") + }), WithNodeKey("key_of_lambda_01")) + p.AddLambda("lambda_02", InvokableLambda(func(ctx context.Context, kvs map[string]any) (map[string]any, error) { + return nil, nil + })) + chain.AppendParallel(p) + c, err := chain.Compile(ctx) + assert.Nil(t, err) + _, err = c.Invoke(ctx, map[string]any{}) + fmt.Printf("%+v\n", err) + assert.Error(t, err) + assert.Contains(t, err.Error(), "key_of_lambda_01") + }) + }) + + }) + +} diff --git a/compose/checkpoint.go b/compose/checkpoint.go new file mode 100644 index 0000000..c174994 --- /dev/null +++ b/compose/checkpoint.go @@ -0,0 +1,404 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/internal/serialization" + "github.com/cloudwego/eino/schema" +) + +func init() { + schema.RegisterName[*checkpoint]("_eino_checkpoint") + schema.RegisterName[*dagChannel]("_eino_dag_channel") + schema.RegisterName[*pregelChannel]("_eino_pregel_channel") + schema.RegisterName[dependencyState]("_eino_dependency_state") + _ = serialization.GenericRegister[channel]("_eino_channel") +} + +// RegisterSerializableType registers a custom type for eino serialization. +// This allows eino to properly serialize and deserialize custom types. +// Both custom interfaces and structs need to be registered using this function. +// Types only need to be registered once - pointers and other references will be handled automatically. +// All built-in eino types are already registered. +// Parameters: +// - name: A unique identifier for the type being registered (should not start with "_eino") +// - T: The generic type parameter representing the type to register +// Returns: +// - error: An error if registration fails (e.g., if the type is already registered) +// Deprecated: RegisterSerializableType is deprecated. Use schema.RegisterName[T](name) instead. +func RegisterSerializableType[T any](name string) (err error) { + return serialization.GenericRegister[T](name) +} + +type CheckPointStore = core.CheckPointStore + +type Serializer interface { + Marshal(v any) ([]byte, error) + Unmarshal(data []byte, v any) error +} + +// WithCheckPointStore sets the checkpoint store implementation for a graph. +func WithCheckPointStore(store CheckPointStore) GraphCompileOption { + return func(o *graphCompileOptions) { + o.checkPointStore = store + } +} + +// WithSerializer sets the serializer used to persist checkpoint state. +func WithSerializer(serializer Serializer) GraphCompileOption { + return func(o *graphCompileOptions) { + o.serializer = serializer + } +} + +// WithCheckPointID sets the checkpoint ID to load from and write to by default. +func WithCheckPointID(checkPointID string) Option { + return Option{ + checkPointID: &checkPointID, + } +} + +// WithWriteToCheckPointID specifies a different checkpoint ID to write to. +// If not provided, the checkpoint ID from WithCheckPointID will be used for writing. +// This is useful for scenarios where you want to load from an existed checkpoint +// but save the progress to a new, separate checkpoint. +func WithWriteToCheckPointID(checkPointID string) Option { + return Option{ + writeToCheckPointID: &checkPointID, + } +} + +// WithForceNewRun forces the graph to run from the beginning, ignoring any checkpoints. +func WithForceNewRun() Option { + return Option{ + forceNewRun: true, + } +} + +// StateModifier modifies state during checkpoint operations for a given node path. +type StateModifier func(ctx context.Context, path NodePath, state any) error + +// WithStateModifier installs a state modifier invoked during checkpoint read/write. +func WithStateModifier(sm StateModifier) Option { + return Option{ + stateModifier: sm, + } +} + +type checkpoint struct { + Channels map[string]channel + Inputs map[string] /*node key*/ any /*input*/ + State any + SkipPreHandler map[string]bool + RerunNodes []string + + SubGraphs map[string]*checkpoint + + InterruptID2Addr map[string]Address + InterruptID2State map[string]core.InterruptState +} + +type stateModifierKey struct{} +type checkPointKey struct{} // *checkpoint + +func getStateModifier(ctx context.Context) StateModifier { + if sm, ok := ctx.Value(stateModifierKey{}).(StateModifier); ok { + return sm + } + return nil +} + +func setStateModifier(ctx context.Context, modifier StateModifier) context.Context { + return context.WithValue(ctx, stateModifierKey{}, modifier) +} + +func getCheckPointFromStore(ctx context.Context, id string, cpr *checkPointer) (cp *checkpoint, err error) { + cp, existed, err := cpr.get(ctx, id) + if err != nil { + return nil, err + } + if !existed { + return nil, nil + } + + return cp, nil +} + +func setCheckPointToCtx(ctx context.Context, cp *checkpoint) context.Context { + ctx = core.PopulateInterruptState(ctx, cp.InterruptID2Addr, cp.InterruptID2State) + return context.WithValue(ctx, checkPointKey{}, cp) +} + +func getCheckPointFromCtx(ctx context.Context) *checkpoint { + if cp, ok := ctx.Value(checkPointKey{}).(*checkpoint); ok { + return cp + } + return nil +} + +func forwardCheckPoint(ctx context.Context, nodeKey string) context.Context { + cp := getCheckPointFromCtx(ctx) + if cp == nil { + return ctx + } + + if subCP, ok := cp.SubGraphs[nodeKey]; ok { + delete(cp.SubGraphs, nodeKey) // only forward once + return context.WithValue(ctx, checkPointKey{}, subCP) + } + return context.WithValue(ctx, checkPointKey{}, (*checkpoint)(nil)) +} + +func newCheckPointer( + inputPairs, outputPairs map[string]streamConvertPair, + store CheckPointStore, + serializer Serializer, +) *checkPointer { + if serializer == nil { + serializer = &serialization.InternalSerializer{} + } + return &checkPointer{ + sc: newStreamConverter(inputPairs, outputPairs), + store: store, + serializer: serializer, + } +} + +type checkPointer struct { + sc *streamConverter + store CheckPointStore + serializer Serializer +} + +func (c *checkPointer) get(ctx context.Context, id string) (*checkpoint, bool, error) { + data, existed, err := c.store.Get(ctx, id) + if err != nil || existed == false { + return nil, existed, err + } + + cp := &checkpoint{} + err = c.serializer.Unmarshal(data, cp) + if err != nil { + return nil, false, err + } + + return cp, true, nil +} + +func (c *checkPointer) set(ctx context.Context, id string, cp *checkpoint) error { + normalizeCheckpointTypedNilInputs(cp) + + data, err := c.serializer.Marshal(cp) + if err != nil { + return err + } + + return c.store.Set(ctx, id, data) +} + +func normalizeCheckpointTypedNilInputs(cp *checkpoint) { + if cp == nil { + return + } + for key, input := range cp.Inputs { + if isTypedNil(input) { + cp.Inputs[key] = nil + } + } + for _, sub := range cp.SubGraphs { + normalizeCheckpointTypedNilInputs(sub) + } +} + +func isTypedNil(v any) bool { + if v == nil { + return false + } + rv := reflect.ValueOf(v) + switch rv.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice: + return rv.IsNil() + default: + return false + } +} + +// MigrateCheckpointState is an advanced compatibility utility for checkpoint upgrades. +// +// It decodes checkpoint bytes using the given serializer, applies migrate to checkpoint.State and +// all nested SubGraphs' states, then re-encodes the checkpoint. +// +// Typical use cases: +// - Resume-time migration when you changed your graph state type/schema and need to load old +// checkpoints without discarding them. +// - Framework-level backward compatibility (e.g. ADK upgrading checkpoints across versions). +// +// Migrate callback contract: +// - Returns (newState, changed, error). +// - If changed is false, the state is left as-is. +// - If error is non-nil, migration stops and the error is returned to the caller. +// +// The original bytes are returned only if no state was changed anywhere in the checkpoint tree. +func MigrateCheckpointState(data []byte, serializer Serializer, migrate func(state any) (any, bool, error)) ([]byte, error) { + cp := &checkpoint{} + if err := serializer.Unmarshal(data, cp); err != nil { + return nil, err + } + changed, err := migrateCheckpoint(cp, migrate) + if err != nil { + return nil, err + } + if !changed { + return data, nil + } + return serializer.Marshal(cp) +} + +// migrateCheckpoint recursively applies migrate to cp.State and all SubGraphs. +func migrateCheckpoint(cp *checkpoint, migrate func(state any) (any, bool, error)) (bool, error) { + anyChanged := false + if cp.State != nil { + newState, changed, err := migrate(cp.State) + if err != nil { + return false, err + } + if changed { + cp.State = newState + anyChanged = true + } + } + for _, sub := range cp.SubGraphs { + changed, err := migrateCheckpoint(sub, migrate) + if err != nil { + return false, err + } + if changed { + anyChanged = true + } + } + return anyChanged, nil +} + +// convertCheckPoint if value in checkpoint is streamReader, convert it to non-stream +func (c *checkPointer) convertCheckPoint(cp *checkpoint, isStream bool) (err error) { + for _, ch := range cp.Channels { + err = ch.convertValues(func(m map[string]any) error { + return c.sc.convertOutputs(isStream, m) + }) + if err != nil { + return err + } + } + + err = c.sc.convertInputs(isStream, cp.Inputs) + if err != nil { + return err + } + + return nil +} + +// convertCheckPoint convert values in checkpoint to streamReader if needed +func (c *checkPointer) restoreCheckPoint(cp *checkpoint, isStream bool) (err error) { + for _, ch := range cp.Channels { + err = ch.convertValues(func(m map[string]any) error { + return c.sc.restoreOutputs(isStream, m) + }) + if err != nil { + return err + } + } + + err = c.sc.restoreInputs(isStream, cp.Inputs) + if err != nil { + return err + } + + return nil +} + +func newStreamConverter(inputPairs, outputPairs map[string]streamConvertPair) *streamConverter { + return &streamConverter{ + inputPairs: inputPairs, + outputPairs: outputPairs, + } +} + +type streamConverter struct { + inputPairs, outputPairs map[string]streamConvertPair +} + +func (s *streamConverter) convertInputs(isStream bool, values map[string]any) error { + return convert(values, s.inputPairs, isStream) +} + +func (s *streamConverter) restoreInputs(isStream bool, values map[string]any) error { + return restore(values, s.inputPairs, isStream) +} + +func (s *streamConverter) convertOutputs(isStream bool, values map[string]any) error { + return convert(values, s.outputPairs, isStream) +} + +func (s *streamConverter) restoreOutputs(isStream bool, values map[string]any) error { + return restore(values, s.outputPairs, isStream) +} + +func convert(values map[string]any, convPairs map[string]streamConvertPair, isStream bool) error { + if !isStream { + return nil + } + for key, v := range values { + convPair, ok := convPairs[key] + if !ok { + return fmt.Errorf("checkpoint conv stream fail, node[%s] have not been registered", key) + } + sr, ok := v.(streamReader) + if !ok { + return fmt.Errorf("checkpoint conv stream fail, value of [%s] isn't stream", key) + } + nValue, err := convPair.concatStream(sr) + if err != nil { + return err + } + values[key] = nValue + } + return nil +} + +func restore(values map[string]any, convPairs map[string]streamConvertPair, isStream bool) error { + if !isStream { + return nil + } + for key, v := range values { + convPair, ok := convPairs[key] + if !ok { + return fmt.Errorf("checkpoint restore stream fail, node[%s] have not been registered", key) + } + sr, err := convPair.restoreStream(v) + if err != nil { + return err + } + values[key] = sr + } + return nil +} diff --git a/compose/checkpoint_migrate_test.go b/compose/checkpoint_migrate_test.go new file mode 100644 index 0000000..6b64b60 --- /dev/null +++ b/compose/checkpoint_migrate_test.go @@ -0,0 +1,116 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" +) + +type stubSerializer struct { + unmarshal func(data []byte, v any) error + marshal func(v any) ([]byte, error) +} + +func (s stubSerializer) Marshal(v any) ([]byte, error) { + return s.marshal(v) +} + +func (s stubSerializer) Unmarshal(data []byte, v any) error { + return s.unmarshal(data, v) +} + +func TestMigrateCheckpointState_UnmarshalError(t *testing.T) { + in := []byte("in") + codec := stubSerializer{ + unmarshal: func(_ []byte, _ any) error { return errors.New("bad") }, + marshal: func(_ any) ([]byte, error) { return []byte("unused"), nil }, + } + _, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) { + return state, false, nil + }) + assert.Error(t, err) +} + +func TestMigrateCheckpointState_NoChangeReturnsOriginalBytes(t *testing.T) { + in := []byte("in") + cp := &checkpoint{State: "s"} + codec := stubSerializer{ + unmarshal: func(_ []byte, v any) error { + *(v.(*checkpoint)) = *cp + return nil + }, + marshal: func(_ any) ([]byte, error) { + return []byte("marshaled"), nil + }, + } + out, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) { + return state, false, nil + }) + assert.NoError(t, err) + assert.Equal(t, in, out) +} + +func TestMigrateCheckpointState_ChangeTriggersMarshal(t *testing.T) { + in := []byte("in") + cp := &checkpoint{State: "s"} + var sawState any + codec := stubSerializer{ + unmarshal: func(_ []byte, v any) error { + *(v.(*checkpoint)) = *cp + return nil + }, + marshal: func(v any) ([]byte, error) { + sawState = v.(*checkpoint).State + return []byte("marshaled"), nil + }, + } + out, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) { + return "s2", true, nil + }) + assert.NoError(t, err) + assert.Equal(t, []byte("marshaled"), out) + assert.Equal(t, "s2", sawState) +} + +func TestMigrateCheckpointState_MigrateErrorStops(t *testing.T) { + in := []byte("in") + cp := &checkpoint{ + State: "root", + SubGraphs: map[string]*checkpoint{ + "sub": {State: "sub"}, + }, + } + codec := stubSerializer{ + unmarshal: func(_ []byte, v any) error { + *(v.(*checkpoint)) = *cp + return nil + }, + marshal: func(_ any) ([]byte, error) { + return []byte("marshaled"), nil + }, + } + _, err := MigrateCheckpointState(in, codec, func(state any) (any, bool, error) { + if state == "sub" { + return nil, false, errors.New("boom") + } + return state, false, nil + }) + assert.Error(t, err) +} diff --git a/compose/checkpoint_test.go b/compose/checkpoint_test.go new file mode 100644 index 0000000..ccba397 --- /dev/null +++ b/compose/checkpoint_test.go @@ -0,0 +1,2059 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "bytes" + "context" + "encoding/gob" + "errors" + "io" + "sync" + "testing" + "time" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/internal/callbacks" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/internal/serialization" + "github.com/cloudwego/eino/schema" +) + +type inMemoryStore struct { + m map[string][]byte +} + +func (i *inMemoryStore) Get(_ context.Context, checkPointID string) ([]byte, bool, error) { + v, ok := i.m[checkPointID] + return v, ok, nil +} + +func (i *inMemoryStore) Set(_ context.Context, checkPointID string, checkPoint []byte) error { + i.m[checkPointID] = checkPoint + return nil +} + +func newInMemoryStore() *inMemoryStore { + return &inMemoryStore{ + m: make(map[string][]byte), + } +} + +type checkpointTypedNilInput struct { + Value string +} + +type checkpointGobSerializer struct{} + +func (g *checkpointGobSerializer) Marshal(v any) ([]byte, error) { + buf := new(bytes.Buffer) + err := gob.NewEncoder(buf).Encode(v) + if err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func (g *checkpointGobSerializer) Unmarshal(data []byte, v any) error { + return gob.NewDecoder(bytes.NewBuffer(data)).Decode(v) +} + +func TestCheckpointGobSerializerNormalizesTypedNilInput(t *testing.T) { + store := newInMemoryStore() + cpr := newCheckPointer(nil, nil, store, &checkpointGobSerializer{}) + cp := &checkpoint{ + Channels: map[string]channel{}, + Inputs: map[string]any{"node": (*checkpointTypedNilInput)(nil)}, + SkipPreHandler: map[string]bool{}, + } + + err := cpr.set(context.Background(), "cp", cp) + assert.NoError(t, err) + if err != nil { + return + } + + got, existed, err := cpr.get(context.Background(), "cp") + assert.NoError(t, err) + assert.True(t, existed) + assert.Contains(t, got.Inputs, "node") + assert.Nil(t, got.Inputs["node"]) +} + +func TestRunnableInvokeAcceptsNilForPointerInput(t *testing.T) { + lambda := InvokableLambda(func(ctx context.Context, input *checkpointTypedNilInput) (string, error) { + assert.Nil(t, input) + return "ok", nil + }) + + out, err := lambda.executor.i(context.Background(), nil) + assert.NoError(t, err) + assert.Equal(t, "ok", out) +} + +type testStruct struct { + A string +} + +func init() { + schema.Register[testStruct]() + schema.RegisterName[*checkpointTypedNilInput]("compose_checkpoint_typed_nil_input") +} + +func TestSimpleCheckPoint(t *testing.T) { + store := newInMemoryStore() + + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) { + return &testStruct{A: ""} + })) + + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + assert.NoError(t, err) + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "2", nil + }), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) { + return in + state.A, nil + })) + assert.NoError(t, err) + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", "2") + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor), WithCheckPointStore(store), WithInterruptAfterNodes([]string{"1"}), WithInterruptBeforeNodes([]string{"2"}), WithGraphName("root")) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "start", WithCheckPointID("1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, &testStruct{A: ""}, info.State) + assert.Equal(t, []string{"2"}, info.BeforeNodes) + assert.Equal(t, []string{"1"}, info.AfterNodes) + assert.Empty(t, info.RerunNodesExtra) + assert.Empty(t, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + })) + + rCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + result, err := r.Invoke(rCtx, "start", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "start1state2", result) + + /* _, err = r.Stream(ctx, "start", WithCheckPointID("2")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, &testStruct{A: ""}, info.State) + assert.Equal(t, []string{"2"}, info.BeforeNodes) + assert.Equal(t, []string{"1"}, info.AfterNodes) + assert.Empty(t, info.RerunNodesExtra) + assert.Empty(t, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + })) + + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + streamResult, err := r.Stream(rCtx, "start", WithCheckPointID("2")) + assert.NoError(t, err) + result = "" + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + result += chunk + } + + assert.Equal(t, "start1state2", result)*/ +} + +func TestCustomStructInAn2y(t *testing.T) { + store := newInMemoryStore() + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) { + return &testStruct{A: ""} + })) + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output *testStruct, err error) { + return &testStruct{A: input + "1"}, nil + }), WithOutputKey("1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) { + return input["1"].(*testStruct).A + "2", nil + }), WithStatePreHandler(func(ctx context.Context, in map[string]any, state *testStruct) (map[string]any, error) { + in["1"].(*testStruct).A += state.A + return in, nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", "2") + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(store), WithInterruptAfterNodes([]string{"1"}), + WithGraphName("root")) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "start", WithCheckPointID("1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, &testStruct{A: ""}, info.State) + assert.Equal(t, []string{"1"}, info.AfterNodes) + assert.Empty(t, info.RerunNodesExtra) + assert.Empty(t, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + })) + rCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + result, err := r.Invoke(rCtx, "start", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "start1state2", result) + + _, err = r.Stream(ctx, "start", WithCheckPointID("2")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, &testStruct{A: ""}, info.State) + assert.Equal(t, []string{"1"}, info.AfterNodes) + assert.Empty(t, info.RerunNodesExtra) + assert.Empty(t, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + })) + + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + streamResult, err := r.Stream(rCtx, "start", WithCheckPointID("2")) + assert.NoError(t, err) + result = "" + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + result += chunk + } + + assert.Equal(t, "start1state2", result) +} + +func TestSubGraph(t *testing.T) { + subG := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) { + return &testStruct{A: ""} + })) + err := subG.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + assert.NoError(t, err) + err = subG.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "2", nil + }), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) { + return in + state.A, nil + })) + assert.NoError(t, err) + + err = subG.AddEdge(START, "1") + assert.NoError(t, err) + err = subG.AddEdge("1", "2") + assert.NoError(t, err) + err = subG.AddEdge("2", END) + assert.NoError(t, err) + + g := NewGraph[string, string]() + err = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + assert.NoError(t, err) + err = g.AddGraphNode("2", subG, WithGraphCompileOptions(WithInterruptAfterNodes([]string{"1"}))) + assert.NoError(t, err) + err = g.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "3", nil + })) + assert.NoError(t, err) + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", "2") + assert.NoError(t, err) + err = g.AddEdge("2", "3") + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithGraphName("root")) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "start", WithCheckPointID("1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + + rCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + result, err := r.Invoke(rCtx, "start", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "start11state23", result) + + _, err = r.Stream(ctx, "start", WithCheckPointID("2")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]any), + SubGraphs: map[string]*InterruptInfo{}, + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + streamResult, err := r.Stream(rCtx, "start", WithCheckPointID("2")) + assert.NoError(t, err) + result = "" + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + result += chunk + } + + assert.Equal(t, "start11state23", result) +} + +type testGraphCallback struct { + onStartTimes int + onEndTimes int + onStreamStartTimes int + onStreamEndTimes int + onErrorTimes int +} + +func (t *testGraphCallback) OnStart(ctx context.Context, info *callbacks.RunInfo, _ callbacks.CallbackInput) context.Context { + if info.Component == ComponentOfGraph { + t.onStartTimes++ + } + return ctx +} + +func (t *testGraphCallback) OnEnd(ctx context.Context, info *callbacks.RunInfo, _ callbacks.CallbackOutput) context.Context { + if info.Component == ComponentOfGraph { + t.onEndTimes++ + } + return ctx +} + +func (t *testGraphCallback) OnError(ctx context.Context, info *callbacks.RunInfo, _ error) context.Context { + if info.Component == ComponentOfGraph { + t.onErrorTimes++ + } + return ctx +} + +func (t *testGraphCallback) OnStartWithStreamInput(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + input.Close() + if info.Component == ComponentOfGraph { + t.onStreamStartTimes++ + } + return ctx +} + +func (t *testGraphCallback) OnEndWithStreamOutput(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + output.Close() + if info.Component == ComponentOfGraph { + t.onStreamEndTimes++ + } + return ctx +} + +func TestNestedSubGraph(t *testing.T) { + sSubG := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) { + return &testStruct{A: ""} + })) + err := sSubG.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + assert.NoError(t, err) + err = sSubG.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "2", nil + }), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) { + return in + state.A, nil + })) + assert.NoError(t, err) + + err = sSubG.AddEdge(START, "1") + assert.NoError(t, err) + err = sSubG.AddEdge("1", "2") + assert.NoError(t, err) + err = sSubG.AddEdge("2", END) + assert.NoError(t, err) + + subG := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) { + return &testStruct{A: ""} + })) + err = subG.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + assert.NoError(t, err) + err = subG.AddGraphNode("2", sSubG, WithGraphCompileOptions(WithInterruptAfterNodes([]string{"1"})), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) { + return in + state.A, nil + }), WithOutputKey("2")) + assert.NoError(t, err) + err = subG.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "3", nil + }), WithOutputKey("3")) + assert.NoError(t, err) + err = subG.AddLambdaNode("4", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) { + return input["2"].(string) + "4\n" + input["3"].(string) + "4\n" + input["state"].(string) + "4\n", nil + }), WithStatePreHandler(func(ctx context.Context, in map[string]any, state *testStruct) (map[string]any, error) { + in["state"] = state.A + return in, nil + })) + assert.NoError(t, err) + err = subG.AddEdge(START, "1") + assert.NoError(t, err) + err = subG.AddEdge("1", "2") + assert.NoError(t, err) + err = subG.AddEdge("1", "3") + assert.NoError(t, err) + err = subG.AddEdge("3", "4") + assert.NoError(t, err) + err = subG.AddEdge("2", "4") + assert.NoError(t, err) + err = subG.AddEdge("4", END) + assert.NoError(t, err) + + g := NewGraph[string, string]() + err = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + assert.NoError(t, err) + err = g.AddGraphNode("2", subG, WithGraphCompileOptions(WithInterruptAfterNodes([]string{"1", "3"}), WithInterruptBeforeNodes([]string{"4"}))) + assert.NoError(t, err) + err = g.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "3", nil + })) + assert.NoError(t, err) + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", "2") + assert.NoError(t, err) + err = g.AddEdge("2", "3") + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithGraphName("root")) + assert.NoError(t, err) + + tGCB := &testGraphCallback{} + _, err = r.Invoke(ctx, "start", WithCheckPointID("1"), WithCallbacks(tGCB)) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + + rCtx := ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Invoke(rCtx, "start", WithCheckPointID("1"), WithCallbacks(tGCB)) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + AfterNodes: []string{"3"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + Parent: &InterruptCtx{ + ID: "runnable:root", + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Invoke(rCtx, "start", WithCheckPointID("1"), WithCallbacks(tGCB)) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + BeforeNodes: []string{"4"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state2"}) + result, err := r.Invoke(rCtx, "start", WithCheckPointID("1"), WithCallbacks(tGCB)) + assert.NoError(t, err) + assert.Equal(t, `start11state1state24 +start1134 +state24 +3`, result) + + _, err = r.Stream(ctx, "start", WithCheckPointID("2"), WithCallbacks(tGCB)) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Stream(rCtx, "start", WithCheckPointID("2"), WithCallbacks(tGCB)) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + AfterNodes: []string{"3"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Stream(rCtx, "start", WithCheckPointID("2"), WithCallbacks(tGCB)) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + BeforeNodes: []string{"4"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state2"}) + streamResult, err := r.Stream(rCtx, "start", WithCheckPointID("2"), WithCallbacks(tGCB)) + assert.NoError(t, err) + result = "" + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + result += chunk + } + assert.Equal(t, `start11state1state24 +start1134 +state24 +3`, result) + + assert.Equal(t, 10, tGCB.onStartTimes) // 3+sSubG*1*3+subG*2*2+g*0 + assert.Equal(t, 3, tGCB.onEndTimes) // success*3 + assert.Equal(t, 10, tGCB.onStreamStartTimes) // 3+sSubG*1*3+subG*2*2+g*0 + assert.Equal(t, 3, tGCB.onStreamEndTimes) // success*3 + assert.Equal(t, 14, tGCB.onErrorTimes) // 2*(sSubG*1*3+subG*2*2+g*0) + + // dag + r, err = g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithNodeTriggerMode(AllPredecessor), + WithGraphName("root")) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "start", WithCheckPointID("1")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Invoke(rCtx, "start", WithCheckPointID("1")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + AfterNodes: []string{"3"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + ID: "runnable:root;node:2;node:2", + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Invoke(rCtx, "start", WithCheckPointID("1")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + BeforeNodes: []string{"4"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state2"}) + result, err = r.Invoke(rCtx, "start", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, `start11state1state24 +start1134 +state24 +3`, result) + + _, err = r.Stream(ctx, "start", WithCheckPointID("2")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Stream(rCtx, "start", WithCheckPointID("2")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + AfterNodes: []string{"3"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: ""}, + AfterNodes: []string{"1"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state"}) + _, err = r.Stream(rCtx, "start", WithCheckPointID("2")) + assert.NotNil(t, err) + info, ok = ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, map[string]*InterruptInfo{ + "2": { + State: &testStruct{A: "state"}, + BeforeNodes: []string{"4"}, + RerunNodesExtra: make(map[string]interface{}), + SubGraphs: make(map[string]*InterruptInfo), + }, + }, info.SubGraphs) + assert.True(t, info.InterruptContexts[0].EqualsWithoutID(&InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + { + Type: AddressSegmentNode, + ID: "2", + }, + }, + Info: &testStruct{ + A: "state", + }, + IsRootCause: true, + Parent: &InterruptCtx{ + Address: Address{ + { + Type: AddressSegmentRunnable, + ID: "root", + }, + }, + }, + })) + rCtx = ResumeWithData(ctx, info.InterruptContexts[0].ID, &testStruct{A: "state2"}) + streamResult, err = r.Stream(rCtx, "start", WithCheckPointID("2")) + assert.NoError(t, err) + result = "" + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + result += chunk + } + assert.Equal(t, `start11state1state24 +start1134 +state24 +3`, result) +} + +func TestDAGInterrupt(t *testing.T) { + g := NewGraph[string, map[string]any]() + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + time.Sleep(time.Millisecond * 100) + return input, nil + }), WithOutputKey("1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + time.Sleep(time.Millisecond * 200) + return input, nil + }), WithOutputKey("2")) + assert.NoError(t, err) + err = g.AddPassthroughNode("3") + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge(START, "2") + assert.NoError(t, err) + err = g.AddEdge("1", "3") + assert.NoError(t, err) + err = g.AddEdge("2", "3") + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore()), WithInterruptAfterNodes([]string{"1", "2"})) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "input", WithCheckPointID("1")) + info, existed := ExtractInterruptInfo(err) + assert.True(t, existed) + assert.Equal(t, []string{"1", "2"}, info.AfterNodes) + + result, err := r.Invoke(ctx, "", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, map[string]any{"1": "input", "2": "input"}, result) +} + +func TestRerunNodeInterrupt(t *testing.T) { + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testStruct) { + return &testStruct{} + })) + + times := 0 + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + defer func() { times++ }() + if times%2 == 0 { + return "", NewInterruptAndRerunErr("test extra") + } + return input, nil + }), WithStatePreHandler(func(ctx context.Context, in string, state *testStruct) (string, error) { + return state.A, nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "input", WithCheckPointID("1")) + info, existed := ExtractInterruptInfo(err) + assert.True(t, existed) + assert.Equal(t, []string{"1"}, info.RerunNodes) + + result, err := r.Invoke(ctx, "", WithCheckPointID("1"), WithStateModifier(func(ctx context.Context, path NodePath, state any) error { + state.(*testStruct).A = "state" + return nil + })) + assert.NoError(t, err) + assert.Equal(t, "state", result) + + _, err = r.Stream(ctx, "input", WithCheckPointID("2")) + info, existed = ExtractInterruptInfo(err) + assert.True(t, existed) + assert.Equal(t, []string{"1"}, info.RerunNodes) + assert.Equal(t, "test extra", info.RerunNodesExtra["1"].(string)) + + streamResult, err := r.Stream(ctx, "", WithCheckPointID("2"), WithStateModifier(func(ctx context.Context, path NodePath, state any) error { + state.(*testStruct).A = "state" + return nil + })) + assert.NoError(t, err) + chunk, err := streamResult.Recv() + assert.NoError(t, err) + assert.Equal(t, "state", chunk) + _, err = streamResult.Recv() + assert.Equal(t, io.EOF, err) +} + +type myInterface interface { + A() +} + +func TestInterfaceResume(t *testing.T) { + g := NewGraph[myInterface, string]() + times := 0 + assert.NoError(t, g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input myInterface) (output string, err error) { + if times == 0 { + times++ + return "", NewInterruptAndRerunErr("test extra") + } + return "success", nil + }))) + assert.NoError(t, g.AddEdge(START, "1")) + assert.NoError(t, g.AddEdge("1", END)) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, nil, WithCheckPointID("1")) + info, existed := ExtractInterruptInfo(err) + assert.True(t, existed) + assert.Equal(t, []string{"1"}, info.RerunNodes) + result, err := r.Invoke(ctx, nil, WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "success", result) +} + +func TestEarlyFailCallback(t *testing.T) { + g := NewGraph[string, string]() + assert.NoError(t, g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }))) + assert.NoError(t, g.AddEdge(START, "1")) + assert.NoError(t, g.AddEdge("1", END)) + + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + tGCB := &testGraphCallback{} + _, _ = r.Invoke(ctx, "", WithCallbacks(tGCB), WithRuntimeMaxSteps(1)) + assert.Equal(t, 1, tGCB.onStartTimes) + assert.Equal(t, 1, tGCB.onErrorTimes) + assert.Equal(t, 0, tGCB.onEndTimes) +} + +func TestGraphStartInterrupt(t *testing.T) { + subG := NewGraph[string, string]() + _ = subG.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "sub1", nil + })) + _ = subG.AddEdge(START, "1") + _ = subG.AddEdge("1", END) + + g := NewGraph[string, string]() + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + _ = g.AddGraphNode("2", subG, WithGraphCompileOptions(WithInterruptBeforeNodes([]string{"1"}))) + _ = g.AddEdge(START, "1") + _ = g.AddEdge("1", "2") + _ = g.AddEdge("2", END) + + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "input", WithCheckPointID("1")) + info, existed := ExtractInterruptInfo(err) + assert.True(t, existed) + assert.Equal(t, []string{"1"}, info.SubGraphs["2"].BeforeNodes) + result, err := r.Invoke(ctx, "", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "input1sub1", result) +} + +func TestWithForceNewRun(t *testing.T) { + g := NewGraph[string, string]() + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + _ = g.AddEdge(START, "1") + _ = g.AddEdge("1", END) + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(&failStore{t: t})) + assert.NoError(t, err) + result, err := r.Invoke(ctx, "input", WithCheckPointID("1"), WithForceNewRun()) + assert.NoError(t, err) + assert.Equal(t, "input1", result) +} + +type failStore struct { + t *testing.T +} + +func (f *failStore) Get(_ context.Context, _ string) ([]byte, bool, error) { + f.t.Fatalf("cannot call store") + return nil, false, errors.New("fail") +} + +func (f *failStore) Set(_ context.Context, _ string, _ []byte) error { + f.t.Fatalf("cannot call store") + return errors.New("fail") +} + +func TestPreHandlerInterrupt(t *testing.T) { + type state struct{} + assert.NoError(t, serialization.GenericRegister[state]("_eino_TestPreHandlerInterrupt_state")) + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) state { + return state{} + })) + times := 0 + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + }), WithStatePreHandler(func(ctx context.Context, in string, state state) (string, error) { + if times == 0 { + times++ + return "", NewInterruptAndRerunErr("") + } + return in, nil + })) + _ = g.AddEdge(START, "1") + _ = g.AddEdge("1", END) + ctx := context.Background() + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + _, err = r.Invoke(ctx, "input", WithCheckPointID("1")) + info, existed := ExtractInterruptInfo(err) + assert.True(t, existed) + assert.Equal(t, []string{"1"}, info.RerunNodes) + result, err := r.Invoke(ctx, "", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "1", result) +} + +func TestCancelInterrupt(t *testing.T) { + g := NewGraph[string, string]() + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + time.Sleep(3 * time.Second) + return input + "1", nil + })) + _ = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "2", nil + })) + _ = g.AddEdge(START, "1") + _ = g.AddEdge("1", "2") + _ = g.AddEdge("2", END) + ctx := context.Background() + + // pregel + r, err := g.Compile(ctx, WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + // interrupt after nodes + canceledCtx, cancel := WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel(WithGraphInterruptTimeout(time.Hour)) + }() + _, err = r.Invoke(canceledCtx, "input", WithCheckPointID("1")) + assert.Error(t, err) + info, success := ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, []string{"1"}, info.AfterNodes) + result, err := r.Invoke(ctx, "input", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "input12", result) + // infinite timeout + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel() + }() + _, err = r.Invoke(canceledCtx, "input", WithCheckPointID("2")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, []string{"1"}, info.AfterNodes) + result, err = r.Invoke(ctx, "input", WithCheckPointID("2")) + assert.NoError(t, err) + assert.Equal(t, "input12", result) + + // interrupt rerun nodes - with auto-enabled PersistRerunInput, input is preserved + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + _, err = r.Invoke(canceledCtx, "input", WithCheckPointID("3")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, []string{"1"}, info.RerunNodes) + result, err = r.Invoke(ctx, "input", WithCheckPointID("3")) + assert.NoError(t, err) + assert.Equal(t, "input12", result) + + // dag + g = NewGraph[string, string]() + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + time.Sleep(3 * time.Second) + return input + "1", nil + })) + _ = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "2", nil + })) + _ = g.AddEdge(START, "1") + _ = g.AddEdge("1", "2") + _ = g.AddEdge("2", END) + r, err = g.Compile(ctx, WithNodeTriggerMode(AllPredecessor), WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + // interrupt after nodes + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel(WithGraphInterruptTimeout(time.Hour)) + }() + _, err = r.Invoke(canceledCtx, "input", WithCheckPointID("1")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, []string{"1"}, info.AfterNodes) + result, err = r.Invoke(ctx, "input", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, "input12", result) + // infinite timeout + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel() + }() + _, err = r.Invoke(canceledCtx, "input", WithCheckPointID("2")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, []string{"1"}, info.AfterNodes) + result, err = r.Invoke(ctx, "input", WithCheckPointID("2")) + assert.NoError(t, err) + assert.Equal(t, "input12", result) + + // interrupt rerun nodes - with auto-enabled PersistRerunInput, input is preserved + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(300 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + _, err = r.Invoke(canceledCtx, "input", WithCheckPointID("3")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, []string{"1"}, info.RerunNodes) + result, err = r.Invoke(ctx, "input", WithCheckPointID("3")) + assert.NoError(t, err) + assert.Equal(t, "input12", result) + + // dag multi canceled nodes + gg := NewGraph[string, map[string]any]() + _ = gg.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "1", nil + })) + _ = gg.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + time.Sleep(3 * time.Second) + return input + "2", nil + }), WithOutputKey("2")) + _ = gg.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + time.Sleep(3 * time.Second) + return input + "3", nil + }), WithOutputKey("3")) + _ = gg.AddLambdaNode("4", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) { + return input, nil + })) + _ = gg.AddEdge(START, "1") + _ = gg.AddEdge("1", "2") + _ = gg.AddEdge("1", "3") + _ = gg.AddEdge("2", "4") + _ = gg.AddEdge("3", "4") + _ = gg.AddEdge("4", END) + ctx = context.Background() + rr, err := gg.Compile(ctx, WithNodeTriggerMode(AllPredecessor), WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + // interrupt after nodes + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel(WithGraphInterruptTimeout(time.Hour)) + }() + _, err = rr.Invoke(canceledCtx, "input", WithCheckPointID("1")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, 2, len(info.AfterNodes)) + result2, err := rr.Invoke(ctx, "input", WithCheckPointID("1")) + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "2": "input12", + "3": "input13", + }, result2) + + // interrupt rerun nodes - with auto-enabled PersistRerunInput, input is preserved + canceledCtx, cancel = WithGraphInterrupt(ctx) + go func() { + time.Sleep(500 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + _, err = rr.Invoke(canceledCtx, "input", WithCheckPointID("2")) + assert.Error(t, err) + info, success = ExtractInterruptInfo(err) + assert.True(t, success) + assert.Equal(t, 2, len(info.RerunNodes)) + result2, err = rr.Invoke(ctx, "input", WithCheckPointID("2")) + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "2": "input12", + "3": "input13", + }, result2) +} + +func TestPersistRerunInputNonStream(t *testing.T) { + store := newInMemoryStore() + + var mu sync.Mutex + var receivedInput string + var callCount int + + g := NewGraph[string, string]() + + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + mu.Lock() + callCount++ + currentCount := callCount + receivedInput = input + mu.Unlock() + + if currentCount == 1 { + time.Sleep(2 * time.Second) + } + return input + "_processed", nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, + WithNodeTriggerMode(AllPredecessor), + WithCheckPointStore(store), + ) + assert.NoError(t, err) + + canceledCtx, cancel := WithGraphInterrupt(ctx) + go func() { + time.Sleep(100 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + + _, err = r.Invoke(canceledCtx, "test_input", WithCheckPointID("cp1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, []string{"1"}, info.RerunNodes) + + mu.Lock() + assert.Equal(t, "test_input", receivedInput) + mu.Unlock() + + result, err := r.Invoke(ctx, "", WithCheckPointID("cp1")) + assert.NoError(t, err) + assert.Equal(t, "test_input_processed", result) + + mu.Lock() + assert.Equal(t, "test_input", receivedInput) + assert.Equal(t, 2, callCount) + mu.Unlock() +} + +func TestPersistRerunInputStream(t *testing.T) { + store := newInMemoryStore() + + var mu sync.Mutex + var receivedInput string + var callCount int + + g := NewGraph[string, string]() + + err := g.AddLambdaNode("1", TransformableLambda(func(ctx context.Context, input *schema.StreamReader[string]) (output *schema.StreamReader[string], err error) { + mu.Lock() + callCount++ + currentCount := callCount + mu.Unlock() + + var sb string + for { + chunk, err := input.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + sb += chunk + } + + mu.Lock() + receivedInput = sb + mu.Unlock() + + if currentCount == 1 { + time.Sleep(2 * time.Second) + } + + return schema.StreamReaderFromArray([]string{sb + "_processed"}), nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, + WithNodeTriggerMode(AllPredecessor), + WithCheckPointStore(store), + ) + assert.NoError(t, err) + + inputStream := schema.StreamReaderFromArray([]string{"chunk1", "chunk2", "chunk3"}) + + canceledCtx, cancel := WithGraphInterrupt(ctx) + go func() { + time.Sleep(100 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + + _, err = r.Transform(canceledCtx, inputStream, WithCheckPointID("cp1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, []string{"1"}, info.RerunNodes) + + mu.Lock() + assert.Equal(t, "chunk1chunk2chunk3", receivedInput) + mu.Unlock() + + emptyInputStream := schema.StreamReaderFromArray([]string{}) + + resultStream, err := r.Transform(ctx, emptyInputStream, WithCheckPointID("cp1")) + assert.NoError(t, err) + + var result string + for { + chunk, err := resultStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + result += chunk + } + + assert.Equal(t, "chunk1chunk2chunk3_processed", result) + + mu.Lock() + assert.Equal(t, "chunk1chunk2chunk3", receivedInput) + assert.Equal(t, 2, callCount) + mu.Unlock() +} + +type testPersistRerunInputState struct { + Prefix string +} + +func TestPersistRerunInputWithPreHandler(t *testing.T) { + store := newInMemoryStore() + + var mu sync.Mutex + var receivedInput string + var callCount int + + schema.Register[testPersistRerunInputState]() + + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) *testPersistRerunInputState { + return &testPersistRerunInputState{Prefix: "prefix_"} + })) + + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + mu.Lock() + callCount++ + currentCount := callCount + receivedInput = input + mu.Unlock() + + if currentCount == 1 { + time.Sleep(2 * time.Second) + } + return input + "_processed", nil + }), WithStatePreHandler(func(ctx context.Context, in string, s *testPersistRerunInputState) (string, error) { + return s.Prefix + in, nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, + WithNodeTriggerMode(AllPredecessor), + WithCheckPointStore(store), + ) + assert.NoError(t, err) + + canceledCtx, cancel := WithGraphInterrupt(ctx) + go func() { + time.Sleep(100 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + + _, err = r.Invoke(canceledCtx, "test_input", WithCheckPointID("cp1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + if ok { + assert.Equal(t, []string{"1"}, info.RerunNodes) + } + + mu.Lock() + assert.Equal(t, "prefix_test_input", receivedInput) + mu.Unlock() + + result, err := r.Invoke(ctx, "", WithCheckPointID("cp1")) + assert.NoError(t, err) + assert.Equal(t, "prefix_test_input_processed", result) + + mu.Lock() + assert.Equal(t, "prefix_test_input", receivedInput) + assert.Equal(t, 2, callCount) + mu.Unlock() +} + +func TestPersistRerunInputBackwardCompatibility(t *testing.T) { + store := newInMemoryStore() + + var receivedInput string + var callCount int + + g := NewGraph[string, string]() + + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + callCount++ + receivedInput = input + if len(input) > 0 { + return "", StatefulInterrupt(ctx, "interrupt", input) + } + + _, _, restoredInput := GetInterruptState[string](ctx) + return restoredInput + "_processed", nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, + WithNodeTriggerMode(AllPredecessor), + WithCheckPointStore(store), + ) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "test_input", WithCheckPointID("cp1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, []string{"1"}, info.RerunNodes) + + assert.Equal(t, "test_input", receivedInput) + + result, err := r.Invoke(ctx, "", WithCheckPointID("cp1")) + assert.NoError(t, err) + assert.Equal(t, "test_input_processed", result) + assert.Equal(t, "", receivedInput) + assert.Equal(t, 2, callCount) +} + +func TestPersistRerunInputSubGraph(t *testing.T) { + store := newInMemoryStore() + + var mu sync.Mutex + var receivedInput string + var callCount int + + subG := NewGraph[string, string]() + err := subG.AddLambdaNode("sub1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + mu.Lock() + callCount++ + currentCount := callCount + receivedInput = input + mu.Unlock() + + if currentCount == 1 { + time.Sleep(2 * time.Second) + } + return input + "_sub_processed", nil + })) + assert.NoError(t, err) + err = subG.AddEdge(START, "sub1") + assert.NoError(t, err) + err = subG.AddEdge("sub1", END) + assert.NoError(t, err) + + g := NewGraph[string, string]() + err = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input + "_main", nil + })) + assert.NoError(t, err) + err = g.AddGraphNode("2", subG) + assert.NoError(t, err) + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", "2") + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, + WithNodeTriggerMode(AllPredecessor), + WithCheckPointStore(store), + ) + assert.NoError(t, err) + + canceledCtx, cancel := WithGraphInterrupt(ctx) + go func() { + time.Sleep(100 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + + _, err = r.Invoke(canceledCtx, "test", WithCheckPointID("cp1")) + assert.NotNil(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok, "Expected interrupt error, got: %v", err) + if len(info.SubGraphs) > 0 { + assert.Contains(t, info.SubGraphs, "2") + subInfo := info.SubGraphs["2"] + assert.Equal(t, []string{"sub1"}, subInfo.RerunNodes) + } else { + assert.Equal(t, []string{"2"}, info.RerunNodes) + } + + mu.Lock() + assert.Equal(t, "test_main", receivedInput) + mu.Unlock() + + result, err := r.Invoke(ctx, "", WithCheckPointID("cp1")) + assert.NoError(t, err) + assert.Equal(t, "test_main_sub_processed", result) + + mu.Lock() + assert.Equal(t, "test_main", receivedInput) + assert.Equal(t, 2, callCount) + mu.Unlock() +} + +type longRunningToolInput struct { + Input string `json:"input"` +} + +func TestToolsNodeWithExternalGraphInterrupt(t *testing.T) { + store := newInMemoryStore() + ctx := context.Background() + + var mu sync.Mutex + var callCount int + + longRunningToolInfo := &schema.ToolInfo{ + Name: "long_running_tool", + Desc: "A tool that takes a long time to run", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: "string", Desc: "input"}, + }), + } + + longRunningTool := newCheckpointTestTool(longRunningToolInfo, func(ctx context.Context, in *longRunningToolInput) (string, error) { + mu.Lock() + callCount++ + currentCount := callCount + mu.Unlock() + + if currentCount == 1 { + time.Sleep(2 * time.Second) + } + return "result_" + in.Input, nil + }) + + toolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{longRunningTool}, + }) + assert.NoError(t, err) + + g := NewGraph[*schema.Message, []*schema.Message]() + err = g.AddToolsNode("tools", toolsNode) + assert.NoError(t, err) + err = g.AddEdge(START, "tools") + assert.NoError(t, err) + err = g.AddEdge("tools", END) + assert.NoError(t, err) + + r, err := g.Compile(ctx, + WithNodeTriggerMode(AllPredecessor), + WithCheckPointStore(store), + ) + assert.NoError(t, err) + + inputMsg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{{ + ID: "call_1", + Type: "function", + Function: schema.FunctionCall{ + Name: "long_running_tool", + Arguments: `{"input": "test"}`, + }, + }}, + } + + canceledCtx, cancel := WithGraphInterrupt(ctx) + go func() { + time.Sleep(100 * time.Millisecond) + cancel(WithGraphInterruptTimeout(0)) + }() + + _, err = r.Invoke(canceledCtx, inputMsg, WithCheckPointID("cp1")) + assert.Error(t, err) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok, "Expected interrupt error, got: %v", err) + if ok { + assert.Equal(t, []string{"tools"}, info.RerunNodes) + } + + result, err := r.Invoke(ctx, &schema.Message{}, WithCheckPointID("cp1")) + assert.NoError(t, err) + assert.Len(t, result, 1) + assert.Equal(t, `"result_test"`, result[0].Content) + + mu.Lock() + assert.Equal(t, 2, callCount) + mu.Unlock() +} + +type checkpointTestTool[I, O any] struct { + info *schema.ToolInfo + fn func(ctx context.Context, in I) (O, error) +} + +func newCheckpointTestTool[I, O any](info *schema.ToolInfo, f func(ctx context.Context, in I) (O, error)) tool.InvokableTool { + return &checkpointTestTool[I, O]{ + info: info, + fn: f, + } +} + +func (f *checkpointTestTool[I, O]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return f.info, nil +} + +func (f *checkpointTestTool[I, O]) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + t := generic.NewInstance[I]() + err := sonic.UnmarshalString(argumentsInJSON, t) + if err != nil { + return "", err + } + o, err := f.fn(ctx, t) + if err != nil { + return "", err + } + return sonic.MarshalString(o) +} diff --git a/compose/component_to_graph_node.go b/compose/component_to_graph_node.go new file mode 100644 index 0000000..4bd27fe --- /dev/null +++ b/compose/component_to_graph_node.go @@ -0,0 +1,212 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" +) + +func toComponentNode[I, O, TOption any]( + node any, + componentType component, + invoke Invoke[I, O, TOption], + stream Stream[I, O, TOption], + collect Collect[I, O, TOption], + transform Transform[I, O, TOption], + opts ...GraphAddNodeOpt, +) (*graphNode, *graphAddNodeOpts) { + meta := parseExecutorInfoFromComponent(componentType, node) + info, options := getNodeInfo(opts...) + run := runnableLambda(invoke, stream, collect, transform, + !meta.isComponentCallbackEnabled, + ) + + gn := toNode(info, run, nil, meta, node, opts...) + + return gn, options +} + +func toEmbeddingNode(node embedding.Embedder, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfEmbedding, + node.EmbedStrings, + nil, + nil, + nil, + opts...) +} + +func toRetrieverNode(node retriever.Retriever, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfRetriever, + node.Retrieve, + nil, + nil, + nil, + opts...) +} + +func toLoaderNode(node document.Loader, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfLoader, + node.Load, + nil, + nil, + nil, + opts...) +} + +func toIndexerNode(node indexer.Indexer, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfIndexer, + node.Store, + nil, + nil, + nil, + opts...) +} + +func toChatModelNode(node model.BaseChatModel, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfChatModel, + node.Generate, + node.Stream, + nil, + nil, + opts...) +} + +func toAgenticModelNode(node model.AgenticModel, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfAgenticModel, + node.Generate, + node.Stream, + nil, nil, + opts..., + ) +} + +func toChatTemplateNode(node prompt.ChatTemplate, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfPrompt, + node.Format, + nil, + nil, + nil, + opts...) +} + +func toAgenticChatTemplateNode(node prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfAgenticPrompt, + node.Format, + nil, nil, nil, + opts..., + ) +} + +func toDocumentTransformerNode(node document.Transformer, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + components.ComponentOfTransformer, + node.Transform, + nil, + nil, + nil, + opts...) +} + +func toToolsNode(node *ToolsNode, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + ComponentOfToolsNode, + node.Invoke, + node.Stream, + nil, + nil, + opts...) +} + +func toAgenticToolsNode(node *AgenticToolsNode, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + return toComponentNode( + node, + ComponentOfAgenticToolsNode, + node.Invoke, + node.Stream, + nil, nil, + opts..., + ) +} + +func toLambdaNode(node *Lambda, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + info, options := getNodeInfo(opts...) + + gn := toNode(info, node.executor, nil, node.executor.meta, node, opts...) + + return gn, options +} + +func toAnyGraphNode(node AnyGraph, opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + meta := parseExecutorInfoFromComponent(node.component(), node) + info, options := getNodeInfo(opts...) + + gn := toNode(info, nil, node, meta, node, opts...) + + return gn, options +} + +func toPassthroughNode(opts ...GraphAddNodeOpt) (*graphNode, *graphAddNodeOpts) { + node := composablePassthrough() + info, options := getNodeInfo(opts...) + gn := toNode(info, node, nil, node.meta, node, opts...) + return gn, options +} + +func toNode(nodeInfo *nodeInfo, executor *composableRunnable, graph AnyGraph, + meta *executorMeta, instance any, opts ...GraphAddNodeOpt) *graphNode { + + if meta == nil { + meta = &executorMeta{} + } + + gn := &graphNode{ + nodeInfo: nodeInfo, + + cr: executor, + g: graph, + executorMeta: meta, + + instance: instance, + opts: opts, + } + + return gn +} diff --git a/compose/dag.go b/compose/dag.go new file mode 100644 index 0000000..98b85ae --- /dev/null +++ b/compose/dag.go @@ -0,0 +1,195 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "fmt" +) + +func dagChannelBuilder(controlDependencies []string, dataDependencies []string, zeroValue func() any, emptyStream func() streamReader) channel { + deps := make(map[string]dependencyState, len(controlDependencies)) + for _, dep := range controlDependencies { + deps[dep] = dependencyStateWaiting + } + indirect := make(map[string]bool, len(dataDependencies)) + for _, dep := range dataDependencies { + indirect[dep] = false + } + + return &dagChannel{ + Values: make(map[string]any), + ControlPredecessors: deps, + DataPredecessors: indirect, + zeroValue: zeroValue, + emptyStream: emptyStream, + } +} + +type dependencyState uint8 + +const ( + dependencyStateWaiting dependencyState = iota + dependencyStateReady + dependencyStateSkipped +) + +type dagChannel struct { + zeroValue func() any + emptyStream func() streamReader + + ControlPredecessors map[string]dependencyState + Values map[string]any + DataPredecessors map[string]bool // if all dependencies have been skipped, indirect dependencies won't effect. + Skipped bool + + mergeConfig FanInMergeConfig +} + +func (ch *dagChannel) setMergeConfig(cfg FanInMergeConfig) { + ch.mergeConfig.StreamMergeWithSourceEOF = cfg.StreamMergeWithSourceEOF +} + +func (ch *dagChannel) load(c channel) error { + dc, ok := c.(*dagChannel) + if !ok { + return fmt.Errorf("load dag channel fail, got %T, want *dagChannel", c) + } + ch.ControlPredecessors = dc.ControlPredecessors + ch.DataPredecessors = dc.DataPredecessors + ch.Skipped = dc.Skipped + ch.Values = dc.Values + return nil +} + +func (ch *dagChannel) reportValues(ins map[string]any) error { + if ch.Skipped { + return nil + } + + for k, v := range ins { + if _, ok := ch.DataPredecessors[k]; !ok { + continue + } + ch.DataPredecessors[k] = true + ch.Values[k] = v + } + return nil +} + +func (ch *dagChannel) reportDependencies(dependencies []string) { + if ch.Skipped { + return + } + + for _, dep := range dependencies { + if _, ok := ch.ControlPredecessors[dep]; ok { + ch.ControlPredecessors[dep] = dependencyStateReady + } + } + return +} + +func (ch *dagChannel) reportSkip(keys []string) bool { + for _, k := range keys { + if _, ok := ch.ControlPredecessors[k]; ok { + ch.ControlPredecessors[k] = dependencyStateSkipped + } + if _, ok := ch.DataPredecessors[k]; ok { + ch.DataPredecessors[k] = true + } + } + + allSkipped := true + for _, state := range ch.ControlPredecessors { + if state != dependencyStateSkipped { + allSkipped = false + break + } + } + ch.Skipped = allSkipped + + return allSkipped +} + +func (ch *dagChannel) get(isStream bool, name string, edgeHandler *edgeHandlerManager) ( + any, bool, error) { + if ch.Skipped { + return nil, false, nil + } + + if len(ch.ControlPredecessors)+len(ch.DataPredecessors) == 0 { + return nil, false, nil + } + + for _, state := range ch.ControlPredecessors { + if state == dependencyStateWaiting { + return nil, false, nil + } + } + for _, ready := range ch.DataPredecessors { + if !ready { + return nil, false, nil + } + } + + defer func() { + ch.Values = make(map[string]any) + for k := range ch.ControlPredecessors { + ch.ControlPredecessors[k] = dependencyStateWaiting + } + for k := range ch.DataPredecessors { + ch.DataPredecessors[k] = false + } + }() + + valueList := make([]any, len(ch.Values)) + names := make([]string, len(ch.Values)) + i := 0 + for k, value := range ch.Values { + resolvedV, err := edgeHandler.handle(k, name, value, isStream) + if err != nil { + return nil, false, err + } + valueList[i] = resolvedV + names[i] = k + i++ + } + + if len(valueList) == 0 { + if isStream { + return ch.emptyStream(), true, nil + } + return ch.zeroValue(), true, nil + } + if len(valueList) == 1 { + return valueList[0], true, nil + } + + mergeOpts := &mergeOptions{ + streamMergeWithSourceEOF: ch.mergeConfig.StreamMergeWithSourceEOF, + names: names, + } + v, err := mergeValues(valueList, mergeOpts) + if err != nil { + return nil, false, err + } + return v, true, nil +} + +func (ch *dagChannel) convertValues(fn func(map[string]any) error) error { + return fn(ch.Values) +} diff --git a/compose/dag_test.go b/compose/dag_test.go new file mode 100644 index 0000000..6185ec4 --- /dev/null +++ b/compose/dag_test.go @@ -0,0 +1,223 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "io" + "testing" +) + +func TestDAG(t *testing.T) { + var err error + + g := NewGraph[string, string]() + err = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithOutputKey("1")) + if err != nil { + t.Fatal(err) + } + + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithOutputKey("2")) + if err != nil { + t.Fatal(err) + } + + err = g.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) { + if _, ok := input["1"]; !ok { + return "", fmt.Errorf("node 1 output fail: %+v", input) + } + if _, ok := input["2"]; !ok { + return "", fmt.Errorf("node 2 output fail: %+v", input) + } + return input["1"].(string) + input["2"].(string), nil + })) + if err != nil { + t.Fatal(err) + } + + err = g.AddLambdaNode("4", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + if err != nil { + t.Fatal(err) + } + + err = g.AddLambdaNode("5", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + if err != nil { + t.Fatal(err) + } + + err = g.AddLambdaNode("6", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithOutputKey("6")) + if err != nil { + t.Fatal(err) + } + + err = g.AddLambdaNode("7", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) { + if _, ok := input["1"]; !ok { + return "", fmt.Errorf("7:node 1 output fail: %+v", input) + } + if _, ok := input["6"]; !ok { + return "", fmt.Errorf("7:node 6 output fail: %+v", input) + } + return input["1"].(string) + input["6"].(string), nil + })) + if err != nil { + t.Fatal(err) + } + + err = g.AddEdge("1", "3") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("2", "3") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("3", "4") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("4", "5") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("4", "6") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("6", "7") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", "7") + if err != nil { + t.Fatal(err) + } + + err = g.AddEdge(START, "1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "2") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("7", END) + if err != nil { + t.Fatal(err) + } + + runner, err := g.Compile(context.Background(), WithNodeTriggerMode(AllPredecessor)) + if err != nil { + t.Fatal(err) + } + + // success + ctx := context.Background() + out, err := runner.Invoke(ctx, "hello") + if err != nil { + t.Fatal(err) + } + if out != "hellohellohello" { + t.Fatalf("node7 fail") + } + + result, err := runner.Invoke(ctx, "1") + if err != nil { + t.Fatal(err) + } + if result != "111" { + t.Fatalf("runner invoke fail, output: %s", result) + } + streamResult, err := runner.Stream(ctx, "1") + if err != nil { + t.Fatal(err) + } + defer streamResult.Close() + ret := "" + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + ret += chunk + } + if ret != "111" { + t.Fatalf("runner stream fail, output: %s", ret) + } + + // loop + gg := NewGraph[string, map[string]any]() + err = gg.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithOutputKey("1")) + if err != nil { + t.Fatal(err) + } + + err = gg.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input map[string]any) (output string, err error) { + return input["1"].(string), nil + })) + if err != nil { + t.Fatal(err) + } + + err = gg.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithOutputKey("3")) + if err != nil { + t.Fatal(err) + } + + err = gg.AddEdge("1", "2") + if err != nil { + t.Fatal(err) + } + err = gg.AddEdge("2", "3") + if err != nil { + t.Fatal(err) + } + err = gg.AddEdge("3", "2") + if err != nil { + t.Fatal(err) + } + err = gg.AddEdge(START, "1") + if err != nil { + t.Fatal(err) + } + err = gg.AddEdge("3", END) + if err != nil { + t.Fatal(err) + } + + _, err = gg.compile(ctx, &graphCompileOptions{nodeTriggerMode: AllPredecessor}) + if err == nil { + t.Fatal("cannot validate loop") + } +} diff --git a/compose/doc.go b/compose/doc.go new file mode 100644 index 0000000..5057f46 --- /dev/null +++ b/compose/doc.go @@ -0,0 +1,19 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package compose provides graph and workflow primitives to build +// composable, interruptible execution pipelines with callback support. +package compose diff --git a/compose/error.go b/compose/error.go new file mode 100644 index 0000000..18d7078 --- /dev/null +++ b/compose/error.go @@ -0,0 +1,111 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "errors" + "fmt" + "reflect" + "strings" +) + +// ErrExceedMaxSteps graph will throw this error when the number of steps exceeds the maximum number of steps. +var ErrExceedMaxSteps = errors.New("exceeds max steps") + +func newUnexpectedInputTypeErr(expected reflect.Type, got reflect.Type) error { + return fmt.Errorf("unexpected input type. expected: %v, got: %v", expected, got) +} + +type defaultImplAction string + +const ( + actionInvokeByStream defaultImplAction = "InvokeByStream" + actionInvokeByCollect defaultImplAction = "InvokeByCollect" + actionInvokeByTransform defaultImplAction = "InvokeByTransform" + actionStreamByInvoke defaultImplAction = "StreamByInvoke" + actionStreamByTransform defaultImplAction = "StreamByTransform" + actionStreamByCollect defaultImplAction = "StreamByCollect" + actionCollectByTransform defaultImplAction = "CollectByTransform" + actionCollectByInvoke defaultImplAction = "CollectByInvoke" + actionCollectByStream defaultImplAction = "CollectByStream" + actionTransformByStream defaultImplAction = "TransformByStream" + actionTransformByCollect defaultImplAction = "TransformByCollect" + actionTransformByInvoke defaultImplAction = "TransformByInvoke" +) + +func newStreamReadError(err error) error { + return fmt.Errorf("failed to read from stream. error: %w", err) +} + +func newGraphRunError(err error) error { + return &internalError{ + typ: internalErrorTypeGraphRun, + nodePath: NodePath{}, + origError: err, + } +} + +func wrapGraphNodeError(nodeKey string, err error) error { + if ok := isInterruptError(err); ok { + return err + } + var ie *internalError + ok := errors.As(err, &ie) + if !ok { + return &internalError{ + typ: internalErrorTypeNodeRun, + nodePath: NodePath{path: []string{nodeKey}}, + origError: err, + } + } + ie.nodePath.path = append([]string{nodeKey}, ie.nodePath.path...) + return ie +} + +type internalErrorType string + +const ( + internalErrorTypeNodeRun = "NodeRunError" + internalErrorTypeGraphRun = "GraphRunError" +) + +type internalError struct { + typ internalErrorType + nodePath NodePath + origError error +} + +func (i *internalError) Error() string { + sb := strings.Builder{} + sb.WriteString(string("[" + i.typ + "] ")) + sb.WriteString(i.origError.Error()) + if len(i.nodePath.path) > 0 { + sb.WriteString("\n------------------------\n") + sb.WriteString("node path: [") + for j := 0; j < len(i.nodePath.path)-1; j++ { + sb.WriteString(i.nodePath.path[j] + ", ") + } + sb.WriteString(i.nodePath.path[len(i.nodePath.path)-1]) + sb.WriteString("]") + } + sb.WriteString("") + return sb.String() +} + +func (i *internalError) Unwrap() error { + return i.origError +} diff --git a/compose/error_test.go b/compose/error_test.go new file mode 100644 index 0000000..e7f140b --- /dev/null +++ b/compose/error_test.go @@ -0,0 +1,137 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestCommonError(t *testing.T) { + g := NewGraph[string, string]() + assert.NoError(t, g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", errors.New("my error") + }))) + assert.NoError(t, g.AddEdge(START, "1")) + assert.NoError(t, g.AddEdge("1", END)) + + ctx := context.Background() + r, err := g.Compile(ctx) + assert.NoError(t, err) + + // node error + _, err = r.Invoke(ctx, "input") + var ie *internalError + + assert.True(t, errors.As(err, &ie)) + assert.Equal(t, "my error", ie.origError.Error()) + + // wrapper error + sr, sw := schema.Pipe[string](0) + sw.Close() + _, err = r.Transform(ctx, sr) + assert.True(t, errors.As(err, &ie)) + assert.ErrorContains(t, ie.origError, "stream reader is empty, concat fail") + assert.Equal(t, []string{"1"}, ie.nodePath.path) + println(err.Error()) +} + +func TestSubGraphNodeError(t *testing.T) { + subG := NewGraph[string, string]() + assert.NoError(t, subG.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", errors.New("my error") + }))) + assert.NoError(t, subG.AddEdge(START, "1")) + assert.NoError(t, subG.AddEdge("1", END)) + + g := NewGraph[string, string]() + assert.NoError(t, g.AddGraphNode("a", subG)) + assert.NoError(t, g.AddEdge(START, "a")) + assert.NoError(t, g.AddEdge("a", END)) + + ctx := context.Background() + r, err := g.Compile(ctx) + assert.NoError(t, err) + _, err = r.Invoke(ctx, "input") + var ie *internalError + assert.True(t, errors.As(err, &ie)) + assert.Equal(t, "my error", ie.origError.Error()) + assert.Equal(t, []string{"a", "1"}, ie.nodePath.path) +} + +func TestContextCancelDuringRun(t *testing.T) { + // Create a graph with a long-running node to test context cancellation + g := NewGraph[string, string]() + + // Add a node that waits for some time (long enough to be cancelled) + assert.NoError(t, g.AddLambdaNode("slow_node", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + select { + case <-ctx.Done(): + // Return context's error when cancelled + return "", ctx.Err() + case <-time.After(200 * time.Millisecond): + return input + "_processed", nil + } + }))) + + assert.NoError(t, g.AddEdge(START, "slow_node")) + assert.NoError(t, g.AddEdge("slow_node", END)) + + // Create a context that we can cancel + ctx, cancel := context.WithCancel(context.Background()) + + // Compile the graph + r, err := g.Compile(ctx) + assert.NoError(t, err) + + // Run the invoke in a goroutine + resultCh := make(chan error) + go func() { + _, err := r.Invoke(ctx, "input") + resultCh <- err + }() + + // Cancel the context after a short delay + time.Sleep(50 * time.Millisecond) + cancel() + + // Get the result + err = <-resultCh + + // Verify the error is related to context cancellation + assert.Error(t, err) + + // Check error type and content + var ie *internalError + assert.True(t, errors.As(err, &ie)) + + // Error path should contain the node + assert.Equal(t, []string{"slow_node"}, ie.nodePath.path) + + // Original error should be context.Canceled + assert.ErrorIs(t, ie.origError, context.Canceled) + + // Test unwrap capability + unwrappedErr := ie.Unwrap() + assert.ErrorIs(t, unwrappedErr, context.Canceled) +} diff --git a/compose/field_mapping.go b/compose/field_mapping.go new file mode 100644 index 0000000..aaf8de7 --- /dev/null +++ b/compose/field_mapping.go @@ -0,0 +1,774 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "errors" + "fmt" + "reflect" + "runtime/debug" + "strings" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +type FieldMapping struct { + fromNodeKey string + from string + to string + + customExtractor func(input any) (any, error) +} + +// String returns the string representation of the FieldMapping. +func (m *FieldMapping) String() string { + var sb strings.Builder + sb.WriteString("[from ") + + if m.from != "" { + sb.WriteString(m.from) + sb.WriteString("(field) of ") + } + + sb.WriteString(m.fromNodeKey) + + if m.to != "" { + sb.WriteString(" to ") + sb.WriteString(m.to) + sb.WriteString("(field)") + } + + sb.WriteString("]") + return sb.String() +} + +// FromField creates a FieldMapping that maps a single predecessor field to the entire successor input. +// This is an exclusive mapping - once set, no other field mappings can be added since the successor input +// has already been fully mapped. +// Field: either the field of a struct, or the key of a map. +func FromField(from string) *FieldMapping { + return &FieldMapping{ + from: from, + } +} + +// ToField creates a FieldMapping that maps the entire predecessor output to a single successor field. +// Field: either the field of a struct, or the key of a map. +func ToField(to string, opts ...FieldMappingOption) *FieldMapping { + fm := &FieldMapping{ + to: to, + } + for _, opt := range opts { + opt(fm) + } + return fm +} + +// MapFields creates a FieldMapping that maps a single predecessor field to a single successor field. +// Field: either the field of a struct, or the key of a map. +func MapFields(from, to string) *FieldMapping { + return &FieldMapping{ + from: from, + to: to, + } +} + +func (m *FieldMapping) FromNodeKey() string { + return m.fromNodeKey +} + +func (m *FieldMapping) FromPath() FieldPath { + return splitFieldPath(m.from) +} + +func (m *FieldMapping) ToPath() FieldPath { + return splitFieldPath(m.to) +} + +func (m *FieldMapping) Equals(o *FieldMapping) bool { + if m == nil { + return o == nil + } + + if o == nil || m.customExtractor != nil || o.customExtractor != nil { + return false + } + + return m.from == o.from && m.to == o.to && m.fromNodeKey == o.fromNodeKey +} + +// FieldPath represents a path to a nested field in a struct or map. +// Each element in the path is either: +// - a struct field name +// - a map key +// +// Example paths: +// - []string{"user"} // top-level field +// - []string{"user", "name"} // nested struct field +// - []string{"users", "admin"} // map key access +type FieldPath []string + +func (fp *FieldPath) join() string { + return strings.Join(*fp, pathSeparator) +} + +func splitFieldPath(path string) FieldPath { + p := strings.Split(path, pathSeparator) + if len(p) == 1 && p[0] == "" { + return FieldPath{} + } + + return p +} + +// pathSeparator is a special character (Unit Separator) used internally to join path elements. +// This character is chosen because it's extremely unlikely to appear in user-defined field names or map keys. +const pathSeparator = "\x1F" + +// FromFieldPath creates a FieldMapping that maps a single predecessor field path to the entire successor input. +// This is an exclusive mapping - once set, no other field mappings can be added since the successor input +// has already been fully mapped. +// +// Example: +// +// // Maps the 'name' field from nested 'user.profile' to the entire successor input +// FromFieldPath(FieldPath{"user", "profile", "name"}) +// +// Note: The field path elements must not contain the internal path separator character ('\x1F'). +func FromFieldPath(fromFieldPath FieldPath) *FieldMapping { + return &FieldMapping{ + from: fromFieldPath.join(), + } +} + +// ToFieldPath creates a FieldMapping that maps the entire predecessor output to a single successor field path. +// +// Example: +// +// // Maps the entire predecessor output to response.data.userName +// ToFieldPath(FieldPath{"response", "data", "userName"}) +// +// Note: The field path elements must not contain the internal path separator character ('\x1F'). +func ToFieldPath(toFieldPath FieldPath, opts ...FieldMappingOption) *FieldMapping { + fm := &FieldMapping{ + to: toFieldPath.join(), + } + for _, opt := range opts { + opt(fm) + } + return fm +} + +// MapFieldPaths creates a FieldMapping that maps a single predecessor field path to a single successor field path. +// +// Example: +// +// // Maps user.profile.name to response.userName +// MapFieldPaths( +// FieldPath{"user", "profile", "name"}, +// FieldPath{"response", "userName"}, +// ) +// +// Note: The field path elements must not contain the internal path separator character ('\x1F'). +func MapFieldPaths(fromFieldPath, toFieldPath FieldPath) *FieldMapping { + return &FieldMapping{ + from: fromFieldPath.join(), + to: toFieldPath.join(), + } +} + +// FieldMappingOption is a functional option for configuring a FieldMapping. +type FieldMappingOption func(*FieldMapping) + +// WithCustomExtractor sets a custom extractor function for the FieldMapping. +// The extractor function is used to extract a value from the 'source' of the FieldMapping. +// NOTE: if specified in this way, Eino can only check the validity of the field mapping at request time.. +func WithCustomExtractor(extractor func(input any) (any, error)) FieldMappingOption { + return func(m *FieldMapping) { + m.customExtractor = extractor + } +} + +func (m *FieldMapping) targetPath() FieldPath { + return splitFieldPath(m.to) +} + +func buildFieldMappingConverter[I any]() func(input any) (any, error) { + return func(input any) (any, error) { + in, ok := input.(map[string]any) + if !ok { + panic(newUnexpectedInputTypeErr(reflect.TypeOf(map[string]any{}), reflect.TypeOf(input))) + } + + return convertTo(in, generic.TypeOf[I]()), nil + } +} + +func buildStreamFieldMappingConverter[I any]() func(input streamReader) streamReader { + return func(input streamReader) streamReader { + s, ok := unpackStreamReader[map[string]any](input) + if !ok { + panic("mappingStreamAssign incoming streamReader chunk type not map[string]any") + } + + return packStreamReader(schema.StreamReaderWithConvert(s, func(v map[string]any) (I, error) { + t := convertTo(v, generic.TypeOf[I]()) + return t.(I), nil + })) + } +} + +func convertTo(mappings map[string]any, typ reflect.Type) any { + tValue := newInstanceByType(typ) + if !tValue.CanAddr() { + tValue = newInstanceByType(reflect.PointerTo(typ)).Elem() + } + + for mapping, taken := range mappings { + tValue = assignOne(tValue, taken, mapping) + } + + return tValue.Interface() +} + +func assignOne(destValue reflect.Value, taken any, to string) reflect.Value { + if len(to) == 0 { // assign to output directly + destValue.Set(reflect.ValueOf(taken)) + return destValue + } + + var ( + toPaths = splitFieldPath(to) + originalDestValue = destValue + parentMap reflect.Value + parentKey string + ) + + for { + path := toPaths[0] + toPaths = toPaths[1:] + if len(toPaths) == 0 { + toSet := reflect.ValueOf(taken) + + if destValue.Type() == reflect.TypeOf((*any)(nil)).Elem() { + existingMap, ok := destValue.Interface().(map[string]any) + if ok { + destValue = reflect.ValueOf(existingMap) + } else { + mapValue := reflect.MakeMap(reflect.TypeOf(map[string]any{})) + destValue.Set(mapValue) + destValue = mapValue + } + } + + if destValue.Kind() == reflect.Map { + key := reflect.ValueOf(path) + keyType := destValue.Type().Key() + if keyType != strType { + key = key.Convert(keyType) + } + + if !toSet.IsValid() { + toSet = reflect.Zero(destValue.Type().Elem()) + } + destValue.SetMapIndex(key, toSet) + + if parentMap.IsValid() { + parentMap.SetMapIndex(reflect.ValueOf(parentKey), destValue) + } + + return originalDestValue + } + + ptrValue := destValue + for destValue.Kind() == reflect.Ptr { + destValue = destValue.Elem() + } + + if !toSet.IsValid() { + // just skip it, because this 'nil' is the zero value of the corresponding struct field + } else { + field := destValue.FieldByName(path) + field.Set(toSet) + } + + if parentMap.IsValid() { + parentMap.SetMapIndex(reflect.ValueOf(parentKey), ptrValue) + } + + return originalDestValue + } + + if destValue.Type() == reflect.TypeOf((*any)(nil)).Elem() { + existingMap, ok := destValue.Interface().(map[string]any) + if ok { + destValue = reflect.ValueOf(existingMap) + } else { + mapValue := reflect.MakeMap(reflect.TypeOf(map[string]any{})) + destValue.Set(mapValue) + destValue = mapValue + } + } + + if destValue.Kind() == reflect.Map { + keyValue := reflect.ValueOf(path) + valueValue := destValue.MapIndex(keyValue) + if !valueValue.IsValid() { + valueValue = newInstanceByType(destValue.Type().Elem()) + destValue.SetMapIndex(keyValue, valueValue) + } + + if parentMap.IsValid() { + parentMap.SetMapIndex(reflect.ValueOf(parentKey), destValue) + } + + parentMap = destValue + parentKey = path + destValue = valueValue + + continue + } + + ptrValue := destValue + for destValue.Kind() == reflect.Ptr { + destValue = destValue.Elem() + } + + field := destValue.FieldByName(path) + instantiateIfNeeded(field) + + if parentMap.IsValid() { + parentMap.SetMapIndex(reflect.ValueOf(parentKey), ptrValue) + parentMap = reflect.Value{} + parentKey = "" + } + + destValue = field + } +} + +func instantiateIfNeeded(field reflect.Value) { + if field.Kind() == reflect.Ptr { + if field.IsNil() { + field.Set(reflect.New(field.Type().Elem())) + } + } else if field.Kind() == reflect.Map { + if field.IsNil() { + field.Set(reflect.MakeMap(field.Type())) + } + } +} + +func newInstanceByType(typ reflect.Type) reflect.Value { + switch typ.Kind() { + case reflect.Map: + return reflect.MakeMap(typ) + case reflect.Slice, reflect.Array: + slice := reflect.New(typ).Elem() + slice.Set(reflect.MakeSlice(typ, 0, 0)) + return slice + case reflect.Ptr: + typ = typ.Elem() + origin := reflect.New(typ) + nested := newInstanceByType(typ) + origin.Elem().Set(nested) + + return origin + default: + return reflect.New(typ).Elem() + } +} + +func checkAndExtractFromField(fromField string, input reflect.Value) (reflect.Value, error) { + f := input.FieldByName(fromField) + if !f.IsValid() { + return reflect.Value{}, fmt.Errorf("field mapping from a struct field, but field not found. field=%v, inputType=%v", fromField, input.Type()) + } + + if !f.CanInterface() { + return reflect.Value{}, fmt.Errorf("field mapping from a struct field, but field not exported. field= %v, inputType=%v", fromField, input.Type()) + } + + return f, nil +} + +type errMapKeyNotFound struct { + mapKey string +} + +func (e *errMapKeyNotFound) Error() string { + return fmt.Sprintf("key=%s", e.mapKey) +} + +type errInterfaceNotValidForFieldMapping struct { + interfaceType reflect.Type + actualType reflect.Type +} + +func (e *errInterfaceNotValidForFieldMapping) Error() string { + return fmt.Sprintf("field mapping from an interface type, but actual type is not struct, struct ptr or map. InterfaceType= %v, ActualType= %v", e.interfaceType, e.actualType) +} + +func checkAndExtractFromMapKey(fromMapKey string, input reflect.Value) (reflect.Value, error) { + key := reflect.ValueOf(fromMapKey) + if input.Type().Key() != strType { + key = key.Convert(input.Type().Key()) + } + + v := input.MapIndex(key) + if !v.IsValid() { + return reflect.Value{}, fmt.Errorf("field mapping from a map key, but key not found in input. %w", &errMapKeyNotFound{mapKey: fromMapKey}) + } + + return v, nil +} + +func checkAndExtractFieldType(paths []string, typ reflect.Type) (extracted reflect.Type, remainingPaths FieldPath, err error) { + extracted = typ + for i, field := range paths { + for extracted.Kind() == reflect.Ptr { + extracted = extracted.Elem() + } + + if extracted.Kind() == reflect.Map { + if !strType.ConvertibleTo(extracted.Key()) { + return nil, nil, fmt.Errorf("type[%v] is not a map with string or string alias key", extracted) + } + + extracted = extracted.Elem() + continue + } + + if extracted.Kind() == reflect.Struct { + f, ok := extracted.FieldByName(field) + if !ok { + return nil, nil, fmt.Errorf("type[%v] has no field[%s]", extracted, field) + } + + if !f.IsExported() { + return nil, nil, fmt.Errorf("type[%v] has an unexported field[%s]", extracted.String(), field) + } + + extracted = f.Type + continue + } + + if extracted.Kind() == reflect.Interface { + return extracted, paths[i:], nil + } + + return nil, nil, fmt.Errorf("intermediate type[%v] is not valid", extracted) + } + + return extracted, nil, nil +} + +var strType = reflect.TypeOf("") + +func fieldMap(mappings []*FieldMapping, allowMapKeyNotFound bool, uncheckedSourcePaths map[string]FieldPath) func(any) (map[string]any, error) { + return func(input any) (result map[string]any, err error) { + result = make(map[string]any, len(mappings)) + var inputValue reflect.Value + loop: + for _, mapping := range mappings { + if mapping.customExtractor != nil { + result[mapping.to], err = mapping.customExtractor(input) + if err != nil { + return nil, err + } + continue + } + + if len(mapping.from) == 0 { + result[mapping.to] = input + continue + } + + fromPath := splitFieldPath(mapping.from) + + if !inputValue.IsValid() { + inputValue = reflect.ValueOf(input) + } + + var ( + pathInputValue = inputValue + pathInputType = inputValue.Type() + taken = input + ) + + for i, path := range fromPath { + for pathInputValue.Kind() == reflect.Ptr { + pathInputValue = pathInputValue.Elem() + } + + if !pathInputValue.IsValid() { + return nil, fmt.Errorf("intermediate source value on path=%v is nil for type [%v]", fromPath[:i+1], pathInputType) + } + + if pathInputValue.Kind() == reflect.Map && pathInputValue.IsNil() { + return nil, fmt.Errorf("intermediate source value on path=%v is nil for map type [%v]", fromPath[:i+1], pathInputType) + } + + taken, pathInputType, err = takeOne(pathInputValue, pathInputType, path) + if err != nil { + // we deferred check from Compile time to request time for interface types, so we won't panic here + var interfaceNotValidErr *errInterfaceNotValidForFieldMapping + if errors.As(err, &interfaceNotValidErr) { + return nil, err + } + + // map key not found can only be a request time error, so we won't panic here + var mapKeyNotFoundErr *errMapKeyNotFound + if errors.As(err, &mapKeyNotFoundErr) { + if allowMapKeyNotFound { + continue loop + } + return nil, err + } + + if uncheckedSourcePaths != nil { + uncheckedPath, ok := uncheckedSourcePaths[mapping.from] + if ok && len(uncheckedPath) >= len(fromPath)-i { + // the err happens on the mapping source path which is unchecked at request time, so we won't panic here + return nil, err + } + } + + panic(safe.NewPanicErr(err, debug.Stack())) + } + + if i < len(fromPath)-1 { + pathInputValue = reflect.ValueOf(taken) + } + } + + result[mapping.to] = taken + } + + return result, nil + } +} + +func streamFieldMap(mappings []*FieldMapping, uncheckedSourcePaths map[string]FieldPath) func(streamReader) streamReader { + return func(input streamReader) streamReader { + return packStreamReader(schema.StreamReaderWithConvert(input.toAnyStreamReader(), fieldMap(mappings, true, uncheckedSourcePaths))) + } +} + +func takeOne(inputValue reflect.Value, inputType reflect.Type, from string) (taken any, takenType reflect.Type, err error) { + var f reflect.Value + switch k := inputValue.Kind(); k { + case reflect.Map: + f, err = checkAndExtractFromMapKey(from, inputValue) + if err != nil { + return nil, nil, err + } + + return f.Interface(), f.Type(), nil + case reflect.Struct: + f, err = checkAndExtractFromField(from, inputValue) + if err != nil { + return nil, nil, err + } + + return f.Interface(), f.Type(), nil + default: + if inputType.Kind() == reflect.Interface { + return nil, nil, &errInterfaceNotValidForFieldMapping{ + interfaceType: inputType, + actualType: inputValue.Type(), + } + } + + panic("when take one value from source, value not map or struct, and type not interface") + } +} + +func isFromAll(mappings []*FieldMapping) bool { + for _, mapping := range mappings { + if len(mapping.from) == 0 && mapping.customExtractor == nil { + return true + } + } + return false +} + +func fromFields(mappings []*FieldMapping) bool { + for _, mapping := range mappings { + if len(mapping.from) == 0 || mapping.customExtractor != nil { + return false + } + } + + return true +} + +func isToAll(mappings []*FieldMapping) bool { + for _, mapping := range mappings { + if len(mapping.to) == 0 { + return true + } + } + return false +} + +func validateStructOrMap(t reflect.Type) bool { + switch t.Kind() { + case reflect.Map: + return true + case reflect.Ptr: + t = t.Elem() + fallthrough + case reflect.Struct: + return true + default: + return false + } +} + +func validateFieldMapping(predecessorType reflect.Type, successorType reflect.Type, mappings []*FieldMapping) ( + // type checkers that are deferred to request-time + typeHandler *handlerPair, + // the remaining predecessor field paths that are not checked at compile time because of interface type found + uncheckedSourcePath map[string]FieldPath, + err error) { + // check if mapping is legal + if isFromAll(mappings) && isToAll(mappings) { + // unreachable + panic(fmt.Errorf("invalid field mappings: from all fields to all, use common edge instead")) + } else if !isToAll(mappings) && (!validateStructOrMap(successorType) && successorType != reflect.TypeOf((*any)(nil)).Elem()) { + // if user has not provided a specific struct type, graph cannot construct any struct in the runtime + return nil, nil, fmt.Errorf("static check fail: successor input type should be struct or map, actual: %v", successorType) + } else if fromFields(mappings) && !validateStructOrMap(predecessorType) { + return nil, nil, fmt.Errorf("static check fail: predecessor output type should be struct or map, actual: %v", predecessorType) + } + + var fieldCheckers map[string]handlerPair + + for i := range mappings { + mapping := mappings[i] + + successorFieldType, successorRemaining, err := checkAndExtractFieldType(splitFieldPath(mapping.to), successorType) + if err != nil { + return nil, nil, fmt.Errorf("static check failed for mapping %s: %w", mapping, err) + } + + if len(successorRemaining) > 0 { + if successorFieldType == reflect.TypeOf((*any)(nil)).Elem() { + continue // at request time expand this 'any' to 'map[string]any' + } + return nil, nil, fmt.Errorf("static check failed for mapping %s, the successor has intermediate interface type %v", mapping, successorFieldType) + } + + if mapping.customExtractor != nil { // custom extractor applies to request-time data, so skip compile-time check + continue + } + + predecessorFieldType, predecessorRemaining, err := checkAndExtractFieldType(splitFieldPath(mapping.from), predecessorType) + if err != nil { + return nil, nil, fmt.Errorf("static check failed for mapping %s: %w", mapping, err) + } + + if len(predecessorRemaining) > 0 { + if uncheckedSourcePath == nil { + uncheckedSourcePath = make(map[string]FieldPath) + } + uncheckedSourcePath[mapping.from] = predecessorRemaining + } + + checker := func(a any) (any, error) { + trueInType := reflect.TypeOf(a) + if trueInType == nil { + switch successorFieldType.Kind() { + case reflect.Map, reflect.Slice, reflect.Ptr, reflect.Interface: + default: + return nil, fmt.Errorf("runtime check failed for mapping %s, field[%v]-[%v] is absolutely not assignable", mapping, trueInType, successorFieldType) + } + } else { + if !trueInType.AssignableTo(successorFieldType) { + return nil, fmt.Errorf("runtime check failed for mapping %s, field[%v]-[%v] is absolutely not assignable", mapping, trueInType, successorFieldType) + } + } + + return a, nil + } + + if len(predecessorRemaining) > 0 { + // can't check if types match at compile time, because there is interface type at some point along the source path. Defer to request time + if fieldCheckers == nil { + fieldCheckers = make(map[string]handlerPair) + } + fieldCheckers[mapping.to] = handlerPair{ + invoke: checker, + transform: func(input streamReader) streamReader { + return packStreamReader(schema.StreamReaderWithConvert(input.toAnyStreamReader(), checker)) + }, + } + } else { + at := checkAssignable(predecessorFieldType, successorFieldType) + if at == assignableTypeMustNot { + return nil, nil, fmt.Errorf("static check failed for mapping %s, field[%v]-[%v] is absolutely not assignable", mapping, predecessorFieldType, successorFieldType) + } else if at == assignableTypeMay { + // can't decide if types match, because the successorFieldType implements predecessorFieldType, which is an interface type + if fieldCheckers == nil { + fieldCheckers = make(map[string]handlerPair) + } + fieldCheckers[mapping.to] = handlerPair{ + invoke: checker, + transform: func(input streamReader) streamReader { + return packStreamReader(schema.StreamReaderWithConvert(input.toAnyStreamReader(), checker)) + }, + } + } + } + + } + + if len(fieldCheckers) == 0 { + return nil, uncheckedSourcePath, nil + } + + checker := func(value map[string]any) (map[string]any, error) { + var err error + for k, v := range fieldCheckers { + for mapping := range value { + if mapping == k { + value[mapping], err = v.invoke(value[mapping]) + if err != nil { + return nil, err + } + } + } + } + return value, nil + } + return &handlerPair{ + invoke: func(value any) (any, error) { + return checker(value.(map[string]any)) + }, + transform: func(input streamReader) streamReader { + s, ok := unpackStreamReader[map[string]any](input) + if !ok { + // impossible + panic("field mapping edge stream value isn't map[string]any") + } + return packStreamReader(schema.StreamReaderWithConvert(s, checker)) + }, + }, uncheckedSourcePath, nil +} diff --git a/compose/generic_graph.go b/compose/generic_graph.go new file mode 100644 index 0000000..c763811 --- /dev/null +++ b/compose/generic_graph.go @@ -0,0 +1,158 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "reflect" + + "github.com/cloudwego/eino/internal/generic" +) + +type newGraphOptions struct { + withState func(ctx context.Context) any + stateType reflect.Type +} + +// NewGraphOption configures behavior when creating a new graph, such as +// providing local state generation. +type NewGraphOption func(ngo *newGraphOptions) + +// WithGenLocalState registers a function to generate per-run local state +// that can be shared across nodes in the graph. +func WithGenLocalState[S any](gls GenLocalState[S]) NewGraphOption { + return func(ngo *newGraphOptions) { + ngo.withState = func(ctx context.Context) any { + return gls(ctx) + } + ngo.stateType = generic.TypeOf[S]() + } +} + +// NewGraph create a directed graph that can compose components, lambda, chain, parallel etc. +// simultaneously provide flexible and multi-granular aspect governance capabilities. +// I: the input type of graph compiled product +// O: the output type of graph compiled product +// +// To share state between nodes, use WithGenLocalState option: +// +// type testState struct { +// UserInfo *UserInfo +// KVs map[string]any +// } +// +// genStateFunc := func(ctx context.Context) *testState { +// return &testState{} +// } +// +// graph := compose.NewGraph[string, string](WithGenLocalState(genStateFunc)) +// +// // you can use WithStatePreHandler and WithStatePostHandler to do something with state +// graph.AddNode("node1", someNode, compose.WithPreHandler(func(ctx context.Context, in string, state *testState) (string, error) { +// // do something with state +// return in, nil +// }), compose.WithPostHandler(func(ctx context.Context, out string, state *testState) (string, error) { +// // do something with state +// return out, nil +// })) +func NewGraph[I, O any](opts ...NewGraphOption) *Graph[I, O] { + options := &newGraphOptions{} + for _, opt := range opts { + opt(options) + } + + g := &Graph[I, O]{ + newGraphFromGeneric[I, O]( + ComponentOfGraph, + options.withState, + options.stateType, + opts, + ), + } + + return g +} + +// Graph is a generic graph that can be used to compose components. +// I: the input type of graph compiled product +// O: the output type of graph compiled product +type Graph[I, O any] struct { + *graph +} + +// AddEdge adds an edge to the graph, edge means a data flow from startNode to endNode. +// the previous node's output type must be set to the next node's input type. +// NOTE: startNode and endNode must have been added to the graph before adding edge. +// e.g. +// +// graph.AddNode("start_node_key", compose.NewPassthroughNode()) +// graph.AddNode("end_node_key", compose.NewPassthroughNode()) +// +// err := graph.AddEdge("start_node_key", "end_node_key") +func (g *Graph[I, O]) AddEdge(startNode, endNode string) (err error) { + return g.graph.addEdgeWithMappings(startNode, endNode, false, false) +} + +// Compile take the raw graph and compile it into a form ready to be run. +// e.g. +// +// graph, err := compose.NewGraph[string, string]() +// if err != nil {...} +// +// runnable, err := graph.Compile(ctx, compose.WithGraphName("my_graph")) +// if err != nil {...} +// +// runnable.Invoke(ctx, "input") // invoke +// runnable.Stream(ctx, "input") // stream +// runnable.Collect(ctx, inputReader) // collect +// runnable.Transform(ctx, inputReader) // transform +func (g *Graph[I, O]) Compile(ctx context.Context, opts ...GraphCompileOption) (Runnable[I, O], error) { + return compileAnyGraph[I, O](ctx, g, opts...) +} + +func compileAnyGraph[I, O any](ctx context.Context, g AnyGraph, opts ...GraphCompileOption) (Runnable[I, O], error) { + if len(globalGraphCompileCallbacks) > 0 { + opts = append([]GraphCompileOption{WithGraphCompileCallbacks(globalGraphCompileCallbacks...)}, opts...) + } + option := newGraphCompileOptions(opts...) + + cr, err := g.compile(ctx, option) + if err != nil { + return nil, err + } + + cr.meta = &executorMeta{ + component: g.component(), + isComponentCallbackEnabled: true, + componentImplType: "", + } + + cr.nodeInfo = &nodeInfo{ + name: option.graphName, + } + + ctxWrapper := func(ctx context.Context, opts ...Option) context.Context { + return initGraphCallbacks(AppendAddressSegment(ctx, AddressSegmentRunnable, option.graphName), cr.nodeInfo, cr.meta, opts...) + } + + rp, err := toGenericRunnable[I, O](cr, ctxWrapper) + if err != nil { + return nil, err + } + + return rp, nil +} diff --git a/compose/generic_helper.go b/compose/generic_helper.go new file mode 100644 index 0000000..0534dc6 --- /dev/null +++ b/compose/generic_helper.go @@ -0,0 +1,256 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "errors" + "fmt" + "reflect" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +func newGenericHelper[I, O any]() *genericHelper { + return &genericHelper{ + inputStreamFilter: defaultStreamMapFilter[I], + outputStreamFilter: defaultStreamMapFilter[O], + inputConverter: handlerPair{ + invoke: defaultValueChecker[I], + transform: defaultStreamConverter[I], + }, + outputConverter: handlerPair{ + invoke: defaultValueChecker[O], + transform: defaultStreamConverter[O], + }, + inputFieldMappingConverter: handlerPair{ + invoke: buildFieldMappingConverter[I](), + transform: buildStreamFieldMappingConverter[I](), + }, + outputFieldMappingConverter: handlerPair{ + invoke: buildFieldMappingConverter[O](), + transform: buildStreamFieldMappingConverter[O](), + }, + inputStreamConvertPair: defaultStreamConvertPair[I](), + outputStreamConvertPair: defaultStreamConvertPair[O](), + inputZeroValue: zeroValueFromGeneric[I], + outputZeroValue: zeroValueFromGeneric[O], + inputEmptyStream: emptyStreamFromGeneric[I], + outputEmptyStream: emptyStreamFromGeneric[O], + } +} + +type genericHelper struct { + // when set input key, use this method to convert input from map[string]any to T + inputStreamFilter, outputStreamFilter streamMapFilter + // when predecessor's output is assignableTypeMay to current node's input, validate and convert(if needed) types using the following two methods + inputConverter, outputConverter handlerPair + // when current node enable field mapping, convert map input to expected struct using the following two methods + inputFieldMappingConverter, outputFieldMappingConverter handlerPair + // can convert input/output from stream to non-stream or non-stream to stream, used for checkpoint + inputStreamConvertPair, outputStreamConvertPair streamConvertPair + + inputZeroValue, outputZeroValue func() any + inputEmptyStream, outputEmptyStream func() streamReader +} + +func (g *genericHelper) forMapInput() *genericHelper { + return &genericHelper{ + outputStreamFilter: g.outputStreamFilter, + outputConverter: g.outputConverter, + outputFieldMappingConverter: g.outputFieldMappingConverter, + outputStreamConvertPair: g.outputStreamConvertPair, + outputZeroValue: g.outputZeroValue, + outputEmptyStream: g.outputEmptyStream, + + inputStreamFilter: defaultStreamMapFilter[map[string]any], + inputConverter: handlerPair{ + invoke: defaultValueChecker[map[string]any], + transform: defaultStreamConverter[map[string]any], + }, + inputFieldMappingConverter: handlerPair{ + invoke: buildFieldMappingConverter[map[string]any](), + transform: buildStreamFieldMappingConverter[map[string]any](), + }, + inputStreamConvertPair: defaultStreamConvertPair[map[string]any](), + inputZeroValue: zeroValueFromGeneric[map[string]any], + inputEmptyStream: emptyStreamFromGeneric[map[string]any], + } +} + +func (g *genericHelper) forMapOutput() *genericHelper { + return &genericHelper{ + inputStreamFilter: g.inputStreamFilter, + inputConverter: g.inputConverter, + inputFieldMappingConverter: g.inputFieldMappingConverter, + inputStreamConvertPair: g.inputStreamConvertPair, + inputZeroValue: g.inputZeroValue, + inputEmptyStream: g.inputEmptyStream, + + outputStreamFilter: defaultStreamMapFilter[map[string]any], + outputConverter: handlerPair{ + invoke: defaultValueChecker[map[string]any], + transform: defaultStreamConverter[map[string]any], + }, + outputFieldMappingConverter: handlerPair{ + invoke: buildFieldMappingConverter[map[string]any](), + transform: buildStreamFieldMappingConverter[map[string]any](), + }, + outputStreamConvertPair: defaultStreamConvertPair[map[string]any](), + outputZeroValue: zeroValueFromGeneric[map[string]any], + outputEmptyStream: emptyStreamFromGeneric[map[string]any], + } +} + +func (g *genericHelper) forPredecessorPassthrough() *genericHelper { + return &genericHelper{ + inputStreamFilter: g.inputStreamFilter, + outputStreamFilter: g.inputStreamFilter, + inputConverter: g.inputConverter, + outputConverter: g.inputConverter, + inputFieldMappingConverter: g.inputFieldMappingConverter, + outputFieldMappingConverter: g.inputFieldMappingConverter, + inputStreamConvertPair: g.inputStreamConvertPair, + outputStreamConvertPair: g.inputStreamConvertPair, + inputZeroValue: g.inputZeroValue, + outputZeroValue: g.inputZeroValue, + inputEmptyStream: g.inputEmptyStream, + outputEmptyStream: g.inputEmptyStream, + } +} + +func (g *genericHelper) forSuccessorPassthrough() *genericHelper { + return &genericHelper{ + inputStreamFilter: g.outputStreamFilter, + outputStreamFilter: g.outputStreamFilter, + inputConverter: g.outputConverter, + outputConverter: g.outputConverter, + inputFieldMappingConverter: g.outputFieldMappingConverter, + outputFieldMappingConverter: g.outputFieldMappingConverter, + inputStreamConvertPair: g.outputStreamConvertPair, + outputStreamConvertPair: g.outputStreamConvertPair, + inputZeroValue: g.outputZeroValue, + outputZeroValue: g.outputZeroValue, + inputEmptyStream: g.outputEmptyStream, + outputEmptyStream: g.outputEmptyStream, + } +} + +type streamMapFilter func(key string, isr streamReader) (streamReader, bool) + +type valueHandler func(value any) (any, error) +type streamHandler func(streamReader) streamReader + +type handlerPair struct { + invoke valueHandler + transform streamHandler +} + +type streamConvertPair struct { + concatStream func(sr streamReader) (any, error) + restoreStream func(any) (streamReader, error) +} + +func defaultStreamConvertPair[T any]() streamConvertPair { + var t T + return streamConvertPair{ + concatStream: func(sr streamReader) (any, error) { + tsr, ok := unpackStreamReader[T](sr) + if !ok { + return nil, fmt.Errorf("cannot convert sr to streamReader[%T]", t) + } + value, err := concatStreamReader(tsr) + if err != nil { + if errors.Is(err, emptyStreamConcatErr) { + return nil, nil + } + return nil, err + } + return value, nil + }, + restoreStream: func(a any) (streamReader, error) { + if a == nil { + return packStreamReader(schema.StreamReaderFromArray([]T{})), nil + } + value, ok := a.(T) + if !ok { + return nil, fmt.Errorf("cannot convert value[%T] to streamReader[%T]", a, t) + } + return packStreamReader(schema.StreamReaderFromArray([]T{value})), nil + }, + } +} + +func defaultStreamMapFilter[T any](key string, isr streamReader) (streamReader, bool) { + sr, ok := unpackStreamReader[map[string]any](isr) + if !ok { + return nil, false + } + + cvt := func(m map[string]any) (T, error) { + var t T + v, ok_ := m[key] + if !ok_ { + return t, schema.ErrNoValue + } + vv, ok_ := v.(T) + if !ok_ { + return t, fmt.Errorf( + "[defaultStreamMapFilter]fail, key[%s]'s value type[%s] isn't expected type[%s]", + key, reflect.TypeOf(v).String(), + generic.TypeOf[T]().String()) + } + return vv, nil + } + + ret := schema.StreamReaderWithConvert[map[string]any, T](sr, cvt) + + return packStreamReader(ret), true +} + +func defaultStreamConverter[T any](reader streamReader) streamReader { + return packStreamReader(schema.StreamReaderWithConvert(reader.toAnyStreamReader(), func(v any) (T, error) { + vv, ok := v.(T) + if !ok { + var t T + return t, fmt.Errorf("runtime type check fail, expected type: %T, actual type: %T", t, v) + } + return vv, nil + })) +} + +func defaultValueChecker[T any](v any) (any, error) { + nValue, ok := v.(T) + if !ok { + var t T + return nil, fmt.Errorf("runtime type check fail, expected type: %T, actual type: %T", t, v) + } + return nValue, nil +} + +func zeroValueFromGeneric[T any]() any { + var t T + return t +} + +func emptyStreamFromGeneric[T any]() streamReader { + var t T + sr, sw := schema.Pipe[T](1) + sw.Send(t, nil) + sw.Close() + return packStreamReader(sr) +} diff --git a/compose/graph.go b/compose/graph.go new file mode 100644 index 0000000..bcf5ae4 --- /dev/null +++ b/compose/graph.go @@ -0,0 +1,1219 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + "reflect" + "strings" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/internal/gmap" +) + +// START is the start node of the graph. You can add your first edge with START. +const START = "start" + +// END is the end node of the graph. You can add your last edge with END. +const END = "end" + +// graphRunType is a custom type used to control the running mode of the graph. +type graphRunType string + +const ( + // runTypePregel is a running mode of the graph that is suitable for large-scale graph processing tasks. Can have cycles in graph. Compatible with NodeTriggerType.AnyPredecessor. + runTypePregel graphRunType = "Pregel" + // runTypeDAG is a running mode of the graph that represents the graph as a directed acyclic graph, suitable for tasks that can be represented as a directed acyclic graph. Compatible with NodeTriggerType.AllPredecessor. + runTypeDAG graphRunType = "DAG" +) + +// String returns the string representation of the graph run type. +func (g graphRunType) String() string { + return string(g) +} + +type graph struct { + nodes map[string]*graphNode + controlEdges map[string][]string + dataEdges map[string][]string + branches map[string][]*GraphBranch + startNodes []string + endNodes []string + + toValidateMap map[string][]struct { + endNode string + mappings []*FieldMapping + } + + stateType reflect.Type + stateGenerator func(ctx context.Context) any + newOpts []NewGraphOption + + expectedInputType, expectedOutputType reflect.Type + + *genericHelper + + fieldMappingRecords map[string][]*FieldMapping + + buildError error + + cmp component + + compiled bool + + handlerOnEdges map[string]map[string][]handlerPair + handlerPreNode map[string][]handlerPair + handlerPreBranch map[string][][]handlerPair +} + +type newGraphConfig struct { + inputType, outputType reflect.Type + gh *genericHelper + cmp component + stateType reflect.Type + stateGenerator func(ctx context.Context) any + newOpts []NewGraphOption +} + +func newGraphFromGeneric[I, O any]( + cmp component, + stateGenerator func(ctx context.Context) any, + stateType reflect.Type, + opts []NewGraphOption, +) *graph { + return newGraph(&newGraphConfig{ + inputType: generic.TypeOf[I](), + outputType: generic.TypeOf[O](), + gh: newGenericHelper[I, O](), + cmp: cmp, + stateType: stateType, + stateGenerator: stateGenerator, + newOpts: opts, + }) +} + +func newGraph(cfg *newGraphConfig) *graph { + return &graph{ + nodes: make(map[string]*graphNode), + dataEdges: make(map[string][]string), + controlEdges: make(map[string][]string), + branches: make(map[string][]*GraphBranch), + + toValidateMap: make(map[string][]struct { + endNode string + mappings []*FieldMapping + }), + + expectedInputType: cfg.inputType, + expectedOutputType: cfg.outputType, + genericHelper: cfg.gh, + + fieldMappingRecords: make(map[string][]*FieldMapping), + + cmp: cfg.cmp, + + stateType: cfg.stateType, + stateGenerator: cfg.stateGenerator, + newOpts: cfg.newOpts, + + handlerOnEdges: make(map[string]map[string][]handlerPair), + handlerPreNode: make(map[string][]handlerPair), + handlerPreBranch: make(map[string][][]handlerPair), + } +} + +func (g *graph) component() component { + return g.cmp +} + +func isChain(cmp component) bool { + return cmp == ComponentOfChain +} + +func isWorkflow(cmp component) bool { + return cmp == ComponentOfWorkflow +} + +// ErrGraphCompiled is returned when attempting to modify a graph after it has been compiled +var ErrGraphCompiled = errors.New("graph has been compiled, cannot be modified") + +func (g *graph) addNode(key string, node *graphNode, options *graphAddNodeOpts) (err error) { + if g.buildError != nil { + return g.buildError + } + + if g.compiled { + return ErrGraphCompiled + } + + defer func() { + if err != nil { + g.buildError = err + } + }() + + if key == END || key == START { + return fmt.Errorf("node '%s' is reserved, cannot add manually", key) + } + + if _, ok := g.nodes[key]; ok { + return fmt.Errorf("node '%s' already present", key) + } + + // check options + if options.needState { + if g.stateGenerator == nil { + return fmt.Errorf("node '%s' needs state but graph state is not enabled", key) + } + } + + if options.nodeOptions.nodeKey != "" { + if !isChain(g.cmp) { + return errors.New("only chain support node key option") + } + } + // end: check options + + // check pre- / post-handler type + if options.processor != nil { + if options.processor.statePreHandler != nil { + // check state type + if g.stateType != options.processor.preStateType { + return fmt.Errorf("node[%s]'s pre handler state type[%v] is different from graph[%v]", key, options.processor.preStateType, g.stateType) + } + // check input type + if node.inputType() == nil && options.processor.statePreHandler.outputType != reflect.TypeOf((*any)(nil)).Elem() { + return fmt.Errorf("passthrough node[%s]'s pre handler type isn't any", key) + } else if node.inputType() != nil && node.inputType() != options.processor.statePreHandler.outputType { + return fmt.Errorf("node[%s]'s pre handler type[%v] is different from its input type[%v]", key, options.processor.statePreHandler.outputType, node.inputType()) + } + } + if options.processor.statePostHandler != nil { + // check state type + if g.stateType != options.processor.postStateType { + return fmt.Errorf("node[%s]'s post handler state type[%v] is different from graph[%v]", key, options.processor.postStateType, g.stateType) + } + // check input type + if node.outputType() == nil && options.processor.statePostHandler.inputType != reflect.TypeOf((*any)(nil)).Elem() { + return fmt.Errorf("passthrough node[%s]'s post handler type isn't any", key) + } else if node.outputType() != nil && node.outputType() != options.processor.statePostHandler.inputType { + return fmt.Errorf("node[%s]'s post handler type[%v] is different from its output type[%v]", key, options.processor.statePostHandler.inputType, node.outputType()) + } + } + } + + g.nodes[key] = node + + return nil +} + +func (g *graph) addEdgeWithMappings(startNode, endNode string, noControl bool, noData bool, mappings ...*FieldMapping) (err error) { + if g.buildError != nil { + return g.buildError + } + if g.compiled { + return ErrGraphCompiled + } + + if noControl && noData { + return fmt.Errorf("edge[%s]-[%s] cannot be both noDirectDependency and noDataFlow", startNode, endNode) + } + + defer func() { + if err != nil { + g.buildError = err + } + }() + if startNode == END { + return errors.New("END cannot be a start node") + } + if endNode == START { + return errors.New("START cannot be an end node") + } + + if _, ok := g.nodes[startNode]; !ok && startNode != START { + return fmt.Errorf("edge start node '%s' needs to be added to graph first", startNode) + } + if _, ok := g.nodes[endNode]; !ok && endNode != END { + return fmt.Errorf("edge end node '%s' needs to be added to graph first", endNode) + } + + if !noControl { + for i := range g.controlEdges[startNode] { + if g.controlEdges[startNode][i] == endNode { + return fmt.Errorf("control edge[%s]-[%s] have been added yet", startNode, endNode) + } + } + + g.controlEdges[startNode] = append(g.controlEdges[startNode], endNode) + if startNode == START { + g.startNodes = append(g.startNodes, endNode) + } + if endNode == END { + g.endNodes = append(g.endNodes, startNode) + } + } + if !noData { + for i := range g.dataEdges[startNode] { + if g.dataEdges[startNode][i] == endNode { + return fmt.Errorf("data edge[%s]-[%s] have been added yet", startNode, endNode) + } + } + + g.addToValidateMap(startNode, endNode, mappings) + err = g.updateToValidateMap() + if err != nil { + return err + } + g.dataEdges[startNode] = append(g.dataEdges[startNode], endNode) + } + + return nil +} + +// AddEmbeddingNode adds a node that implements embedding.Embedder. +// e.g. +// +// embeddingNode, err := openai.NewEmbedder(ctx, &openai.EmbeddingConfig{ +// Model: "text-embedding-3-small", +// }) +// +// graph.AddEmbeddingNode("embedding_node_key", embeddingNode) +func (g *graph) AddEmbeddingNode(key string, node embedding.Embedder, opts ...GraphAddNodeOpt) error { + gNode, options := toEmbeddingNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddRetrieverNode adds a node that implements retriever.Retriever. +// e.g. +// +// retriever, err := vikingdb.NewRetriever(ctx, &vikingdb.RetrieverConfig{}) +// +// graph.AddRetrieverNode("retriever_node_key", retrieverNode) +func (g *graph) AddRetrieverNode(key string, node retriever.Retriever, opts ...GraphAddNodeOpt) error { + gNode, options := toRetrieverNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddLoaderNode adds a node that implements document.Loader. +// e.g. +// +// loader, err := file.NewLoader(ctx, &file.LoaderConfig{}) +// +// graph.AddLoaderNode("loader_node_key", loader) +func (g *graph) AddLoaderNode(key string, node document.Loader, opts ...GraphAddNodeOpt) error { + gNode, options := toLoaderNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddIndexerNode adds a node that implements indexer.Indexer. +// e.g. +// +// indexer, err := vikingdb.NewIndexer(ctx, &vikingdb.IndexerConfig{}) +// +// graph.AddIndexerNode("indexer_node_key", indexer) +func (g *graph) AddIndexerNode(key string, node indexer.Indexer, opts ...GraphAddNodeOpt) error { + gNode, options := toIndexerNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddChatModelNode add node that implements model.BaseChatModel. +// e.g. +// +// chatModel, err := openai.NewChatModel(ctx, &openai.ChatModelConfig{ +// Model: "gpt-4o", +// }) +// +// graph.AddChatModelNode("chat_model_node_key", chatModel) +func (g *graph) AddChatModelNode(key string, node model.BaseChatModel, opts ...GraphAddNodeOpt) error { + gNode, options := toChatModelNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddAgenticModelNode add node that implements agentic.Model. +// e.g. +// +// model, err := openai.NewAgenticModel(ctx, &openai.AgenticModelConfig{ +// Model: "gpt-4o", +// }) +// +// graph.AddAgenticModelNode("agentic_model_node_key", model) +func (g *graph) AddAgenticModelNode(key string, node model.AgenticModel, opts ...GraphAddNodeOpt) error { + gNode, options := toAgenticModelNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddChatTemplateNode add node that implements prompt.ChatTemplate. +// e.g. +// +// chatTemplate, err := prompt.FromMessages(schema.FString, &schema.Message{ +// Role: schema.System, +// Content: "You are acting as a {role}.", +// }) +// +// graph.AddChatTemplateNode("chat_template_node_key", chatTemplate) +func (g *graph) AddChatTemplateNode(key string, node prompt.ChatTemplate, opts ...GraphAddNodeOpt) error { + gNode, options := toChatTemplateNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddAgenticChatTemplateNode add node that implements prompt.AgenticChatTemplate. +// e.g. +// +// chatTemplate, err := prompt.FromAgenticMessages(schema.FString, &schema.AgenticMessage{}) +// +// graph.AddAgenticChatTemplateNode("chat_template_node_key", chatTemplate) +func (g *graph) AddAgenticChatTemplateNode(key string, node prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) error { + gNode, options := toAgenticChatTemplateNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddToolsNode adds a node that implements ToolsNode. +// e.g. +// +// toolsNode, err := compose.NewToolNode(ctx, &compose.ToolsNodeConfig{}) +// +// graph.AddToolsNode("tools_node_key", toolsNode) +func (g *graph) AddToolsNode(key string, node *ToolsNode, opts ...GraphAddNodeOpt) error { + gNode, options := toToolsNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddAgenticToolsNode adds a node that implements AgenticToolsNode. +// e.g. +// +// toolsNode, err := compose.NewAgenticToolsNode(ctx, &compose.ToolsNodeConfig{}) +// +// graph.AddAgenticToolsNode("tools_node_key", toolsNode) +func (g *graph) AddAgenticToolsNode(key string, node *AgenticToolsNode, opts ...GraphAddNodeOpt) error { + gNode, options := toAgenticToolsNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddDocumentTransformerNode adds a node that implements document.Transformer. +// e.g. +// +// markdownSplitter, err := markdown.NewHeaderSplitter(ctx, &markdown.HeaderSplitterConfig{}) +// +// graph.AddDocumentTransformerNode("document_transformer_node_key", markdownSplitter) +func (g *graph) AddDocumentTransformerNode(key string, node document.Transformer, opts ...GraphAddNodeOpt) error { + gNode, options := toDocumentTransformerNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddLambdaNode add node that implements at least one of Invoke[I, O], Stream[I, O], Collect[I, O], Transform[I, O]. +// due to the lack of supporting method generics, we need to use function generics to generate Lambda run as Runnable[I, O]. +// for Invoke[I, O], use compose.InvokableLambda() +// for Stream[I, O], use compose.StreamableLambda() +// for Collect[I, O], use compose.CollectableLambda() +// for Transform[I, O], use compose.TransformableLambda() +// for arbitrary combinations of 4 kinds of lambda, use compose.AnyLambda() +func (g *graph) AddLambdaNode(key string, node *Lambda, opts ...GraphAddNodeOpt) error { + gNode, options := toLambdaNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddGraphNode add one kind of Graph[I, O]、Chain[I, O]、StateChain[I, O, S] as a node. +// for Graph[I, O], comes from NewGraph[I, O]() +// for Chain[I, O], comes from NewChain[I, O]() +func (g *graph) AddGraphNode(key string, node AnyGraph, opts ...GraphAddNodeOpt) error { + gNode, options := toAnyGraphNode(node, opts...) + return g.addNode(key, gNode, options) +} + +// AddPassthroughNode adds a passthrough node to the graph. +// mostly used in pregel mode of graph. +// e.g. +// +// graph.AddPassthroughNode("passthrough_node_key") +func (g *graph) AddPassthroughNode(key string, opts ...GraphAddNodeOpt) error { + gNode, options := toPassthroughNode(opts...) + return g.addNode(key, gNode, options) +} + +// AddBranch adds a branch to the graph. +// e.g. +// +// condition := func(ctx context.Context, in string) (string, error) { +// return "next_node_key", nil +// } +// endNodes := map[string]bool{"path01": true, "path02": true} +// branch := compose.NewGraphBranch(condition, endNodes) +// +// graph.AddBranch("start_node_key", branch) +func (g *graph) AddBranch(startNode string, branch *GraphBranch) (err error) { + return g.addBranch(startNode, branch, false) +} + +func (g *graph) addBranch(startNode string, branch *GraphBranch, skipData bool) (err error) { + if g.buildError != nil { + return g.buildError + } + + if g.compiled { + return ErrGraphCompiled + } + + defer func() { + if err != nil { + g.buildError = err + } + }() + + if startNode == END { + return errors.New("END cannot be a start node") + } + + if _, ok := g.nodes[startNode]; !ok && startNode != START { + return fmt.Errorf("branch start node '%s' needs to be added to graph first", startNode) + } + + if _, ok := g.handlerPreBranch[startNode]; !ok { + g.handlerPreBranch[startNode] = [][]handlerPair{} + } + branch.idx = len(g.handlerPreBranch[startNode]) + + if startNode != START && g.nodes[startNode].executorMeta.component == ComponentOfPassthrough { + g.nodes[startNode].cr.inputType = branch.inputType + g.nodes[startNode].cr.outputType = branch.inputType + g.nodes[startNode].cr.genericHelper = branch.genericHelper.forPredecessorPassthrough() + } + + // check branch condition type + result := checkAssignable(g.getNodeOutputType(startNode), branch.inputType) + if result == assignableTypeMustNot { + return fmt.Errorf("condition's input type[%s] and start node[%s]'s output type[%s] are mismatched", branch.inputType.String(), startNode, g.getNodeOutputType(startNode).String()) + } else if result == assignableTypeMay { + g.handlerPreBranch[startNode] = append(g.handlerPreBranch[startNode], []handlerPair{branch.inputConverter}) + } else { + g.handlerPreBranch[startNode] = append(g.handlerPreBranch[startNode], []handlerPair{}) + } + + if !skipData { + for endNode := range branch.endNodes { + if _, ok := g.nodes[endNode]; !ok { + if endNode != END { + return fmt.Errorf("branch end node '%s' needs to be added to graph first", endNode) + } + } + + g.addToValidateMap(startNode, endNode, nil) + e := g.updateToValidateMap() + if e != nil { + return e + } + + if startNode == START { + g.startNodes = append(g.startNodes, endNode) + } + if endNode == END { + g.endNodes = append(g.endNodes, startNode) + } + } + } else { + for endNode := range branch.endNodes { + if startNode == START { + g.startNodes = append(g.startNodes, endNode) + } + if endNode == END { + g.endNodes = append(g.endNodes, startNode) + } + } + branch.noDataFlow = true + } + + g.branches[startNode] = append(g.branches[startNode], branch) + + return nil +} + +func (g *graph) addToValidateMap(startNode, endNode string, mapping []*FieldMapping) { + g.toValidateMap[startNode] = append(g.toValidateMap[startNode], struct { + endNode string + mappings []*FieldMapping + }{endNode: endNode, mappings: mapping}) +} + +// updateToValidateMap after update node, check validate map +// check again if nodes in toValidateMap have been updated. because when there are multiple linked passthrough nodes, in the worst scenario, only one node can be updated at a time. +func (g *graph) updateToValidateMap() error { + var startNodeOutputType, endNodeInputType reflect.Type + for { + hasChanged := false + for startNode := range g.toValidateMap { + startNodeOutputType = g.getNodeOutputType(startNode) + + for i := 0; i < len(g.toValidateMap[startNode]); i++ { + endNode := g.toValidateMap[startNode][i] + + endNodeInputType = g.getNodeInputType(endNode.endNode) + if startNodeOutputType == nil && endNodeInputType == nil { + continue + } + + // update toValidateMap + g.toValidateMap[startNode] = append(g.toValidateMap[startNode][:i], g.toValidateMap[startNode][i+1:]...) + i-- + + hasChanged = true + // assume that START and END type isn't empty + if startNodeOutputType != nil && endNodeInputType == nil { + g.nodes[endNode.endNode].cr.inputType = startNodeOutputType + g.nodes[endNode.endNode].cr.outputType = g.nodes[endNode.endNode].cr.inputType + g.nodes[endNode.endNode].cr.genericHelper = g.getNodeGenericHelper(startNode).forSuccessorPassthrough() + } else if startNodeOutputType == nil /* redundant condition && endNodeInputType != nil */ { + g.nodes[startNode].cr.inputType = endNodeInputType + g.nodes[startNode].cr.outputType = g.nodes[startNode].cr.inputType + g.nodes[startNode].cr.genericHelper = g.getNodeGenericHelper(endNode.endNode).forPredecessorPassthrough() + } else if len(endNode.mappings) == 0 { + // common node check + result := checkAssignable(startNodeOutputType, endNodeInputType) + if result == assignableTypeMustNot { + return fmt.Errorf("graph edge[%s]-[%s]: start node's output type[%s] and end node's input type[%s] mismatch", + startNode, endNode.endNode, startNodeOutputType.String(), endNodeInputType.String()) + } else if result == assignableTypeMay { + // add runtime check edges + if _, ok := g.handlerOnEdges[startNode]; !ok { + g.handlerOnEdges[startNode] = make(map[string][]handlerPair) + } + g.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], g.getNodeGenericHelper(endNode.endNode).inputConverter) + } + continue + } + + if len(endNode.mappings) > 0 { + if _, ok := g.handlerOnEdges[startNode]; !ok { + g.handlerOnEdges[startNode] = make(map[string][]handlerPair) + } + g.fieldMappingRecords[endNode.endNode] = append(g.fieldMappingRecords[endNode.endNode], endNode.mappings...) + + // field mapping check + checker, uncheckedSourcePaths, err := validateFieldMapping(g.getNodeOutputType(startNode), g.getNodeInputType(endNode.endNode), endNode.mappings) + if err != nil { + return err + } + + g.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], handlerPair{ + invoke: func(value any) (any, error) { + return fieldMap(endNode.mappings, false, uncheckedSourcePaths)(value) + }, + transform: streamFieldMap(endNode.mappings, uncheckedSourcePaths), + }) + + if checker != nil { + g.handlerOnEdges[startNode][endNode.endNode] = append(g.handlerOnEdges[startNode][endNode.endNode], *checker) + } + } + } + } + if !hasChanged { + break + } + } + + return nil +} + +func (g *graph) getNodeGenericHelper(name string) *genericHelper { + if name == START { + return g.genericHelper.forPredecessorPassthrough() + } else if name == END { + return g.genericHelper.forSuccessorPassthrough() + } + return g.nodes[name].getGenericHelper() +} + +func (g *graph) getNodeInputType(name string) reflect.Type { + if name == START { + return g.inputType() + } else if name == END { + return g.outputType() + } + return g.nodes[name].inputType() +} + +func (g *graph) getNodeOutputType(name string) reflect.Type { + if name == START { + return g.inputType() + } else if name == END { + return g.outputType() + } + return g.nodes[name].outputType() +} + +func (g *graph) inputType() reflect.Type { + return g.expectedInputType +} + +func (g *graph) outputType() reflect.Type { + return g.expectedOutputType +} + +func (g *graph) compile(ctx context.Context, opt *graphCompileOptions) (*composableRunnable, error) { + if g.buildError != nil { + return nil, g.buildError + } + + // get run type + runType := runTypePregel + cb := pregelChannelBuilder + if isChain(g.cmp) || isWorkflow(g.cmp) { + if opt != nil && opt.nodeTriggerMode != "" { + return nil, errors.New(fmt.Sprintf("%s doesn't support node trigger mode option", g.cmp)) + } + } + if (opt != nil && opt.nodeTriggerMode == AllPredecessor) || isWorkflow(g.cmp) { + runType = runTypeDAG + cb = dagChannelBuilder + } + + // get eager type + eager := false + if isWorkflow(g.cmp) || runType == runTypeDAG { + eager = true + } + if opt != nil && opt.eagerDisabled { + eager = false + } + + if len(g.startNodes) == 0 { + return nil, errors.New("start node not set") + } + if len(g.endNodes) == 0 { + return nil, errors.New("end node not set") + } + + // toValidateMap isn't empty means there are nodes that cannot infer type + for _, v := range g.toValidateMap { + if len(v) > 0 { + return nil, fmt.Errorf("some node's input or output types cannot be inferred: %v", g.toValidateMap) + } + } + + for key := range g.fieldMappingRecords { + // not allowed to map multiple fields to the same field + toMap := make(map[string]bool) + for _, mapping := range g.fieldMappingRecords[key] { + if _, ok := toMap[mapping.to]; ok { + return nil, fmt.Errorf("duplicate mapping target field: %s of node[%s]", mapping.to, key) + } + toMap[mapping.to] = true + } + + // add map to input converter + g.handlerPreNode[key] = append(g.handlerPreNode[key], g.getNodeGenericHelper(key).inputFieldMappingConverter) + } + + key2SubGraphs := g.beforeChildGraphsCompile(opt) + chanSubscribeTo := make(map[string]*chanCall) + for name, node := range g.nodes { + node.beforeChildGraphCompile(name, key2SubGraphs) + + r, err := node.compileIfNeeded(ctx) + if err != nil { + return nil, err + } + + chCall := &chanCall{ + action: r, + writeTo: g.dataEdges[name], + controls: g.controlEdges[name], + + preProcessor: node.nodeInfo.preProcessor, + postProcessor: node.nodeInfo.postProcessor, + } + + branches := g.branches[name] + if len(branches) > 0 { + branchRuns := make([]*GraphBranch, 0, len(branches)) + branchRuns = append(branchRuns, branches...) + + chCall.writeToBranches = branchRuns + } + + chanSubscribeTo[name] = chCall + } + + dataPredecessors := make(map[string][]string) + controlPredecessors := make(map[string][]string) + for start, ends := range g.controlEdges { + for _, end := range ends { + if _, ok := controlPredecessors[end]; !ok { + controlPredecessors[end] = []string{start} + } else { + controlPredecessors[end] = append(controlPredecessors[end], start) + } + } + } + for start, ends := range g.dataEdges { + for _, end := range ends { + if _, ok := dataPredecessors[end]; !ok { + dataPredecessors[end] = []string{start} + } else { + dataPredecessors[end] = append(dataPredecessors[end], start) + } + } + } + for start, branches := range g.branches { + for _, branch := range branches { + for end := range branch.endNodes { + if _, ok := controlPredecessors[end]; !ok { + controlPredecessors[end] = []string{start} + } else { + controlPredecessors[end] = append(controlPredecessors[end], start) + } + + if !branch.noDataFlow { + if _, ok := dataPredecessors[end]; !ok { + dataPredecessors[end] = []string{start} + } else { + dataPredecessors[end] = append(dataPredecessors[end], start) + } + } + } + } + } + + inputChannels := &chanCall{ + writeTo: g.dataEdges[START], + controls: g.controlEdges[START], + writeToBranches: make([]*GraphBranch, len(g.branches[START])), + } + copy(inputChannels.writeToBranches, g.branches[START]) + + var mergeConfigs map[string]FanInMergeConfig + if opt != nil { + mergeConfigs = opt.mergeConfigs + } + if mergeConfigs == nil { + mergeConfigs = make(map[string]FanInMergeConfig) + } + + r := &runner{ + chanSubscribeTo: chanSubscribeTo, + controlPredecessors: controlPredecessors, + dataPredecessors: dataPredecessors, + + inputChannels: inputChannels, + + eager: eager, + + chanBuilder: cb, + + inputType: g.inputType(), + outputType: g.outputType(), + genericHelper: g.genericHelper, + + preBranchHandlerManager: &preBranchHandlerManager{h: g.handlerPreBranch}, + preNodeHandlerManager: &preNodeHandlerManager{h: g.handlerPreNode}, + edgeHandlerManager: &edgeHandlerManager{h: g.handlerOnEdges}, + + mergeConfigs: mergeConfigs, + } + + successors := make(map[string][]string) + for ch := range r.chanSubscribeTo { + successors[ch] = getSuccessors(r.chanSubscribeTo[ch]) + } + r.successors = successors + + if g.stateGenerator != nil { + r.runCtx = func(ctx context.Context) context.Context { + var parent *internalState + if p, ok := ctx.Value(stateKey{}).(*internalState); ok { + parent = p + } + + return context.WithValue(ctx, stateKey{}, &internalState{ + state: g.stateGenerator(ctx), + parent: parent, + }) + } + } + + if runType == runTypeDAG { + err := validateDAG(r.chanSubscribeTo, controlPredecessors) + if err != nil { + return nil, err + } + r.dag = true + } + + if opt != nil { + inputPairs := make(map[string]streamConvertPair) + outputPairs := make(map[string]streamConvertPair) + for key, c := range r.chanSubscribeTo { + inputPairs[key] = c.action.inputStreamConvertPair + outputPairs[key] = c.action.outputStreamConvertPair + } + inputPairs[END] = r.outputConvertStreamPair + outputPairs[START] = r.inputConvertStreamPair + r.checkPointer = newCheckPointer(inputPairs, outputPairs, opt.checkPointStore, opt.serializer) + + r.interruptBeforeNodes = opt.interruptBeforeNodes + r.interruptAfterNodes = opt.interruptAfterNodes + r.options = *opt + } + + // default options + if r.dag && r.options.maxRunSteps > 0 { + return nil, fmt.Errorf("cannot set max run steps in dag mode") + } else if !r.dag && r.options.maxRunSteps == 0 { + r.options.maxRunSteps = len(r.chanSubscribeTo) + 10 + } + + g.compiled = true + + g.onCompileFinish(ctx, opt, key2SubGraphs) + + return r.toComposableRunnable(), nil +} + +func getSuccessors(c *chanCall) []string { + ret := make([]string, len(c.writeTo)) + copy(ret, c.writeTo) + ret = append(ret, c.controls...) + for _, branch := range c.writeToBranches { + for node := range branch.endNodes { + ret = append(ret, node) + } + } + return uniqueSlice(ret) +} + +func uniqueSlice(s []string) []string { + seen := make(map[string]struct{}, len(s)) + cur := 0 + for i := range s { + if _, ok := seen[s[i]]; !ok { + seen[s[i]] = struct{}{} + s[cur] = s[i] + cur++ + } + } + return s[:cur] +} + +type subGraphCompileCallback struct { + closure func(ctx context.Context, info *GraphInfo) +} + +// OnFinish is called when the graph is compiled. +func (s *subGraphCompileCallback) OnFinish(ctx context.Context, info *GraphInfo) { + s.closure(ctx, info) +} + +func (g *graph) beforeChildGraphsCompile(opt *graphCompileOptions) map[string]*GraphInfo { + if opt == nil || len(opt.callbacks) == 0 { + return nil + } + + return make(map[string]*GraphInfo) +} + +func (gn *graphNode) beforeChildGraphCompile(nodeKey string, key2SubGraphs map[string]*GraphInfo) { + if gn.g == nil || key2SubGraphs == nil { + return + } + + subGraphCallback := func(ctx2 context.Context, subGraph *GraphInfo) { + key2SubGraphs[nodeKey] = subGraph + } + + gn.nodeInfo.compileOption.callbacks = append(gn.nodeInfo.compileOption.callbacks, &subGraphCompileCallback{closure: subGraphCallback}) +} + +func (g *graph) toGraphInfo(opt *graphCompileOptions, key2SubGraphs map[string]*GraphInfo) *GraphInfo { + gInfo := &GraphInfo{ + CompileOptions: opt.origOpts, + Nodes: make(map[string]GraphNodeInfo, len(g.nodes)), + Edges: gmap.Clone(g.controlEdges), + DataEdges: gmap.Clone(g.dataEdges), + Branches: gmap.Map(g.branches, func(startNode string, branches []*GraphBranch) (string, []GraphBranch) { + branchInfo := make([]GraphBranch, 0, len(branches)) + for _, b := range branches { + branchInfo = append(branchInfo, GraphBranch{ + invoke: b.invoke, + collect: b.collect, + inputType: b.inputType, + genericHelper: b.genericHelper, + endNodes: gmap.Clone(b.endNodes), + }) + } + return startNode, branchInfo + }), + InputType: g.expectedInputType, + OutputType: g.expectedOutputType, + Name: opt.graphName, + GenStateFn: g.stateGenerator, + NewGraphOptions: g.newOpts, + } + + for key := range g.nodes { + gNode := g.nodes[key] + if gNode.executorMeta.component == ComponentOfPassthrough { + gInfo.Nodes[key] = GraphNodeInfo{ + Component: gNode.executorMeta.component, + GraphAddNodeOpts: gNode.opts, + InputType: gNode.cr.inputType, + OutputType: gNode.cr.outputType, + Name: gNode.nodeInfo.name, + InputKey: gNode.cr.nodeInfo.inputKey, + OutputKey: gNode.cr.nodeInfo.outputKey, + } + continue + } + + gNodeInfo := &GraphNodeInfo{ + Component: gNode.executorMeta.component, + Instance: gNode.instance, + GraphAddNodeOpts: gNode.opts, + InputType: gNode.cr.inputType, + OutputType: gNode.cr.outputType, + Name: gNode.nodeInfo.name, + InputKey: gNode.cr.nodeInfo.inputKey, + OutputKey: gNode.cr.nodeInfo.outputKey, + Mappings: g.fieldMappingRecords[key], + } + + if gi, ok := key2SubGraphs[key]; ok { + gNodeInfo.GraphInfo = gi + } + + gInfo.Nodes[key] = *gNodeInfo + } + + return gInfo +} + +func (g *graph) onCompileFinish(ctx context.Context, opt *graphCompileOptions, key2SubGraphs map[string]*GraphInfo) { + if opt == nil { + return + } + + if len(opt.callbacks) == 0 { + return + } + + gInfo := g.toGraphInfo(opt, key2SubGraphs) + + for _, cb := range opt.callbacks { + cb.OnFinish(ctx, gInfo) + } +} + +func (g *graph) getGenericHelper() *genericHelper { + return g.genericHelper +} + +func (g *graph) GetType() string { + return "" +} + +func transferTask(script [][]string, invertedEdges map[string][]string) [][]string { + utilMap := map[string]bool{} + for i := len(script) - 1; i >= 0; i-- { + for j := 0; j < len(script[i]); j++ { + // deduplicate + if _, ok := utilMap[script[i][j]]; ok { + script[i] = append(script[i][:j], script[i][j+1:]...) + j-- + continue + } + utilMap[script[i][j]] = true + + target := i + for k := i + 1; k < len(script); k++ { + hasDependencies := false + for l := range script[k] { + for _, dependency := range invertedEdges[script[i][j]] { + if script[k][l] == dependency { + hasDependencies = true + break + } + } + if hasDependencies { + break + } + } + if hasDependencies { + break + } + target = k + } + if target != i { + script[target] = append(script[target], script[i][j]) + script[i] = append(script[i][:j], script[i][j+1:]...) + j-- + } + } + } + + return script +} + +func validateDAG(chanSubscribeTo map[string]*chanCall, controlPredecessors map[string][]string) error { + m := map[string]int{} + for node := range chanSubscribeTo { + if edges, ok := controlPredecessors[node]; ok { + m[node] = len(edges) + for _, pre := range edges { + if pre == START { + m[node] -= 1 + } + } + } else { + m[node] = 0 + } + } + hasChanged := true + for hasChanged { + hasChanged = false + for node := range m { + if m[node] == 0 { + hasChanged = true + for _, subNode := range chanSubscribeTo[node].controls { + if subNode == END { + continue + } + m[subNode]-- + } + for _, subBranch := range chanSubscribeTo[node].writeToBranches { + for subNode := range subBranch.endNodes { + if subNode == END { + continue + } + m[subNode]-- + } + } + m[node] = -1 + } + } + } + + var loopStarts []string + for k, v := range m { + if v > 0 { + loopStarts = append(loopStarts, k) + } + } + if len(loopStarts) > 0 { + return fmt.Errorf("%w: %s", DAGInvalidLoopErr, formatLoops(findLoops(loopStarts, chanSubscribeTo))) + } + return nil +} + +// DAGInvalidLoopErr indicates the graph contains a cycle and is invalid. +var DAGInvalidLoopErr = errors.New("DAG is invalid, has loop") + +func findLoops(startNodes []string, chanCalls map[string]*chanCall) [][]string { + controlSuccessors := map[string][]string{} + for node, ch := range chanCalls { + controlSuccessors[node] = append(controlSuccessors[node], ch.controls...) + for _, b := range ch.writeToBranches { + for end := range b.endNodes { + controlSuccessors[node] = append(controlSuccessors[node], end) + } + } + } + + visited := map[string]bool{} + var dfs func(path []string) [][]string + dfs = func(path []string) [][]string { + var ret [][]string + pathEnd := path[len(path)-1] + successors, ok := controlSuccessors[pathEnd] + if !ok { + return nil + } + for _, successor := range successors { + visited[successor] = true + + if successor == END { + continue + } + + var looped bool + for i, node := range path { + if node == successor { + ret = append(ret, append(path[i:], successor)) + looped = true + break + } + } + if looped { + continue + } + + ret = append(ret, dfs(append(path, successor))...) + } + return ret + } + + var ret [][]string + for _, node := range startNodes { + if !visited[node] { + ret = append(ret, dfs([]string{node})...) + } + } + return ret +} + +func formatLoops(loops [][]string) string { + sb := strings.Builder{} + for _, loop := range loops { + if len(loop) == 0 { + continue + } + sb.WriteString("[") + sb.WriteString(loop[0]) + for i := 1; i < len(loop); i++ { + sb.WriteString("->") + sb.WriteString(loop[i]) + } + sb.WriteString("]") + } + return sb.String() +} + +// NewNodePath specifies a path to a node in the graph, which is composed of node keys. +// Starting from the top graph, +// following this set of node keys can lead to a specific node in the top graph or a subgraph. +// +// e.g. +// NewNodePath("sub_graph_node_key", "node_key_within_sub_graph") +func NewNodePath(nodeKeyPath ...string) *NodePath { + return &NodePath{path: nodeKeyPath} +} + +// NodePath represents a path composed of node keys to locate a node. +type NodePath struct { + path []string +} + +// GetPath returns the sequence of node keys in the path. +func (p *NodePath) GetPath() []string { + return p.path +} diff --git a/compose/graph_add_node_options.go b/compose/graph_add_node_options.go new file mode 100644 index 0000000..b1db71f --- /dev/null +++ b/compose/graph_add_node_options.go @@ -0,0 +1,168 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "reflect" + + "github.com/cloudwego/eino/internal/generic" +) + +type graphAddNodeOpts struct { + nodeOptions *nodeOptions + processor *processorOpts + + needState bool +} + +// GraphAddNodeOpt is a functional option type for adding a node to a graph. +// e.g. +// +// graph.AddNode("node_name", node, compose.WithInputKey("input_key"), compose.WithOutputKey("output_key")) +type GraphAddNodeOpt func(o *graphAddNodeOpts) + +type nodeOptions struct { + nodeName string + + nodeKey string + + inputKey string + outputKey string + + graphCompileOption []GraphCompileOption // when this node is itself an AnyGraph, this option will be used to compile the node as a nested graph +} + +// WithNodeName sets the name of the node. +func WithNodeName(n string) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.nodeOptions.nodeName = n + } +} + +// WithNodeKey set the node key, which is used to identify the node in the chain. +// only for use in Chain/StateChain. +func WithNodeKey(key string) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.nodeOptions.nodeKey = key + } +} + +// WithInputKey sets the input key of the node. +// this will change the input value of the node, for example, if the pre node's output is map[string]any{"key01": "value01"}, +// and the current node's input key is "key01", then the current node's input value will be "value01". +func WithInputKey(k string) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.nodeOptions.inputKey = k + } +} + +// WithOutputKey sets the output key of the node. +// this will change the output value of the node, for example, if the current node's output key is "key01", +// then the node's output value will be map[string]any{"key01": value}. +func WithOutputKey(k string) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.nodeOptions.outputKey = k + } +} + +// WithGraphCompileOptions when the node is an AnyGraph, use this option to set compile option for the node. +// e.g. +// +// graph.AddNode("node_name", node, compose.WithGraphCompileOptions(compose.WithGraphName("my_sub_graph"))) +func WithGraphCompileOptions(opts ...GraphCompileOption) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.nodeOptions.graphCompileOption = opts + } +} + +// WithStatePreHandler modify node's input of I according to state S and input or store input information into state, and it's thread-safe. +// notice: this option requires Graph to be created with WithGenLocalState option. +// I: input type of the Node like ChatModel, Lambda, Retriever etc. +// S: state type defined in WithGenLocalState +func WithStatePreHandler[I, S any](pre StatePreHandler[I, S]) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.processor.statePreHandler = convertPreHandler(pre) + o.processor.preStateType = generic.TypeOf[S]() + o.needState = true + } +} + +// WithStatePostHandler modify node's output of O according to state S and output or store output information into state, and it's thread-safe. +// notice: this option requires Graph to be created with WithGenLocalState option. +// O: output type of the Node like ChatModel, Lambda, Retriever etc. +// S: state type defined in WithGenLocalState +func WithStatePostHandler[O, S any](post StatePostHandler[O, S]) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.processor.statePostHandler = convertPostHandler(post) + o.processor.postStateType = generic.TypeOf[S]() + o.needState = true + } +} + +// WithStreamStatePreHandler modify node's streaming input of I according to state S and input or store input information into state, and it's thread-safe. +// notice: this option requires Graph to be created with WithGenLocalState option. +// when to use: when upstream node's output is an actual stream, and you want the current node's input to remain an actual stream after state pre handler. +// caution: while StreamStatePreHandler is thread safe, modifying state within your own goroutine is NOT. +// I: input type of the Node like ChatModel, Lambda, Retriever etc. +// S: state type defined in WithGenLocalState +func WithStreamStatePreHandler[I, S any](pre StreamStatePreHandler[I, S]) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.processor.statePreHandler = streamConvertPreHandler(pre) + o.processor.preStateType = generic.TypeOf[S]() + o.needState = true + } +} + +// WithStreamStatePostHandler modify node's streaming output of O according to state S and output or store output information into state, and it's thread-safe. +// notice: this option requires Graph to be created with WithGenLocalState option. +// when to use: when current node's output is an actual stream, and you want the downstream node's input to remain an actual stream after state post handler. +// caution: while StreamStatePostHandler is thread safe, modifying state within your own goroutine is NOT. +// O: output type of the Node like ChatModel, Lambda, Retriever etc. +// S: state type defined in WithGenLocalState +func WithStreamStatePostHandler[O, S any](post StreamStatePostHandler[O, S]) GraphAddNodeOpt { + return func(o *graphAddNodeOpts) { + o.processor.statePostHandler = streamConvertPostHandler(post) + o.processor.postStateType = generic.TypeOf[S]() + o.needState = true + } +} + +type processorOpts struct { + statePreHandler *composableRunnable + preStateType reflect.Type // used for type validation + statePostHandler *composableRunnable + postStateType reflect.Type // used for type validation +} + +func getGraphAddNodeOpts(opts ...GraphAddNodeOpt) *graphAddNodeOpts { + opt := &graphAddNodeOpts{ + nodeOptions: &nodeOptions{ + nodeName: "", + nodeKey: "", + }, + processor: &processorOpts{ + statePreHandler: nil, + statePostHandler: nil, + }, + } + + for _, fn := range opts { + fn(opt) + } + + return opt +} diff --git a/compose/graph_call_options.go b/compose/graph_call_options.go new file mode 100644 index 0000000..eb2ab16 --- /dev/null +++ b/compose/graph_call_options.go @@ -0,0 +1,264 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + "time" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" +) + +type graphCancelChanKey struct{} +type graphCancelChanVal struct { + ch chan *time.Duration +} + +type graphInterruptOptions struct { + timeout *time.Duration +} + +// GraphInterruptOption configures behavior when interrupting a running graph. +type GraphInterruptOption func(o *graphInterruptOptions) + +// WithGraphInterruptTimeout specifies the max waiting time before generating an interrupt. +// After the max waiting time, the graph will force an interrupt. Any unfinished tasks will be re-run when the graph is resumed. +func WithGraphInterruptTimeout(timeout time.Duration) GraphInterruptOption { + return func(o *graphInterruptOptions) { + o.timeout = &timeout + } +} + +// WithGraphInterrupt creates a context with graph cancellation support. +// When the returned context is used to invoke a graph or workflow, calling the interrupt function will trigger an interrupt. +// The graph will wait for current tasks to complete by default. +// +// Input Persistence: When WithGraphInterrupt is used, ALL nodes (in both root graph and subgraphs) will automatically +// persist their inputs (both streaming and non-streaming) before execution. If the graph is interrupted, these inputs +// are restored when the graph resumes from a checkpoint, ensuring interrupted nodes receive their original inputs. +// +// This behavior differs from internal interrupts triggered via compose.Interrupt() within a node's function body. +// Internal interrupts do NOT automatically persist inputs - the node author must manage input persistence manually, +// either by saving it in the global graph state or using compose.StatefulInterrupt() to store it in local interrupt state. +// WithGraphInterrupt enables automatic input persistence because external interrupts can occur at any point during +// node execution, making it impossible for the node to prepare for the interrupt. +// +// Why input persistence is not enabled by default for internal interrupts: Enabling it universally would break +// existing code that relies on checking "input == nil" to determine whether the node is running for the first time +// or resuming from an interrupt. The recommended approach is to use compose.GetInterruptState() to explicitly +// determine whether the current execution is a first run or a resume. +func WithGraphInterrupt(parent context.Context) (ctx context.Context, interrupt func(opts ...GraphInterruptOption)) { + ch := make(chan *time.Duration, 1) + ctx = context.WithValue(parent, graphCancelChanKey{}, &graphCancelChanVal{ + ch: ch, + }) + return ctx, func(opts ...GraphInterruptOption) { + o := &graphInterruptOptions{} + for _, opt := range opts { + opt(o) + } + ch <- o.timeout + close(ch) + } +} + +func getGraphCancel(ctx context.Context) *graphCancelChanVal { + val, ok := ctx.Value(graphCancelChanKey{}).(*graphCancelChanVal) + if !ok { + return nil + } + return val +} + +// Option is a functional option type for calling a graph. +type Option struct { + options []any + handler []callbacks.Handler + + paths []*NodePath + + maxRunSteps int + checkPointID *string + writeToCheckPointID *string + forceNewRun bool + stateModifier StateModifier +} + +func (o Option) deepCopy() Option { + nOptions := make([]any, len(o.options)) + copy(nOptions, o.options) + nHandler := make([]callbacks.Handler, len(o.handler)) + copy(nHandler, o.handler) + nPaths := make([]*NodePath, len(o.paths)) + for i, path := range o.paths { + nPath := *path + nPaths[i] = &nPath + } + return Option{ + options: nOptions, + handler: nHandler, + paths: nPaths, + maxRunSteps: o.maxRunSteps, + } +} + +// DesignateNode sets the key of the node to which the option will be applied. +// notice: only effective at the top graph. +// e.g. +// +// embeddingOption := compose.WithEmbeddingOption(embedding.WithModel("text-embedding-3-small")) +// runnable.Invoke(ctx, "input", embeddingOption.DesignateNode("embedding_node_key")) +func (o Option) DesignateNode(nodeKey ...string) Option { + nKeys := make([]*NodePath, len(nodeKey)) + for i, k := range nodeKey { + nKeys[i] = NewNodePath(k) + } + return o.DesignateNodeWithPath(nKeys...) +} + +// DesignateNodeWithPath sets the path of the node(s) to which the option will be applied. +// You can specify a node in the subgraph through `NodePath` to make the option only take effect at this node. +// +// e.g. +// nodePath := NewNodePath("sub_graph_node_key", "node_key_within_sub_graph") +// DesignateNodeWithPath(nodePath) +func (o Option) DesignateNodeWithPath(path ...*NodePath) Option { + o.paths = append(o.paths, path...) + return o +} + +// WithEmbeddingOption is a functional option type for embedding component. +// e.g. +// +// embeddingOption := compose.WithEmbeddingOption(embedding.WithModel("text-embedding-3-small")) +// runnable.Invoke(ctx, "input", embeddingOption) +func WithEmbeddingOption(opts ...embedding.Option) Option { + return withComponentOption(opts...) +} + +// WithRetrieverOption is a functional option type for retriever component. +// e.g. +// +// retrieverOption := compose.WithRetrieverOption(retriever.WithIndex("my_index")) +// runnable.Invoke(ctx, "input", retrieverOption) +func WithRetrieverOption(opts ...retriever.Option) Option { + return withComponentOption(opts...) +} + +// WithLoaderOption is a functional option type for loader component. +// e.g. +// +// loaderOption := compose.WithLoaderOption(document.WithCollection("my_collection")) +// runnable.Invoke(ctx, "input", loaderOption) +func WithLoaderOption(opts ...document.LoaderOption) Option { + return withComponentOption(opts...) +} + +// WithDocumentTransformerOption is a functional option type for document transformer component. +func WithDocumentTransformerOption(opts ...document.TransformerOption) Option { + return withComponentOption(opts...) +} + +// WithIndexerOption is a functional option type for indexer component. +// e.g. +// +// indexerOption := compose.WithIndexerOption(indexer.WithSubIndexes([]string{"my_sub_index"})) +// runnable.Invoke(ctx, "input", indexerOption) +func WithIndexerOption(opts ...indexer.Option) Option { + return withComponentOption(opts...) +} + +// WithChatModelOption is a functional option type for chat model component. +// e.g. +// +// chatModelOption := compose.WithChatModelOption(model.WithTemperature(0.7)) +// runnable.Invoke(ctx, "input", chatModelOption) +func WithChatModelOption(opts ...model.Option) Option { + return withComponentOption(opts...) +} + +// WithChatTemplateOption is a functional option type for chat template component. +func WithChatTemplateOption(opts ...prompt.Option) Option { + return withComponentOption(opts...) +} + +// WithToolsNodeOption is a functional option type for tools node component. +func WithToolsNodeOption(opts ...ToolsNodeOption) Option { + return withComponentOption(opts...) +} + +// WithLambdaOption is a functional option type for lambda component. +func WithLambdaOption(opts ...any) Option { + return Option{ + options: opts, + paths: make([]*NodePath, 0), + } +} + +// WithCallbacks set callback handlers for all components in a single call. +// e.g. +// +// runnable.Invoke(ctx, "input", compose.WithCallbacks(&myCallbacks{})) +func WithCallbacks(cbs ...callbacks.Handler) Option { + return Option{ + handler: cbs, + } +} + +// WithRuntimeMaxSteps sets the maximum number of steps for the graph runtime. +// e.g. +// +// runnable.Invoke(ctx, "input", compose.WithRuntimeMaxSteps(20)) +func WithRuntimeMaxSteps(maxSteps int) Option { + return Option{ + maxRunSteps: maxSteps, + } +} + +func withComponentOption[TOption any](opts ...TOption) Option { + o := make([]any, 0, len(opts)) + for i := range opts { + o = append(o, opts[i]) + } + return Option{ + options: o, + paths: make([]*NodePath, 0), + } +} + +func convertOption[TOption any](opts ...any) ([]TOption, error) { + if len(opts) == 0 { + return nil, nil + } + ret := make([]TOption, 0, len(opts)) + for i := range opts { + o, ok := opts[i].(TOption) + if !ok { + return nil, fmt.Errorf("unexpected component option type, expected:%s, actual:%s", reflect.TypeOf((*TOption)(nil)).Elem().String(), reflect.TypeOf(opts[i]).String()) + } + ret = append(ret, o) + } + return ret, nil +} diff --git a/compose/graph_call_options_test.go b/compose/graph_call_options_test.go new file mode 100644 index 0000000..708860b --- /dev/null +++ b/compose/graph_call_options_test.go @@ -0,0 +1,496 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/retriever" + mockDocument "github.com/cloudwego/eino/internal/mock/components/document" + mockEmbedding "github.com/cloudwego/eino/internal/mock/components/embedding" + mockRetriever "github.com/cloudwego/eino/internal/mock/components/retriever" + "github.com/cloudwego/eino/schema" +) + +var optionSuccess = true +var idx int + +func checkOption(opts ...model.Option) bool { + if len(opts) != 2 { + return false + } + o := model.GetCommonOptions(&model.Options{}, opts...) + if o.TopP == nil || *o.TopP != 1.0 { + return false + } + if o.Model == nil { + return false + } + if idx == 0 { + idx = 1 + if o.Model == nil || *o.Model != "123" { + return false + } + } else { + idx = 0 + if o.Model == nil || *o.Model != "456" { + return false + } + } + + return true +} + +type testModel struct{} + +func (t *testModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +func (t *testModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + if !checkOption(opts...) { + optionSuccess = false + } + return &schema.Message{}, nil +} + +func (t *testModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + if !checkOption(opts...) { + optionSuccess = false + } + sr, sw := schema.Pipe[*schema.Message](1) + sw.Send(nil, nil) + sw.Close() + return sr, nil +} + +func TestCallOption(t *testing.T) { + g := NewGraph[[]*schema.Message, *schema.Message]() + err := g.AddLambdaNode("1", InvokableLambdaWithOption(func(ctx context.Context, input []*schema.Message, opts ...string) (output []*schema.Message, err error) { + if len(opts) != 1 || opts[0] != "1" { + t.Fatalf("lambda option length isn't 1 or content isn't '1': %v", opts) + } + return input, nil + })) + assert.Nil(t, err) + + err = g.AddChatModelNode("2", &testModel{}) + assert.Nil(t, err) + + err = g.AddLambdaNode("-", InvokableLambda(func(ctx context.Context, input *schema.Message) (output []*schema.Message, err error) { + return []*schema.Message{input}, nil + })) + assert.Nil(t, err) + + err = g.AddChatModelNode("3", &testModel{}) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", "2") + assert.Nil(t, err) + + err = g.AddEdge("2", "-") + assert.Nil(t, err) + + err = g.AddEdge("-", "3") + assert.Nil(t, err) + + err = g.AddEdge("3", END) + assert.Nil(t, err) + + ctx := context.Background() + + r, err := g.Compile(ctx) + assert.Nil(t, err) + + sessionKey := struct{}{} + startCnt := 0 + endCnt := 0 + opts := []Option{ + WithChatModelOption( + model.WithModel("123"), + ).DesignateNode("2"), + WithChatModelOption( + model.WithModel("456"), + ).DesignateNode("3"), + WithChatModelOption( + model.WithTopP(1.0), + ), + WithCallbacks(callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + startCnt++ + return context.WithValue(ctx, sessionKey, "start") + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + if ctx.Value(sessionKey).(string) == "start" { + endCnt++ + return context.WithValue(ctx, sessionKey, "end") + } + return ctx + }).Build()).DesignateNode("3"), + WithLambdaOption("1").DesignateNode("1"), + } + + _, err = r.Invoke(ctx, []*schema.Message{}, + opts...) + if err != nil { + t.Fatal(err) + } + if !optionSuccess { + t.Fatal("invoke option fail") + } + if startCnt != 1 { + t.Fatal("node callback fail") + } + if endCnt != 1 { + t.Fatal("node callback fail") + } + _, err = r.Stream(ctx, []*schema.Message{}, + opts...) + if err != nil { + t.Fatal(err) + } + if !optionSuccess { + t.Fatal("stream option fail") + } + + srOfCollect, swOfCollect := schema.Pipe[[]*schema.Message](1) + swOfCollect.Send([]*schema.Message{}, nil) + swOfCollect.Close() + _, err = r.Collect(ctx, srOfCollect, opts...) + assert.Nil(t, err) + + if !optionSuccess { + t.Fatal("collect option fail") + } + + srOfTransform, swOfTransform := schema.Pipe[[]*schema.Message](1) + swOfTransform.Send([]*schema.Message{}, nil) + swOfTransform.Close() + _, err = r.Transform(ctx, srOfTransform, opts...) + assert.Nil(t, err) + + if !optionSuccess { + t.Fatal("transform option fail") + } +} + +func TestCallOptionsOneByOne(t *testing.T) { + ctx := context.Background() + t.Run("common_option", func(t *testing.T) { + type option struct { + uid int64 + } + + opt := withComponentOption(&option{uid: 100}) + assert.Len(t, opt.options, 1) + assert.IsType(t, &option{}, opt.options[0]) + assert.Equal(t, &option{uid: 100}, opt.options[0]) + }) + + t.Run("embedding_option", func(t *testing.T) { + ctrl := gomock.NewController(t) + inst := mockEmbedding.NewMockEmbedder(ctrl) + var opt *embedding.Options + inst.EXPECT().EmbedStrings(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) { + opt = embedding.GetCommonOptions(&embedding.Options{}, opts...) + return nil, nil + }).Times(1) + ch := NewChain[map[string]any, map[string]any]() + ch.AppendEmbedding(inst, WithInputKey("input"), WithOutputKey("output")) + r, err := ch.Compile(ctx) + assert.NoError(t, err) + outs, err := r.Invoke(ctx, + map[string]any{"input": []string{}}, + WithEmbeddingOption(embedding.WithModel("123")), + ) + assert.NoError(t, err) + assert.Contains(t, outs, "output") + + assert.NotNil(t, opt.Model) + assert.Equal(t, "123", *opt.Model) + }) + + t.Run("retriever_option", func(t *testing.T) { + ctrl := gomock.NewController(t) + inst := mockRetriever.NewMockRetriever(ctrl) + var opt *retriever.Options + inst.EXPECT().Retrieve(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + opt = retriever.GetCommonOptions(&retriever.Options{}, opts...) + return nil, nil + }). + Times(1) + ch := NewChain[map[string]any, map[string]any]() + ch.AppendRetriever(inst, WithInputKey("input"), WithOutputKey("output")) + r, err := ch.Compile(ctx) + assert.NoError(t, err) + outs, err := r.Invoke(ctx, + map[string]any{"input": "hi"}, + WithRetrieverOption(retriever.WithIndex("123")), + ) + assert.NoError(t, err) + assert.Contains(t, outs, "output") + + assert.NotNil(t, opt.Index) + assert.Equal(t, "123", *opt.Index) + }) + + t.Run("loader_option", func(t *testing.T) { + ctrl := gomock.NewController(t) + inst := mockDocument.NewMockLoader(ctrl) + type implOption struct { + uid int64 + } + + type implOptFn func(o *implOption) + + withUID := func(uid int64) document.LoaderOption { + return document.WrapLoaderImplSpecificOptFn[implOption](func(i *implOption) { + i.uid = uid + }) + } + + var opt *implOption + + inst.EXPECT().Load(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, src document.Source, opts ...document.LoaderOption) ([]*schema.Document, error) { + opt = document.GetLoaderImplSpecificOptions[implOption](&implOption{uid: 1}, opts...) + return nil, nil + }). + Times(1) + ch := NewChain[map[string]any, map[string]any]() + ch.AppendLoader(inst, WithInputKey("input"), WithOutputKey("output")) + r, err := ch.Compile(ctx) + assert.NoError(t, err) + outs, err := r.Invoke(ctx, + map[string]any{"input": document.Source{}}, + WithLoaderOption(withUID(123)), + ) + assert.NoError(t, err) + assert.Contains(t, outs, "output") + + assert.Equal(t, int64(123), opt.uid) + }) +} + +func TestCallOptionInSubGraph(t *testing.T) { + ctx := context.Background() + + type child1Option string + type child2Option string + type parentOption string + type grandparentOption string + + child1 := NewGraph[string, string]() + err := child1.AddLambdaNode("1", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...child1Option) (output string, err error) { + if len(opts) != 1 || opts[0] != "child1-1" { + t.Fatal("child1-1 option error") + } + return input + " child1-1", nil + }), WithNodeName("child1-1")) + assert.NoError(t, err) + err = child1.AddEdge(START, "1") + assert.NoError(t, err) + err = child1.AddEdge("1", END) + assert.NoError(t, err) + + child2 := NewGraph[string, string]() + err = child2.AddLambdaNode("1", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...child2Option) (output string, err error) { + if len(opts) != 1 || opts[0] != "child2-1" { + t.Fatal("child2-1 option error") + } + return input + " child2-1", nil + }), WithNodeName("child2-1")) + assert.NoError(t, err) + err = child2.AddEdge(START, "1") + assert.NoError(t, err) + err = child2.AddEdge("1", END) + assert.NoError(t, err) + + parent := NewGraph[string, string]() + err = parent.AddLambdaNode("1", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...parentOption) (output string, err error) { + if len(opts) != 1 || opts[0] != "parent-1" { + t.Fatal("parent-1 option error") + } + return input + " parent-1", nil + }), WithNodeName("parent-1")) + assert.NoError(t, err) + err = parent.AddGraphNode("2", child1, WithNodeName("child1")) + assert.NoError(t, err) + err = parent.AddGraphNode("3", child2, WithNodeName("child2")) + assert.NoError(t, err) + err = parent.AddEdge(START, "1") + assert.NoError(t, err) + err = parent.AddEdge("1", "2") + assert.NoError(t, err) + err = parent.AddEdge("2", "3") + assert.NoError(t, err) + err = parent.AddEdge("3", END) + assert.NoError(t, err) + + grandParent := NewGraph[string, string]() + err = grandParent.AddLambdaNode("1", InvokableLambdaWithOption(func(ctx context.Context, input string, opts ...grandparentOption) (output string, err error) { + if len(opts) != 1 || opts[0] != "grandparent-1" { + t.Fatal("grandparent-1 option error") + } + return input + " grandparent-1", nil + }), WithNodeName("grandparent-1")) + assert.NoError(t, err) + err = grandParent.AddGraphNode("2", parent, WithNodeName("parent")) + assert.NoError(t, err) + err = grandParent.AddEdge(START, "1") + assert.NoError(t, err) + err = grandParent.AddEdge("1", "2") + assert.NoError(t, err) + err = grandParent.AddEdge("2", END) + assert.NoError(t, err) + + r, err := grandParent.Compile(ctx, WithGraphName("grandparent")) + assert.NoError(t, err) + + grandCommonTimes := 0 + grandCommonCB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + switch grandCommonTimes { + case 0: + if info.Name != "grandparent" || info.Component != ComponentOfGraph { + t.Fatal("grandparent common callback 0 error") + } + case 1: + if info.Name != "grandparent-1" { + t.Fatal("grandparent common callback 1 error") + } + case 2: + if info.Name != "parent" { + t.Fatal("grandparent common callback 2 error") + } + case 3: + if info.Name != "parent-1" { + t.Fatal("grandparent common callback 3 error") + } + case 4: + if info.Name != "child1" { + t.Fatal("grandparent common callback 4 error") + } + case 5: + if info.Name != "child1-1" { + t.Fatal("grandparent common callback 5 error") + } + case 6: + if info.Name != "child2" { + t.Fatal("grandparent common callback 6 error") + } + case 7: + if info.Name != "child2-1" { + t.Fatal("grandparent common callback 7 error") + } + default: + t.Fatal("grandparent common callback too many") + } + grandCommonTimes++ + return ctx + }).Build() + grand1CB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Name != "grandparent-1" { + t.Fatal("grandparent common callback 0 error") + } + return ctx + }).Build() + parentCommonCBTimes := 0 + parentCommonCB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + switch parentCommonCBTimes { + case 0: + if info.Name != "parent" { + t.Fatal("parent common callback 0 error") + } + case 1: + if info.Name != "parent-1" { + t.Fatal("parent common callback 1 error") + } + case 2: + if info.Name != "child1" { + t.Fatal("parent common callback 2 error") + } + case 3: + if info.Name != "child1-1" { + t.Fatal("parent common callback 3 error") + } + case 4: + if info.Name != "child2" { + t.Fatal("parent common callback 4 error") + } + case 5: + if info.Name != "child2-1" { + t.Fatal("parent common callback 5 error") + } + default: + t.Fatal("parent common callback too many") + } + parentCommonCBTimes++ + return ctx + }).Build() + child1CommonCBTimes := 0 + child1CommonCB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + switch child1CommonCBTimes { + case 0: + if info.Name != "child1" { + t.Fatal("child1 common callback 0 error") + } + case 1: + if info.Name != "child1-1" { + t.Fatal("child1 common callback 1 error") + } + default: + t.Fatal("child1 common callback too many") + } + child1CommonCBTimes++ + return ctx + }).Build() + child2CB := callbacks.NewHandlerBuilder().OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Name != "child2-1" { + t.Fatal("child2-1 common callback 0 error") + } + return ctx + }).Build() + + result, err := r.Invoke(ctx, "input", + WithCallbacks(grandCommonCB), + WithCallbacks(parentCommonCB).DesignateNodeWithPath(NewNodePath("2")), + WithCallbacks(grand1CB).DesignateNode("1"), + WithCallbacks(child1CommonCB).DesignateNodeWithPath(NewNodePath("2", "2")), + WithCallbacks(child2CB).DesignateNodeWithPath(NewNodePath("2", "3", "1")), + WithLambdaOption(grandparentOption("grandparent-1")).DesignateNodeWithPath(NewNodePath("1")), + WithLambdaOption(parentOption("parent-1")).DesignateNodeWithPath(NewNodePath("2", "1")), + WithLambdaOption(child1Option("child1-1")).DesignateNodeWithPath(NewNodePath("2", "2", "1")), + WithLambdaOption(child2Option("child2-1")).DesignateNodeWithPath(NewNodePath("2", "3", "1")), + ) + assert.NoError(t, err) + assert.Equal(t, result, "input grandparent-1 parent-1 child1-1 child2-1") +} diff --git a/compose/graph_compile_options.go b/compose/graph_compile_options.go new file mode 100644 index 0000000..2e0476f --- /dev/null +++ b/compose/graph_compile_options.go @@ -0,0 +1,133 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +type graphCompileOptions struct { + maxRunSteps int + graphName string + nodeTriggerMode NodeTriggerMode // default to AnyPredecessor (pregel) + + callbacks []GraphCompileCallback + + origOpts []GraphCompileOption + + checkPointStore CheckPointStore + serializer Serializer + interruptBeforeNodes []string + interruptAfterNodes []string + + eagerDisabled bool + + mergeConfigs map[string]FanInMergeConfig +} + +func newGraphCompileOptions(opts ...GraphCompileOption) *graphCompileOptions { + option := &graphCompileOptions{} + + for _, o := range opts { + o(option) + } + + option.origOpts = opts + + return option +} + +// GraphCompileOption options for compiling AnyGraph. +type GraphCompileOption func(*graphCompileOptions) + +// WithMaxRunSteps sets the maximum number of steps that a graph can run. +// This is useful to prevent infinite loops in graphs with cycles. +// If the number of steps exceeds maxSteps, the graph execution will be terminated with an error. +func WithMaxRunSteps(maxSteps int) GraphCompileOption { + return func(o *graphCompileOptions) { + o.maxRunSteps = maxSteps + } +} + +// WithGraphName sets a name for the graph. +// The name is used for debugging and logging purposes. +// If not set, a default name will be used. +func WithGraphName(graphName string) GraphCompileOption { + return func(o *graphCompileOptions) { + o.graphName = graphName + } +} + +// WithEagerExecution enables the eager execution mode for the graph. +// In eager mode, nodes will be executed immediately once they are ready to run, +// without waiting for the completion of a super step, ref: https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine +// Note: Eager mode is not allowed when the graph's trigger mode is set to AnyPredecessor. +// Workflow uses eager mode by default. +// Deprecated: Eager execution is automatically enabled by default when a node's trigger mode is set to AllPredecessor. +// If you were using this option previously, it can be safely removed without changing behavior. +func WithEagerExecution() GraphCompileOption { + return func(o *graphCompileOptions) { + return + } +} + +// WithEagerExecutionDisabled disables the eager execution mode for the graph. +// By default, eager execution is enabled for Workflow and Graph with the AllPredecessor trigger mode. +// After using this option, nodes will wait for the completion of a super step instead of execute immediately once they are ready to run. +// ref: https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine +func WithEagerExecutionDisabled() GraphCompileOption { + return func(o *graphCompileOptions) { + o.eagerDisabled = true + } +} + +// WithNodeTriggerMode sets the trigger mode for nodes in the graph. +// The trigger mode determines when a node is triggered during graph execution, ref: https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine +// AnyPredecessor by default. +func WithNodeTriggerMode(triggerMode NodeTriggerMode) GraphCompileOption { + return func(o *graphCompileOptions) { + o.nodeTriggerMode = triggerMode + } +} + +// WithGraphCompileCallbacks sets callbacks for graph compilation. +func WithGraphCompileCallbacks(cbs ...GraphCompileCallback) GraphCompileOption { + return func(o *graphCompileOptions) { + o.callbacks = append(o.callbacks, cbs...) + } +} + +// FanInMergeConfig defines the configuration for fan-in merge operations. +// It allows specifying how multiple inputs are merged into a single input. +// StreamMergeWithSourceEOF indicates whether to emit a SourceEOF error for each stream +// when it ends, before the final merged output is produced. This is useful for +// tracking the completion of individual input streams in a named stream merge. +type FanInMergeConfig struct { + StreamMergeWithSourceEOF bool //indicates whether to emit a SourceEOF error for each stream +} + +// WithFanInMergeConfig sets the fan-in merge configurations +// for the graph nodes that receive inputs from multiple sources. +func WithFanInMergeConfig(confs map[string]FanInMergeConfig) GraphCompileOption { + return func(o *graphCompileOptions) { + o.mergeConfigs = confs + } +} + +// InitGraphCompileCallbacks set global graph compile callbacks, +// which ONLY will be added to top level graph compile options +func InitGraphCompileCallbacks(cbs []GraphCompileCallback) { + globalGraphCompileCallbacks = cbs +} + +var globalGraphCompileCallbacks []GraphCompileCallback diff --git a/compose/graph_manager.go b/compose/graph_manager.go new file mode 100644 index 0000000..46df348 --- /dev/null +++ b/compose/graph_manager.go @@ -0,0 +1,556 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "runtime/debug" + "time" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/internal/safe" +) + +type channel interface { + reportValues(map[string]any) error + reportDependencies([]string) + reportSkip([]string) bool + get(bool, string, *edgeHandlerManager) (any, bool, error) + convertValues(fn func(map[string]any) error) error + load(channel) error + + setMergeConfig(FanInMergeConfig) +} + +type edgeHandlerManager struct { + h map[string]map[string][]handlerPair +} + +func (e *edgeHandlerManager) handle(from, to string, value any, isStream bool) (any, error) { + if _, ok := e.h[from]; !ok { + return value, nil + } + if _, ok := e.h[from][to]; !ok { + return value, nil + } + if isStream { + for _, v := range e.h[from][to] { + value = v.transform(value.(streamReader)) + } + } else { + for _, v := range e.h[from][to] { + var err error + value, err = v.invoke(value) + if err != nil { + return nil, err + } + } + } + return value, nil +} + +type preNodeHandlerManager struct { + h map[string][]handlerPair +} + +func (p *preNodeHandlerManager) handle(nodeKey string, value any, isStream bool) (any, error) { + if _, ok := p.h[nodeKey]; !ok { + return value, nil + } + if isStream { + for _, v := range p.h[nodeKey] { + value = v.transform(value.(streamReader)) + } + } else { + for _, v := range p.h[nodeKey] { + var err error + value, err = v.invoke(value) + if err != nil { + return nil, err + } + } + } + return value, nil +} + +type preBranchHandlerManager struct { + h map[string][][]handlerPair +} + +func (p *preBranchHandlerManager) handle(nodeKey string, idx int, value any, isStream bool) (any, error) { + if _, ok := p.h[nodeKey]; !ok { + return value, nil + } + if isStream { + for _, v := range p.h[nodeKey][idx] { + value = v.transform(value.(streamReader)) + } + } else { + for _, v := range p.h[nodeKey][idx] { + var err error + value, err = v.invoke(value) + if err != nil { + return nil, err + } + } + } + return value, nil +} + +type channelManager struct { + isStream bool + channels map[string]channel + + successors map[string][]string + dataPredecessors map[string]map[string]struct{} + controlPredecessors map[string]map[string]struct{} + + edgeHandlerManager *edgeHandlerManager + preNodeHandlerManager *preNodeHandlerManager +} + +func (c *channelManager) loadChannels(channels map[string]channel) error { + for key, ch := range c.channels { + if nCh, ok := channels[key]; ok { + if err := ch.load(nCh); err != nil { + return fmt.Errorf("load channel[%s] fail: %w", key, err) + } + } + } + return nil +} + +func (c *channelManager) updateValues(_ context.Context, values map[string] /*to*/ map[string] /*from*/ any) error { + for target, fromMap := range values { + toChannel, ok := c.channels[target] + if !ok { + return fmt.Errorf("target channel doesn't existed: %s", target) + } + dps, ok := c.dataPredecessors[target] + if !ok { + dps = map[string]struct{}{} + } + nFromMap := make(map[string]any, len(fromMap)) + for from, value := range fromMap { + if _, ok = dps[from]; ok { + nFromMap[from] = fromMap[from] + } else { + if sr, okk := value.(streamReader); okk { + sr.close() + } + } + } + + err := toChannel.reportValues(nFromMap) + if err != nil { + return fmt.Errorf("update target channel[%s] fail: %w", target, err) + } + } + return nil +} + +func (c *channelManager) updateDependencies(_ context.Context, dependenciesMap map[string][]string) error { + for target, dependencies := range dependenciesMap { + toChannel, ok := c.channels[target] + if !ok { + return fmt.Errorf("target channel doesn't existed: %s", target) + } + cps, ok := c.controlPredecessors[target] + if !ok { + cps = map[string]struct{}{} + } + var deps []string + for _, from := range dependencies { + if _, ok = cps[from]; ok { + deps = append(deps, from) + } + } + + toChannel.reportDependencies(deps) + } + return nil +} + +func (c *channelManager) getFromReadyChannels(_ context.Context) (map[string]any, error) { + result := make(map[string]any) + for target, ch := range c.channels { + v, ready, err := ch.get(c.isStream, target, c.edgeHandlerManager) + if err != nil { + return nil, fmt.Errorf("get value from ready channel[%s] fail: %w", target, err) + } + if ready { + v, err = c.preNodeHandlerManager.handle(target, v, c.isStream) + if err != nil { + return nil, err + } + result[target] = v + } + } + return result, nil +} + +func (c *channelManager) updateAndGet(ctx context.Context, values map[string]map[string]any, dependencies map[string][]string) (map[string]any, error) { + err := c.updateValues(ctx, values) + if err != nil { + return nil, fmt.Errorf("update channel fail: %w", err) + } + err = c.updateDependencies(ctx, dependencies) + if err != nil { + return nil, fmt.Errorf("update channel fail: %w", err) + } + return c.getFromReadyChannels(ctx) +} + +func (c *channelManager) reportBranch(from string, skippedNodes []string) error { + var nKeys []string + for _, node := range skippedNodes { + skipped := c.channels[node].reportSkip([]string{from}) + if skipped { + nKeys = append(nKeys, node) + } + } + + for i := 0; i < len(nKeys); i++ { + key := nKeys[i] + + if key == END { + continue + } + if _, ok := c.successors[key]; !ok { + return fmt.Errorf("unknown node: %s", key) + } + for _, successor := range c.successors[key] { + skipped := c.channels[successor].reportSkip([]string{key}) + if skipped { + nKeys = appendIfNotExist(nKeys, successor) + } + // todo: detect if end node has been skipped? + } + } + return nil +} + +func appendIfNotExist(s []string, elem string) []string { + for _, i := range s { + if i == elem { + return s + } + } + return append(s, elem) +} + +type task struct { + ctx context.Context + nodeKey string + call *chanCall + input any + originalInput any + output any + option []any + err error + skipPreHandler bool +} + +type taskManager struct { + runWrapper runnableCallWrapper + opts []Option + needAll bool + + num uint32 + done *internal.UnboundedChan[*task] + runningTasks map[string]*task + + cancelCh chan *time.Duration + canceled bool + deadline *time.Time + + persistRerunInput bool +} + +func (t *taskManager) execute(currentTask *task) { + defer func() { + panicInfo := recover() + if panicInfo != nil { + currentTask.output = nil + currentTask.err = safe.NewPanicErr(panicInfo, debug.Stack()) + } + + t.done.Send(currentTask) + }() + + ctx := initNodeCallbacks(currentTask.ctx, currentTask.nodeKey, currentTask.call.action.nodeInfo, currentTask.call.action.meta, t.opts...) + currentTask.output, currentTask.err = t.runWrapper(ctx, currentTask.call.action, currentTask.input, currentTask.option...) +} + +func (t *taskManager) submit(tasks []*task) error { + if len(tasks) == 0 { + return nil + } + + // synchronously execute one task, if there are no other tasks in the task pool and meet one of the following conditions： + // 1. the new task is the only one + // 2. the task manager mode is set to needAll + for i := 0; i < len(tasks); i++ { + currentTask := tasks[i] + + if t.persistRerunInput { + if sr, ok := currentTask.input.(streamReader); ok { + copies := sr.copy(2) + currentTask.originalInput, currentTask.input = copies[0], copies[1] + } else { + currentTask.originalInput = currentTask.input + } + } + + err := runPreHandler(currentTask, t.runWrapper) + if err != nil { + // pre-handler error, regarded as a failure of the task itself + currentTask.err = err + tasks = append(tasks[:i], tasks[i+1:]...) + i-- + t.num++ + t.done.Send(currentTask) + } + + t.runningTasks[currentTask.nodeKey] = currentTask + } + if len(tasks) == 0 { + // all tasks' pre-handler failed + return nil + } + + var syncTask *task + if t.num == 0 && (len(tasks) == 1 || t.needAll) && t.cancelCh == nil /*if graph can be interrupted by user, shouldn't sync run task*/ { + syncTask = tasks[0] + tasks = tasks[1:] + } + for _, currentTask := range tasks { + t.num += 1 + go t.execute(currentTask) + } + if syncTask != nil { + t.num += 1 + t.execute(syncTask) + } + return nil +} + +func (t *taskManager) wait() (tasks []*task, canceled bool, canceledTasks []*task) { + if t.needAll { + tasks, canceledTasks = t.waitAll() + return tasks, t.canceled, canceledTasks + } + + ta, success, canceled := t.waitOne() + if canceled { + // has canceled and timeout, return canceled tasks + for _, rta := range t.runningTasks { + canceledTasks = append(canceledTasks, rta) + } + t.runningTasks = make(map[string]*task) + t.num = 0 + return nil, true, canceledTasks + } + if t.canceled { + // has canceled, but not timeout, wait all + tasks, canceledTasks = t.waitAll() + return append(tasks, ta), true, canceledTasks + } + if !success { + return []*task{}, t.canceled, nil + } + + return []*task{ta}, t.canceled, nil +} + +func (t *taskManager) waitOne() (ta *task, success bool, canceled bool) { + if t.num == 0 { + return nil, false, false + } + + if t.cancelCh == nil { + ta, _ = t.done.Receive() + } else { + ta, _, canceled = t.receive(t.done.Receive) + } + + t.num-- + + if canceled { + return nil, false, true + } + + delete(t.runningTasks, ta.nodeKey) + + if ta.originalInput != nil && (ta.err == nil || !isInterruptError(ta.err)) { + if sr, ok := ta.originalInput.(streamReader); ok { + sr.close() + } + ta.originalInput = nil + } + + if ta.err != nil { + // biz error, jump post processor + return ta, true, false + } + runPostHandler(ta, t.runWrapper) + return ta, true, false +} + +func (t *taskManager) waitAll() (successTasks []*task, canceledTasks []*task) { + result := make([]*task, 0, t.num) + for { + ta, success, canceled := t.waitOne() + if canceled { + for _, rt := range t.runningTasks { + canceledTasks = append(canceledTasks, rt) + } + t.runningTasks = make(map[string]*task) + t.num = 0 + return result, canceledTasks + } + if !success { + return result, nil + } + result = append(result, ta) + } +} + +func (t *taskManager) receive(recv func() (*task, bool)) (ta *task, closed bool, canceled bool) { + if t.deadline != nil { + // have canceled, receive in a certain time + return receiveWithDeadline(recv, *t.deadline) + } + if t.canceled { + // canceled without timeout + ta, closed = recv() + return ta, closed, false + } + if t.cancelCh != nil { + // have not canceled, receive while listening + ta, closed, canceled, t.canceled, t.deadline = receiveWithListening(recv, t.cancelCh) + return ta, closed, canceled + } + // won't cancel + ta, closed = recv() + return ta, closed, false +} + +func receiveWithDeadline(recv func() (*task, bool), deadline time.Time) (ta *task, closed bool, canceled bool) { + now := time.Now() + if deadline.Before(now) { + return nil, false, true + } + + timeout := deadline.Sub(now) + + resultCh := make(chan struct{}, 1) + + go func() { + ta, closed = recv() + resultCh <- struct{}{} + }() + + timeoutCh := time.After(timeout) + + select { + case <-resultCh: + return ta, closed, false + case <-timeoutCh: + return nil, false, true + } +} + +func receiveWithListening(recv func() (*task, bool), cancel chan *time.Duration) (*task, bool, bool, bool, *time.Time) { + type pair struct { + ta *task + closed bool + } + resultCh := make(chan pair, 1) + var timeoutCh <-chan time.Time + + var deadline *time.Time + canceled := false + go func() { + ta, closed := recv() + resultCh <- pair{ta, closed} + }() + + select { + case p := <-resultCh: + return p.ta, p.closed, false, false, nil + case timeout, ok := <-cancel: + if !ok { + // The cancel channel has been closed — this means a previous call to + // receiveWithListening already consumed the cancel signal (task completed + // at the same time as cancel, and select picked the task result). Since + // cancel was already issued, treat this as an immediate cancel rather than + // blocking forever on resultCh. + return nil, false, true, true, nil + } + canceled = true + if timeout == nil { + break + } + timeoutCh = time.After(*timeout) + dt := time.Now().Add(*timeout) + deadline = &dt + } + + if timeoutCh != nil { + select { + case p := <-resultCh: + return p.ta, p.closed, false, canceled, deadline + case <-timeoutCh: + return nil, false, true, canceled, deadline + } + } + p := <-resultCh + return p.ta, p.closed, false, canceled, nil +} + +func runPreHandler(ta *task, runWrapper runnableCallWrapper) (err error) { + defer func() { + if e := recover(); e != nil { + err = safe.NewPanicErr(fmt.Errorf("panic in pre handler: %v", e), debug.Stack()) + } + }() + if ta.call.preProcessor != nil && !ta.skipPreHandler { + nInput, err := runWrapper(ta.ctx, ta.call.preProcessor, ta.input, ta.option...) + if err != nil { + return fmt.Errorf("run node[%s] pre processor fail: %w", ta.nodeKey, err) + } + ta.input = nInput + } + return nil +} + +func runPostHandler(ta *task, runWrapper runnableCallWrapper) { + defer func() { + if e := recover(); e != nil { + ta.err = safe.NewPanicErr(fmt.Errorf("panic in post handler: %v", e), debug.Stack()) + } + }() + if ta.call.postProcessor != nil { + nOutput, err := runWrapper(ta.ctx, ta.call.postProcessor, ta.output, ta.option...) + if err != nil { + ta.err = fmt.Errorf("run node[%s] post processor fail: %w", ta.nodeKey, err) + } + ta.output = nOutput + } +} diff --git a/compose/graph_node.go b/compose/graph_node.go new file mode 100644 index 0000000..47091eb --- /dev/null +++ b/compose/graph_node.go @@ -0,0 +1,177 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "reflect" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/internal/generic" +) + +// the info of most original executable object directly provided by the user +type executorMeta struct { + + // automatically identified based on the way of addNode + component component + + // indicates whether the executable object user provided could execute the callback aspect itself. + // if it could, the callback in the corresponding graph node won't be executed + // for components, the value comes from callbacks.Checker + isComponentCallbackEnabled bool + + // for components, the value comes from components.Typer + // for lambda, the value comes from the user's explicit config + // if componentImplType is empty, then the class name or func name in the instance will be inferred, but no guarantee. + componentImplType string +} + +type nodeInfo struct { + + // the name of graph node for display purposes, not unique. + // passed from WithNodeName() + name string + + inputKey string + outputKey string + + preProcessor, postProcessor *composableRunnable + + compileOption *graphCompileOptions // if the node is an AnyGraph, it will need compile options of its own +} + +// graphNode the complete information of the node in graph +type graphNode struct { + cr *composableRunnable + + g AnyGraph + + nodeInfo *nodeInfo + executorMeta *executorMeta + + instance any + opts []GraphAddNodeOpt +} + +func (gn *graphNode) getGenericHelper() *genericHelper { + var ret *genericHelper + if gn.g != nil { + ret = gn.g.getGenericHelper() + } else if gn.cr != nil { + ret = gn.cr.genericHelper + } else { + return nil + } + + if gn.nodeInfo != nil { + if len(gn.nodeInfo.inputKey) > 0 { + ret = ret.forMapInput() + } + if len(gn.nodeInfo.outputKey) > 0 { + ret = ret.forMapOutput() + } + } + return ret +} + +func (gn *graphNode) inputType() reflect.Type { + if gn.nodeInfo != nil && len(gn.nodeInfo.inputKey) != 0 { + return generic.TypeOf[map[string]any]() + } + // priority follow compile + if gn.g != nil { + return gn.g.inputType() + } else if gn.cr != nil { + return gn.cr.inputType + } + + return nil +} + +func (gn *graphNode) outputType() reflect.Type { + if gn.nodeInfo != nil && len(gn.nodeInfo.outputKey) != 0 { + return generic.TypeOf[map[string]any]() + } + // priority follow compile + if gn.g != nil { + return gn.g.outputType() + } else if gn.cr != nil { + return gn.cr.outputType + } + + return nil +} + +func (gn *graphNode) compileIfNeeded(ctx context.Context) (*composableRunnable, error) { + var r *composableRunnable + if gn.g != nil { + cr, err := gn.g.compile(ctx, gn.nodeInfo.compileOption) + if err != nil { + return nil, err + } + + r = cr + gn.cr = cr + } else if gn.cr != nil { + r = gn.cr + } else { + return nil, errors.New("no graph or component provided") + } + + r.meta = gn.executorMeta + r.nodeInfo = gn.nodeInfo + + if gn.nodeInfo.outputKey != "" { + r = outputKeyedComposableRunnable(gn.nodeInfo.outputKey, r) + } + + if gn.nodeInfo.inputKey != "" { + r = inputKeyedComposableRunnable(gn.nodeInfo.inputKey, r) + } + + return r, nil +} + +func parseExecutorInfoFromComponent(c component, executor any) *executorMeta { + + componentImplType, ok := components.GetType(executor) + if !ok { + componentImplType = generic.ParseTypeName(reflect.ValueOf(executor)) + } + + return &executorMeta{ + component: c, + isComponentCallbackEnabled: components.IsCallbacksEnabled(executor), + componentImplType: componentImplType, + } +} + +func getNodeInfo(opts ...GraphAddNodeOpt) (*nodeInfo, *graphAddNodeOpts) { + + opt := getGraphAddNodeOpts(opts...) + + return &nodeInfo{ + name: opt.nodeOptions.nodeName, + inputKey: opt.nodeOptions.inputKey, + outputKey: opt.nodeOptions.outputKey, + preProcessor: opt.processor.statePreHandler, + postProcessor: opt.processor.statePostHandler, + compileOption: newGraphCompileOptions(opt.nodeOptions.graphCompileOption...), + }, opt +} diff --git a/compose/graph_run.go b/compose/graph_run.go new file mode 100644 index 0000000..02b4fca --- /dev/null +++ b/compose/graph_run.go @@ -0,0 +1,1055 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + "reflect" + "strings" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/internal/serialization" +) + +type chanCall struct { + action *composableRunnable + writeTo []string + writeToBranches []*GraphBranch + + controls []string // branch must control + + preProcessor, postProcessor *composableRunnable +} + +type chanBuilder func(dependencies []string, indirectDependencies []string, zeroValue func() any, emptyStream func() streamReader) channel + +type runner struct { + chanSubscribeTo map[string]*chanCall + + successors map[string][]string + dataPredecessors map[string][]string + controlPredecessors map[string][]string + + inputChannels *chanCall + + chanBuilder chanBuilder // could be nil + eager bool + dag bool + + runCtx func(ctx context.Context) context.Context + + options graphCompileOptions + + inputType reflect.Type + outputType reflect.Type + + // take effect as a subgraph through toComposableRunnable + inputStreamFilter streamMapFilter + inputConverter handlerPair + inputFieldMappingConverter handlerPair + inputConvertStreamPair, outputConvertStreamPair streamConvertPair + + *genericHelper + + // checks need to do because cannot check at compile + runtimeCheckEdges map[string]map[string]bool + runtimeCheckBranches map[string][]bool + + edgeHandlerManager *edgeHandlerManager + preNodeHandlerManager *preNodeHandlerManager + preBranchHandlerManager *preBranchHandlerManager + + checkPointer *checkPointer + interruptBeforeNodes []string + interruptAfterNodes []string + + mergeConfigs map[string]FanInMergeConfig +} + +func (r *runner) invoke(ctx context.Context, input any, opts ...Option) (any, error) { + return r.run(ctx, false, input, opts...) +} + +func (r *runner) transform(ctx context.Context, input streamReader, opts ...Option) (streamReader, error) { + s, err := r.run(ctx, true, input, opts...) + if err != nil { + return nil, err + } + + return s.(streamReader), nil +} + +type runnableCallWrapper func(context.Context, *composableRunnable, any, ...any) (any, error) + +func runnableInvoke(ctx context.Context, r *composableRunnable, input any, opts ...any) (any, error) { + return r.i(ctx, input, opts...) +} + +func runnableTransform(ctx context.Context, r *composableRunnable, input any, opts ...any) (any, error) { + return r.t(ctx, input.(streamReader), opts...) +} + +func (r *runner) run(ctx context.Context, isStream bool, input any, opts ...Option) (result any, err error) { + haveOnStart := false // delay triggering onGraphStart until state initialization is complete, so that the state can be accessed within onGraphStart. + defer func() { + if !haveOnStart { + ctx, input = onGraphStart(ctx, input, isStream) + } + if err != nil { + ctx, err = onGraphError(ctx, err) + } else { + ctx, result = onGraphEnd(ctx, result, isStream) + } + }() + + var runWrapper runnableCallWrapper + runWrapper = runnableInvoke + if isStream { + runWrapper = runnableTransform + } + + // Initialize channel and task managers. + cm := r.initChannelManager(isStream) + tm := r.initTaskManager(runWrapper, getGraphCancel(ctx), opts...) + maxSteps := r.options.maxRunSteps + + maxSteps, err = r.resolveMaxSteps(maxSteps, opts) + if err != nil { + return nil, err + } + + // Extract and validate options for each node. + optMap, extractErr := extractOption(r.chanSubscribeTo, opts...) + if extractErr != nil { + return nil, newGraphRunError(fmt.Errorf("graph extract option fail: %w", extractErr)) + } + + // Extract CheckPointID + checkPointID, writeToCheckPointID, stateModifier, forceNewRun := getCheckPointInfo(opts...) + if checkPointID != nil && r.checkPointer.store == nil { + return nil, newGraphRunError(fmt.Errorf("receive checkpoint id but have not set checkpoint store")) + } + + // Extract subgraph + path, isSubGraph := getNodePath(ctx) + + // load checkpoint from ctx/store or init graph + initialized := false + var nextTasks []*task + if cp := getCheckPointFromCtx(ctx); cp != nil { + // in subgraph, try to load checkpoint from ctx + initialized = true + + ctx, err = r.restoreCheckPointState(ctx, *path, getStateModifier(ctx), cp, isStream, cm) + if err != nil { + return nil, err + } + + ctx, input = onGraphStart(ctx, input, isStream) + haveOnStart = true + + nextTasks, err = r.restoreTasks(ctx, cp.Inputs, cp.SkipPreHandler, cp.RerunNodes, isStream, optMap) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("restore tasks fail: %w", err)) + } + } else if checkPointID != nil && !forceNewRun { + cp, err = getCheckPointFromStore(ctx, *checkPointID, r.checkPointer) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("load checkpoint from store fail: %w", err)) + } + if cp != nil { + // load checkpoint from store + initialized = true + + ctx = setStateModifier(ctx, stateModifier) + ctx = setCheckPointToCtx(ctx, cp) + + ctx, err = r.restoreCheckPointState(ctx, *NewNodePath(), stateModifier, cp, isStream, cm) + if err != nil { + return nil, err + } + + ctx, input = onGraphStart(ctx, input, isStream) + haveOnStart = true + + nextTasks, err = r.restoreTasks(ctx, cp.Inputs, cp.SkipPreHandler, cp.RerunNodes, isStream, optMap) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("restore tasks fail: %w", err)) + } + } + } + if !initialized { + // have not inited from checkpoint + if r.runCtx != nil { + ctx = r.runCtx(ctx) + } + + ctx, input = onGraphStart(ctx, input, isStream) + haveOnStart = true + + var isEnd bool + nextTasks, result, isEnd, err = r.calculateNextTasks(ctx, []*task{{ + nodeKey: START, + call: r.inputChannels, + output: input, + }}, isStream, cm, optMap) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("calculate next tasks fail: %w", err)) + } + if isEnd { + return result, nil + } + if len(nextTasks) == 0 { + return nil, newGraphRunError(fmt.Errorf("no tasks to execute after graph start")) + } + + if keys := getHitKey(nextTasks, r.interruptBeforeNodes); len(keys) > 0 { + tempInfo := newInterruptTempInfo() + tempInfo.interruptBeforeNodes = append(tempInfo.interruptBeforeNodes, keys...) + return nil, r.handleInterrupt(ctx, + tempInfo, + nextTasks, + cm.channels, + isStream, + isSubGraph, + writeToCheckPointID, + ) + } + } + + // used to reporting NoTask error + var lastCompletedTask []*task + + // Main execution loop. + for step := 0; ; step++ { + // Check for context cancellation. + select { + case <-ctx.Done(): + _, _ = tm.waitAll() + return nil, newGraphRunError(fmt.Errorf("context has been canceled: %w", ctx.Err())) + default: + } + if !r.dag && step >= maxSteps { + return nil, newGraphRunError(ErrExceedMaxSteps) + } + + // 1. submit next tasks + // 2. get completed tasks + // 3. calculate next tasks + + err = tm.submit(nextTasks) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("failed to submit tasks: %w", err)) + } + + var totalCanceledTasks []*task + + completedTasks, canceled, canceledTasks := tm.wait() + totalCanceledTasks = append(totalCanceledTasks, canceledTasks...) + tempInfo := newInterruptTempInfo() + tempInfo.collectCanceledInfo(canceled, canceledTasks, completedTasks) + + err = r.resolveInterruptCompletedTasks(tempInfo, completedTasks) + if err != nil { + return nil, err // err has been wrapped + } + + if len(tempInfo.subGraphInterrupts)+len(tempInfo.interruptRerunNodes) > 0 { + var newCompletedTasks []*task + newCompletedTasks, canceledTasks = tm.waitAll() + totalCanceledTasks = append(totalCanceledTasks, canceledTasks...) + for _, ct := range canceledTasks { + // handle timeout tasks as rerun + tempInfo.interruptRerunNodes = append(tempInfo.interruptRerunNodes, ct.nodeKey) + } + + err = r.resolveInterruptCompletedTasks(tempInfo, newCompletedTasks) + if err != nil { + return nil, err // err has been wrapped + } + + // subgraph has interrupted + // save other completed tasks to channel + // save interrupted subgraph as next task with SkipPreHandler + // report current graph interrupt info + return nil, r.handleInterruptWithSubGraphAndRerunNodes( + ctx, + tempInfo, + append(append(completedTasks, newCompletedTasks...), totalCanceledTasks...), // canceled tasks are handled as rerun + writeToCheckPointID, + isSubGraph, + cm, + isStream, + ) + } + + if len(completedTasks) == 0 { + return nil, newGraphRunError(fmt.Errorf("no tasks to execute, last completed nodes: %v", printTask(lastCompletedTask))) + } + lastCompletedTask = completedTasks + + var isEnd bool + nextTasks, result, isEnd, err = r.calculateNextTasks(ctx, completedTasks, isStream, cm, optMap) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("failed to calculate next tasks: %w", err)) + } + if isEnd { + return result, nil + } + + tempInfo.interruptBeforeNodes = getHitKey(nextTasks, r.interruptBeforeNodes) + + if len(tempInfo.interruptBeforeNodes) > 0 || len(tempInfo.interruptAfterNodes) > 0 { + var newCompletedTasks []*task + newCompletedTasks, canceledTasks = tm.waitAll() + totalCanceledTasks = append(totalCanceledTasks, canceledTasks...) + for _, ct := range canceledTasks { + tempInfo.interruptRerunNodes = append(tempInfo.interruptRerunNodes, ct.nodeKey) + } + + err = r.resolveInterruptCompletedTasks(tempInfo, newCompletedTasks) + if err != nil { + return nil, err // err has been wrapped + } + + if len(tempInfo.subGraphInterrupts)+len(tempInfo.interruptRerunNodes) > 0 { + return nil, r.handleInterruptWithSubGraphAndRerunNodes( + ctx, + tempInfo, + append(append(completedTasks, newCompletedTasks...), totalCanceledTasks...), + writeToCheckPointID, + isSubGraph, + cm, + isStream, + ) + } + + var newNextTasks []*task + newNextTasks, result, isEnd, err = r.calculateNextTasks(ctx, newCompletedTasks, isStream, cm, optMap) + if err != nil { + return nil, newGraphRunError(fmt.Errorf("failed to calculate next tasks: %w", err)) + } + + if isEnd { + return result, nil + } + + tempInfo.interruptBeforeNodes = append(tempInfo.interruptBeforeNodes, getHitKey(newNextTasks, r.interruptBeforeNodes)...) + + // simple interrupt + return nil, r.handleInterrupt(ctx, tempInfo, append(nextTasks, newNextTasks...), cm.channels, isStream, isSubGraph, writeToCheckPointID) + } + } +} + +func (r *runner) resolveMaxSteps(maxSteps int, opts []Option) (int, error) { + if r.dag { + for i := range opts { + if opts[i].maxRunSteps > 0 { + return 0, newGraphRunError(fmt.Errorf("cannot set max run steps in dag")) + } + } + return maxSteps, nil + } + for i := range opts { + if opts[i].maxRunSteps > 0 { + maxSteps = opts[i].maxRunSteps + } + } + if maxSteps < 1 { + return 0, newGraphRunError(errors.New("max run steps limit must be at least 1")) + } + return maxSteps, nil +} + +func (r *runner) restoreCheckPointState( + ctx context.Context, + path NodePath, + sm StateModifier, + cp *checkpoint, + isStream bool, + cm *channelManager, +) (context.Context, error) { + err := r.checkPointer.restoreCheckPoint(cp, isStream) + if err != nil { + return ctx, newGraphRunError(fmt.Errorf("restore checkpoint fail: %w", err)) + } + + err = cm.loadChannels(cp.Channels) + if err != nil { + return ctx, newGraphRunError(err) + } + if sm != nil && cp.State != nil { + err = sm(ctx, path, cp.State) + if err != nil { + return ctx, newGraphRunError(fmt.Errorf("state modifier fail: %w", err)) + } + } + if cp.State != nil { + isResumeTarget, hasData, data := GetResumeContext[any](ctx) + if isResumeTarget && hasData { + cp.State = data + } + + var parent *internalState + if prev := ctx.Value(stateKey{}); prev != nil { + if p, ok := prev.(*internalState); ok { + parent = p + } + } + + ctx = context.WithValue(ctx, stateKey{}, &internalState{state: cp.State, parent: parent}) + } + + return ctx, nil +} + +func newInterruptTempInfo() *interruptTempInfo { + return &interruptTempInfo{ + subGraphInterrupts: map[string]*subGraphInterruptError{}, + interruptRerunExtra: map[string]any{}, + } +} + +type interruptTempInfo struct { + subGraphInterrupts map[string]*subGraphInterruptError + interruptRerunNodes []string + interruptBeforeNodes []string + interruptAfterNodes []string + interruptRerunExtra map[string]any + + signals []*core.InterruptSignal +} + +func (ti *interruptTempInfo) collectCanceledInfo(canceled bool, canceledTasks, completedTasks []*task) { + if !canceled { + return + } + + if len(canceledTasks) > 0 { + for _, t := range canceledTasks { + ti.interruptRerunNodes = append(ti.interruptRerunNodes, t.nodeKey) + } + } else { + for _, t := range completedTasks { + ti.interruptAfterNodes = append(ti.interruptAfterNodes, t.nodeKey) + } + } +} + +func (r *runner) resolveInterruptCompletedTasks(tempInfo *interruptTempInfo, completedTasks []*task) (err error) { + for _, completedTask := range completedTasks { + if completedTask.err != nil { + if info := isSubGraphInterrupt(completedTask.err); info != nil { + tempInfo.subGraphInterrupts[completedTask.nodeKey] = info + tempInfo.signals = append(tempInfo.signals, info.signal) + continue + } + + ire := &core.InterruptSignal{} + if errors.As(completedTask.err, &ire) { + tempInfo.interruptRerunNodes = append(tempInfo.interruptRerunNodes, completedTask.nodeKey) + if ire.Info != nil { + tempInfo.interruptRerunExtra[completedTask.nodeKey] = ire.InterruptInfo.Info + } + + tempInfo.signals = append(tempInfo.signals, ire) + continue + } + + return wrapGraphNodeError(completedTask.nodeKey, completedTask.err) + } + + for _, key := range r.interruptAfterNodes { + if key == completedTask.nodeKey { + tempInfo.interruptAfterNodes = append(tempInfo.interruptAfterNodes, key) + break + } + } + } + return nil +} + +func getHitKey(tasks []*task, keys []string) []string { + var ret []string + for _, t := range tasks { + for _, key := range keys { + if key == t.nodeKey { + ret = append(ret, t.nodeKey) + } + } + } + return ret +} + +func (r *runner) handleInterrupt( + ctx context.Context, + tempInfo *interruptTempInfo, + nextTasks []*task, + channels map[string]channel, + isStream bool, + isSubGraph bool, + checkPointID *string, +) error { + cp := &checkpoint{ + Channels: channels, + Inputs: make(map[string]any), + SkipPreHandler: map[string]bool{}, + } + if r.runCtx != nil { + // current graph has enable state + if state, ok := ctx.Value(stateKey{}).(*internalState); ok { + state.mu.Lock() + copiedState, err := deepCopyState(state.state) + state.mu.Unlock() + if err != nil { + return fmt.Errorf("failed to copy state: %w", err) + } + cp.State = copiedState + } + } + + intInfo := &InterruptInfo{ + State: cp.State, + AfterNodes: tempInfo.interruptAfterNodes, + BeforeNodes: tempInfo.interruptBeforeNodes, + RerunNodes: tempInfo.interruptRerunNodes, + RerunNodesExtra: tempInfo.interruptRerunExtra, + SubGraphs: make(map[string]*InterruptInfo), + } + + info := cp.State + + is, err := core.Interrupt(ctx, info, nil, tempInfo.signals) + if err != nil { + return fmt.Errorf("failed to interrupt: %w", err) + } + + cp.InterruptID2Addr, cp.InterruptID2State = core.SignalToPersistenceMaps(is) + + for _, t := range nextTasks { + cp.Inputs[t.nodeKey] = t.input + } + err = r.checkPointer.convertCheckPoint(cp, isStream) + if err != nil { + return fmt.Errorf("failed to convert checkpoint: %w", err) + } + if isSubGraph { + return &subGraphInterruptError{ + Info: intInfo, + CheckPoint: cp, + signal: is, + } + } else if checkPointID != nil { + err := r.checkPointer.set(ctx, *checkPointID, cp) + if err != nil { + return fmt.Errorf("failed to set checkpoint: %w, checkPointID: %s", err, *checkPointID) + } + } + + intInfo.InterruptContexts = core.ToInterruptContexts(is, nil) + return &interruptError{Info: intInfo} +} + +// deepCopyState creates a deep copy of the state using serialization +func deepCopyState(state any) (any, error) { + if state == nil { + return nil, nil + } + serializer := &serialization.InternalSerializer{} + data, err := serializer.Marshal(state) + if err != nil { + return nil, fmt.Errorf("failed to marshal state: %w", err) + } + + // Create new instance of the same type + stateType := reflect.TypeOf(state) + isPtr := stateType.Kind() == reflect.Ptr + if isPtr { + stateType = stateType.Elem() + } + newStatePtr := reflect.New(stateType).Interface() + if err := serializer.Unmarshal(data, newStatePtr); err != nil { + return nil, fmt.Errorf("failed to unmarshal state: %w", err) + } + if isPtr { + return newStatePtr, nil + } + return reflect.ValueOf(newStatePtr).Elem().Interface(), nil +} + +func (r *runner) handleInterruptWithSubGraphAndRerunNodes( + ctx context.Context, + tempInfo *interruptTempInfo, + completeTasks []*task, + checkPointID *string, + isSubGraph bool, + cm *channelManager, + isStream bool, +) error { + var rerunTasks, subgraphTasks, otherTasks []*task + skipPreHandler := map[string]bool{} + for _, t := range completeTasks { + if _, ok := tempInfo.subGraphInterrupts[t.nodeKey]; ok { + subgraphTasks = append(subgraphTasks, t) + skipPreHandler[t.nodeKey] = true // subgraph won't run pre-handler again, but rerun nodes will + continue + } + rerun := false + for _, key := range tempInfo.interruptRerunNodes { + if key == t.nodeKey { + rerunTasks = append(rerunTasks, t) + rerun = true + break + } + } + if !rerun { + otherTasks = append(otherTasks, t) + } + } + + // forward completed tasks + toValue, controls, err := r.resolveCompletedTasks(ctx, otherTasks, isStream, cm) + if err != nil { + return fmt.Errorf("failed to resolve completed tasks in interrupt: %w", err) + } + err = cm.updateValues(ctx, toValue) + if err != nil { + return fmt.Errorf("failed to update values in interrupt: %w", err) + } + err = cm.updateDependencies(ctx, controls) + if err != nil { + return fmt.Errorf("failed to update dependencies in interrupt: %w", err) + } + + cp := &checkpoint{ + Channels: cm.channels, + Inputs: make(map[string]any), + SkipPreHandler: skipPreHandler, + SubGraphs: make(map[string]*checkpoint), + } + if r.runCtx != nil { + // current graph has enable state + if state, ok := ctx.Value(stateKey{}).(*internalState); ok { + state.mu.Lock() + copiedState, err_ := deepCopyState(state.state) + state.mu.Unlock() + if err_ != nil { + return fmt.Errorf("failed to copy state: %w", err_) + } + cp.State = copiedState + } + } + + intInfo := &InterruptInfo{ + State: cp.State, + BeforeNodes: tempInfo.interruptBeforeNodes, + AfterNodes: tempInfo.interruptAfterNodes, + RerunNodes: tempInfo.interruptRerunNodes, + RerunNodesExtra: tempInfo.interruptRerunExtra, + SubGraphs: make(map[string]*InterruptInfo), + } + + info := cp.State + + is, err := core.Interrupt(ctx, info, nil, tempInfo.signals) + if err != nil { + return fmt.Errorf("failed to interrupt: %w", err) + } + + cp.InterruptID2Addr, cp.InterruptID2State = core.SignalToPersistenceMaps(is) + + for _, t := range subgraphTasks { + cp.RerunNodes = append(cp.RerunNodes, t.nodeKey) + cp.SubGraphs[t.nodeKey] = tempInfo.subGraphInterrupts[t.nodeKey].CheckPoint + intInfo.SubGraphs[t.nodeKey] = tempInfo.subGraphInterrupts[t.nodeKey].Info + } + for _, t := range rerunTasks { + cp.RerunNodes = append(cp.RerunNodes, t.nodeKey) + if t.originalInput != nil { + cp.Inputs[t.nodeKey] = t.originalInput + } + } + err = r.checkPointer.convertCheckPoint(cp, isStream) + if err != nil { + return fmt.Errorf("failed to convert checkpoint: %w", err) + } + if isSubGraph { + return &subGraphInterruptError{ + Info: intInfo, + CheckPoint: cp, + signal: is, + } + } else if checkPointID != nil { + err = r.checkPointer.set(ctx, *checkPointID, cp) + if err != nil { + return fmt.Errorf("failed to set checkpoint: %w, checkPointID: %s", err, *checkPointID) + } + } + intInfo.InterruptContexts = core.ToInterruptContexts(is, nil) + return &interruptError{Info: intInfo} +} + +func (r *runner) calculateNextTasks(ctx context.Context, completedTasks []*task, isStream bool, cm *channelManager, optMap map[string][]any) ([]*task, any, bool, error) { + writeChannelValues, controls, err := r.resolveCompletedTasks(ctx, completedTasks, isStream, cm) + if err != nil { + return nil, nil, false, err + } + nodeMap, err := cm.updateAndGet(ctx, writeChannelValues, controls) + if err != nil { + return nil, nil, false, fmt.Errorf("failed to update and get channels: %w", err) + } + var nextTasks []*task + if len(nodeMap) > 0 { + // Check if we've reached the END node. + if v, ok := nodeMap[END]; ok { + return nil, v, true, nil + } + + // Create and submit the next batch of tasks. + nextTasks, err = r.createTasks(ctx, nodeMap, optMap) + if err != nil { + return nil, nil, false, fmt.Errorf("failed to create tasks: %w", err) + } + } + return nextTasks, nil, false, nil +} + +func (r *runner) createTasks(ctx context.Context, nodeMap map[string]any, optMap map[string][]any) ([]*task, error) { + var nextTasks []*task + for nodeKey, nodeInput := range nodeMap { + call, ok := r.chanSubscribeTo[nodeKey] + if !ok { + return nil, fmt.Errorf("node[%s] has not been registered", nodeKey) + } + + if call.action.nodeInfo != nil && call.action.nodeInfo.compileOption != nil { + ctx = forwardCheckPoint(ctx, nodeKey) + } + + nextTasks = append(nextTasks, &task{ + ctx: AppendAddressSegment(ctx, AddressSegmentNode, nodeKey), + nodeKey: nodeKey, + call: call, + input: nodeInput, + option: optMap[nodeKey], + }) + } + return nextTasks, nil +} + +func getCheckPointInfo(opts ...Option) (checkPointID *string, writeToCheckPointID *string, stateModifier StateModifier, forceNewRun bool) { + for _, opt := range opts { + if opt.checkPointID != nil { + checkPointID = opt.checkPointID + } + if opt.writeToCheckPointID != nil { + writeToCheckPointID = opt.writeToCheckPointID + } + if opt.stateModifier != nil { + stateModifier = opt.stateModifier + } + forceNewRun = opt.forceNewRun + } + if writeToCheckPointID == nil { + writeToCheckPointID = checkPointID + } + return +} + +func (r *runner) restoreTasks( + ctx context.Context, + inputs map[string]any, + skipPreHandler map[string]bool, + rerunNodes []string, + isStream bool, + optMap map[string][]any) ([]*task, error) { + ret := make([]*task, 0, len(inputs)) + for _, key := range rerunNodes { + if _, hasInput := inputs[key]; hasInput { + continue + } + + call, ok := r.chanSubscribeTo[key] + if !ok { + return nil, fmt.Errorf("channel[%s] from checkpoint is not registered", key) + } + if isStream { + inputs[key] = call.action.inputEmptyStream() + } else { + inputs[key] = call.action.inputZeroValue() + } + } + for key, input := range inputs { + call, ok := r.chanSubscribeTo[key] + if !ok { + return nil, fmt.Errorf("channel[%s] from checkpoint is not registered", key) + } + + if call.action.nodeInfo != nil && call.action.nodeInfo.compileOption != nil { + // sub graph + ctx = forwardCheckPoint(ctx, key) + } + + newTask := &task{ + ctx: AppendAddressSegment(ctx, AddressSegmentNode, key), + nodeKey: key, + call: call, + input: input, + option: nil, + skipPreHandler: skipPreHandler[key], + } + if opt, ok := optMap[key]; ok { + newTask.option = opt + } + + ret = append(ret, newTask) + } + return ret, nil +} + +func (r *runner) resolveCompletedTasks(ctx context.Context, completedTasks []*task, isStream bool, cm *channelManager) (map[string]map[string]any, map[string][]string, error) { + writeChannelValues := make(map[string]map[string]any) + newDependencies := make(map[string][]string) + for _, t := range completedTasks { + for _, key := range t.call.controls { + newDependencies[key] = append(newDependencies[key], t.nodeKey) + } + + // update channel & new_next_tasks + vs := copyItem(t.output, len(t.call.writeTo)+len(t.call.writeToBranches)*2) + nextNodeKeys, err := r.calculateBranch(ctx, t.nodeKey, t.call, + vs[len(t.call.writeTo)+len(t.call.writeToBranches):], isStream, cm) + if err != nil { + return nil, nil, fmt.Errorf("calculate next step fail, node: %s, error: %w", t.nodeKey, err) + } + + for _, key := range nextNodeKeys { + newDependencies[key] = append(newDependencies[key], t.nodeKey) + } + nextNodeKeys = append(nextNodeKeys, t.call.writeTo...) + + // If branches generates more than one successor, the inputs need to be copied accordingly. + if len(nextNodeKeys) > 0 { + toCopyNum := len(nextNodeKeys) - len(t.call.writeTo) - len(t.call.writeToBranches) + nVs := copyItem(vs[len(t.call.writeTo)+len(t.call.writeToBranches)-1], toCopyNum+1) + vs = append(vs[:len(t.call.writeTo)+len(t.call.writeToBranches)-1], nVs...) + + for i, next := range nextNodeKeys { + if _, ok := writeChannelValues[next]; !ok { + writeChannelValues[next] = make(map[string]any) + } + writeChannelValues[next][t.nodeKey] = vs[i] + } + } + } + return writeChannelValues, newDependencies, nil +} + +func (r *runner) calculateBranch(ctx context.Context, curNodeKey string, startChan *chanCall, input []any, isStream bool, cm *channelManager) ([]string, error) { + if len(input) < len(startChan.writeToBranches) { + // unreachable + return nil, errors.New("calculate next input length is shorter than branches") + } + + ret := make([]string, 0, len(startChan.writeToBranches)) + + skippedNodes := make(map[string]struct{}) + for i, branch := range startChan.writeToBranches { + // check branch input type if needed + var err error + input[i], err = r.preBranchHandlerManager.handle(curNodeKey, i, input[i], isStream) + if err != nil { + return nil, fmt.Errorf("branch[%s]-[%d] pre handler fail: %w", curNodeKey, branch.idx, err) + } + + // process branch output + var ws []string + if isStream { + ws, err = branch.collect(ctx, input[i].(streamReader)) + if err != nil { + return nil, fmt.Errorf("branch collect run error: %w", err) + } + } else { + ws, err = branch.invoke(ctx, input[i]) + if err != nil { + return nil, fmt.Errorf("branch invoke run error: %w", err) + } + } + + for node := range branch.endNodes { + skipped := true + for _, w := range ws { + if node == w { + skipped = false + break + } + } + if skipped { + skippedNodes[node] = struct{}{} + } + } + + ret = append(ret, ws...) + } + + // When a node has multiple branches, + // there may be a situation where a succeeding node is selected by some branches and discarded by the other branches, + // in which case the succeeding node should not be skipped. + var skippedNodeList []string + for _, selected := range ret { + if _, ok := skippedNodes[selected]; ok { + delete(skippedNodes, selected) + } + } + for skipped := range skippedNodes { + skippedNodeList = append(skippedNodeList, skipped) + } + + err := cm.reportBranch(curNodeKey, skippedNodeList) + if err != nil { + return nil, err + } + return ret, nil +} + +func (r *runner) initTaskManager(runWrapper runnableCallWrapper, cancelVal *graphCancelChanVal, opts ...Option) *taskManager { + tm := &taskManager{ + runWrapper: runWrapper, + opts: opts, + needAll: !r.eager, + done: internal.NewUnboundedChan[*task](), + runningTasks: make(map[string]*task), + persistRerunInput: cancelVal != nil, + } + if cancelVal != nil { + tm.cancelCh = cancelVal.ch + } + return tm +} + +func (r *runner) initChannelManager(isStream bool) *channelManager { + builder := r.chanBuilder + if builder == nil { + builder = pregelChannelBuilder + } + + chs := make(map[string]channel) + for ch := range r.chanSubscribeTo { + chs[ch] = builder(r.controlPredecessors[ch], r.dataPredecessors[ch], r.chanSubscribeTo[ch].action.inputZeroValue, r.chanSubscribeTo[ch].action.inputEmptyStream) + } + + chs[END] = builder(r.controlPredecessors[END], r.dataPredecessors[END], r.outputZeroValue, r.outputEmptyStream) + + dataPredecessors := make(map[string]map[string]struct{}) + for k, vs := range r.dataPredecessors { + dataPredecessors[k] = make(map[string]struct{}) + for _, v := range vs { + dataPredecessors[k][v] = struct{}{} + } + } + controlPredecessors := make(map[string]map[string]struct{}) + for k, vs := range r.controlPredecessors { + controlPredecessors[k] = make(map[string]struct{}) + for _, v := range vs { + controlPredecessors[k][v] = struct{}{} + } + } + + for k, v := range chs { + if cfg, ok := r.mergeConfigs[k]; ok { + v.setMergeConfig(cfg) + } + } + + return &channelManager{ + isStream: isStream, + channels: chs, + successors: r.successors, + dataPredecessors: dataPredecessors, + controlPredecessors: controlPredecessors, + + edgeHandlerManager: r.edgeHandlerManager, + preNodeHandlerManager: r.preNodeHandlerManager, + } +} + +func (r *runner) toComposableRunnable() *composableRunnable { + cr := &composableRunnable{ + i: func(ctx context.Context, input any, opts ...any) (output any, err error) { + tos, err := convertOption[Option](opts...) + if err != nil { + return nil, err + } + return r.invoke(ctx, input, tos...) + }, + t: func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) { + tos, err := convertOption[Option](opts...) + if err != nil { + return nil, err + } + return r.transform(ctx, input, tos...) + }, + + inputType: r.inputType, + outputType: r.outputType, + genericHelper: r.genericHelper, + optionType: nil, // if option type is nil, graph will transmit all options. + } + + return cr +} + +func copyItem(item any, n int) []any { + if n < 2 { + return []any{item} + } + + ret := make([]any, n) + if s, ok := item.(streamReader); ok { + ss := s.copy(n) + for i := range ret { + ret[i] = ss[i] + } + + return ret + } + + for i := range ret { + ret[i] = item + } + + return ret +} + +func printTask(ts []*task) string { + if len(ts) == 0 { + return "[]" + } + sb := strings.Builder{} + sb.WriteString("[") + for i := 0; i < len(ts)-1; i++ { + sb.WriteString(ts[i].nodeKey) + sb.WriteString(", ") + } + sb.WriteString(ts[len(ts)-1].nodeKey) + sb.WriteString("]") + return sb.String() +} diff --git a/compose/graph_test.go b/compose/graph_test.go new file mode 100644 index 0000000..c2020d9 --- /dev/null +++ b/compose/graph_test.go @@ -0,0 +1,2071 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "io" + "reflect" + "sort" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/schema" +) + +func TestSingleGraph(t *testing.T) { + + const ( + nodeOfModel = "model" + nodeOfPrompt = "prompt" + ) + + ctx := context.Background() + g := NewGraph[map[string]any, *schema.Message]() + + pt := prompt.FromMessages(schema.FString, + schema.UserMessage("what's the weather in {location}?"), + ) + + err := g.AddChatTemplateNode("prompt", pt) + assert.NoError(t, err) + + cm := &chatModel{ + msgs: []*schema.Message{ + { + Role: schema.Assistant, + Content: "the weather is good", + }, + }, + } + + err = g.AddChatModelNode(nodeOfModel, cm, WithNodeName("MockChatModel")) + assert.NoError(t, err) + + err = g.AddEdge(START, nodeOfPrompt) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfPrompt, nodeOfModel) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfModel, END) + assert.NoError(t, err) + + r, err := g.Compile(context.Background(), WithMaxRunSteps(10)) + assert.NoError(t, err) + + in := map[string]any{"location": "beijing"} + _, err = r.Invoke(ctx, in) + assert.NoError(t, err) + + // stream + s, err := r.Stream(ctx, in) + assert.NoError(t, err) + + _, err = concatStreamReader(s) + assert.NoError(t, err) + + sr, sw := schema.Pipe[map[string]any](1) + _ = sw.Send(in, nil) + sw.Close() + + // transform + s, err = r.Transform(ctx, sr) + assert.NoError(t, err) + + _, err = concatStreamReader(s) + assert.NoError(t, err) + + // error test + in = map[string]any{"wrong key": 1} + _, err = r.Invoke(ctx, in) + assert.Errorf(t, err, "could not find key: location") + + _, err = r.Stream(ctx, in) + assert.Errorf(t, err, "could not find key: location") + + sr, sw = schema.Pipe[map[string]any](1) + _ = sw.Send(in, nil) + sw.Close() + + _, err = r.Transform(ctx, sr) + assert.Errorf(t, err, "could not find key: location") +} + +type person interface { + Say() string +} + +type doctor struct { + say string +} + +func (d *doctor) Say() string { + return d.say +} + +func TestGraphWithImplementableType(t *testing.T) { + + const ( + node1 = "1st" + node2 = "2nd" + ) + + ctx := context.Background() + + g := NewGraph[string, string]() + + err := g.AddLambdaNode(node1, InvokableLambda(func(ctx context.Context, input string) (output *doctor, err error) { + return &doctor{say: input}, nil + })) + assert.NoError(t, err) + + err = g.AddLambdaNode(node2, InvokableLambda(func(ctx context.Context, input person) (output string, err error) { + return input.Say(), nil + })) + assert.NoError(t, err) + + err = g.AddEdge(START, node1) + assert.NoError(t, err) + + err = g.AddEdge(node1, node2) + assert.NoError(t, err) + + err = g.AddEdge(node2, END) + assert.NoError(t, err) + + r, err := g.Compile(context.Background(), WithMaxRunSteps(10)) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "how are you", WithRuntimeMaxSteps(1)) + assert.Error(t, err) + assert.ErrorContains(t, err, "exceeds max steps") + + _, err = r.Invoke(ctx, "how are you", WithRuntimeMaxSteps(1)) + assert.Error(t, err) + assert.ErrorContains(t, err, "exceeds max steps") + + out, err := r.Invoke(ctx, "how are you") + assert.NoError(t, err) + assert.Equal(t, "how are you", out) + + outStream, err := r.Stream(ctx, "i'm fine") + assert.NoError(t, err) + defer outStream.Close() + + say, err := outStream.Recv() + assert.NoError(t, err) + assert.Equal(t, "i'm fine", say) +} + +func TestNestedGraph(t *testing.T) { + const ( + nodeOfLambda1 = "lambda1" + nodeOfLambda2 = "lambda2" + nodeOfSubGraph = "sub_graph" + nodeOfModel = "model" + nodeOfPrompt = "prompt" + ) + + ctx := context.Background() + g := NewGraph[string, *schema.Message]() + sg := NewGraph[map[string]any, *schema.Message]() + + l1 := InvokableLambda[string, map[string]any]( + func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"location": input}, nil + }) + + l2 := InvokableLambda[*schema.Message, *schema.Message]( + func(ctx context.Context, input *schema.Message) (output *schema.Message, err error) { + input.Content = fmt.Sprintf("after lambda 2: %s", input.Content) + return input, nil + }) + + pt := prompt.FromMessages(schema.FString, + schema.UserMessage("what's the weather in {location}?"), + ) + + err := sg.AddChatTemplateNode("prompt", pt) + assert.NoError(t, err) + + cm := &chatModel{ + msgs: []*schema.Message{ + { + Role: schema.Assistant, + Content: "the weather is good", + }, + }, + } + + err = sg.AddChatModelNode(nodeOfModel, cm, WithNodeName("MockChatModel")) + assert.NoError(t, err) + + err = sg.AddEdge(START, nodeOfPrompt) + assert.NoError(t, err) + + err = sg.AddEdge(nodeOfPrompt, nodeOfModel) + assert.NoError(t, err) + + err = sg.AddEdge(nodeOfModel, END) + assert.NoError(t, err) + + err = g.AddLambdaNode(nodeOfLambda1, l1, WithNodeName("Lambda1")) + assert.NoError(t, err) + + err = g.AddGraphNode(nodeOfSubGraph, sg, WithNodeName("SubGraphName")) + assert.NoError(t, err) + + err = g.AddLambdaNode(nodeOfLambda2, l2, WithNodeName("Lambda2")) + assert.NoError(t, err) + + err = g.AddEdge(START, nodeOfLambda1) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfLambda1, nodeOfSubGraph) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfSubGraph, nodeOfLambda2) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfLambda2, END) + assert.NoError(t, err) + + r, err := g.Compile(context.Background(), + WithMaxRunSteps(10), + WithGraphName("GraphName"), + ) + assert.NoError(t, err) + + ck := "depth" + cb := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + v, ok := ctx.Value(ck).(int) + if ok { + v++ + } + + return context.WithValue(ctx, ck, v) + }). + OnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + input.Close() + + v, ok := ctx.Value(ck).(int) + if ok { + v++ + } + + return context.WithValue(ctx, ck, v) + }).Build() + + // invoke + _, err = r.Invoke(ctx, "london", WithCallbacks(cb)) + assert.NoError(t, err) + + // stream + rs, err := r.Stream(ctx, "london", WithCallbacks(cb)) + assert.NoError(t, err) + for { + _, err = rs.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + } + + // collect + sr, sw := schema.Pipe[string](5) + _ = sw.Send("london", nil) + sw.Close() + + _, err = r.Collect(ctx, sr, WithCallbacks(cb)) + assert.NoError(t, err) + + // transform + sr, sw = schema.Pipe[string](5) + _ = sw.Send("london", nil) + sw.Close() + + rt, err := r.Transform(ctx, sr, WithCallbacks(cb)) + assert.NoError(t, err) + for { + _, err = rt.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + } +} + +type chatModel struct { + msgs []*schema.Message +} + +func (c *chatModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +func (c *chatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return c.msgs[0], nil +} + +func (c *chatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](len(c.msgs)) + go func() { + for _, msg := range c.msgs { + sw.Send(msg, nil) + } + sw.Close() + }() + return sr, nil +} + +func TestValidate(t *testing.T) { + // test unmatched nodes + g := NewGraph[string, string]() + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return "", nil })) + assert.NoError(t, err) + + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return "", nil })) + assert.NoError(t, err) + + err = g.AddEdge("1", "2") + assert.ErrorContains(t, err, "graph edge[1]-[2]: start node's output type[string] and end node's input type[int] mismatch") + + // test unmatched passthrough node + g = NewGraph[string, string]() + err = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return "", nil })) + assert.NoError(t, err) + + err = g.AddPassthroughNode("2") + assert.NoError(t, err) + + err = g.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return "", nil })) + assert.NoError(t, err) + + err = g.AddEdge("1", "2") + assert.NoError(t, err) + + err = g.AddEdge("2", "3") + assert.ErrorContains(t, err, "graph edge[2]-[3]: start node's output type[string] and end node's input type[int] mismatch") + + // test may matched passthrough + g2 := NewGraph[any, string]() + err = g2.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input any) (output any, err error) { return input, nil })) + assert.NoError(t, err) + err = g2.AddPassthroughNode("2") + assert.NoError(t, err) + err = g2.AddLambdaNode("3", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return strconv.Itoa(input), nil })) + assert.NoError(t, err) + err = g2.AddEdge(START, "1") + assert.NoError(t, err) + err = g2.AddEdge("2", "3") + assert.NoError(t, err) + err = g2.AddEdge("1", "2") + assert.NoError(t, err) + err = g2.AddEdge("3", END) + assert.NoError(t, err) + ru, err := g2.Compile(context.Background()) + assert.NoError(t, err) + // success + result, err := ru.Invoke(context.Background(), 1) + assert.NoError(t, err) + assert.Equal(t, result, "1") + // fail + _, err = ru.Invoke(context.Background(), "1") + assert.ErrorContains(t, err, "runtime type check") + + // test unmatched graph type + g = NewGraph[string, string]() + err = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input int) (output string, err error) { return "", nil })) + assert.NoError(t, err) + + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output int, err error) { return 0, nil })) + assert.NoError(t, err) + + err = g.AddEdge("1", "2") + assert.NoError(t, err) + + err = g.AddEdge(START, "1") + assert.ErrorContains(t, err, "graph edge[start]-[1]: start node's output type[string] and end node's input type[int] mismatch") + + // sub graph implement + type A interface { + A() + } + type B interface { + B() + } + + type AB interface { + A + B + } + lA := InvokableLambda(func(ctx context.Context, input A) (output string, err error) { return "", nil }) + lB := InvokableLambda(func(ctx context.Context, input B) (output string, err error) { return "", nil }) + lAB := InvokableLambda(func(ctx context.Context, input string) (output AB, err error) { return nil, nil }) + + p := NewParallel().AddLambda("1", lA).AddLambda("2", lB) + c := NewChain[string, map[string]any]().AppendLambda(lAB).AppendParallel(p) + _, err = c.Compile(context.Background()) + assert.NoError(t, err) + + // error usage + p = NewParallel().AddLambda("1", lA).AddLambda("2", lAB) + c = NewChain[string, map[string]any]().AppendParallel(p) + _, err = c.Compile(context.Background()) + assert.ErrorContains(t, err, "add parallel edge failed, from=start, to=node_0_parallel_0, err: graph edge[start]-[node_0_parallel_0]: start node's output type[string] and end node's input type[compose.A] mismatch") + + // test graph output type check + gg := NewGraph[string, A]() + err = gg.AddLambdaNode("nodeA", InvokableLambda(func(ctx context.Context, input string) (output A, err error) { return nil, nil })) + assert.NoError(t, err) + + err = gg.AddLambdaNode("nodeA2", InvokableLambda(func(ctx context.Context, input string) (output A, err error) { return nil, nil })) + assert.NoError(t, err) + + err = gg.AddLambdaNode("nodeB", InvokableLambda(func(ctx context.Context, input string) (output B, err error) { return nil, nil })) + assert.NoError(t, err) + + err = gg.AddEdge("nodeA", END) + assert.NoError(t, err) + + err = gg.AddEdge("nodeB", END) + assert.ErrorContains(t, err, "graph edge[nodeB]-[end]: start node's output type[compose.B] and end node's input type[compose.A] mismatch") + + err = gg.AddEdge("nodeA2", END) + assert.ErrorContains(t, err, "graph edge[nodeB]-[end]: start node's output type[compose.B] and end node's input type[compose.A] mismatch") + + // test any type + anyG := NewGraph[any, string]() + err = anyG.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node1", nil })) + assert.NoError(t, err) + + err = anyG.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node2", nil })) + assert.NoError(t, err) + + err = anyG.AddEdge(START, "node1") + assert.NoError(t, err) + + err = anyG.AddEdge("node1", "node2") + assert.NoError(t, err) + + err = anyG.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + r, err := anyG.Compile(context.Background()) + assert.NoError(t, err) + + result, err = r.Invoke(context.Background(), "start") + assert.NoError(t, err) + assert.Equal(t, "startnode1node2", result) + + streamResult, err := r.Stream(context.Background(), "start") + assert.NoError(t, err) + + result = "" + for { + chunk, err := streamResult.Recv() + if err != nil { + if err == io.EOF { + break + } + assert.NoError(t, err) + } + result += chunk + } + + assert.Equal(t, "startnode1node2", result) + + // test any type runtime error + anyG = NewGraph[any, string]() + err = anyG.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return 123, nil })) + if err != nil { + t.Fatal(err) + } + err = anyG.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node2", nil })) + if err != nil { + t.Fatal(err) + } + err = anyG.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = anyG.AddEdge("node1", "node2") + if err != nil { + t.Fatal(err) + } + err = anyG.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + r, err = anyG.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + _, err = r.Invoke(context.Background(), "start") + if err == nil || !strings.Contains(err.Error(), "runtime") { + t.Fatal("test any type runtime error fail, error is nil or error doesn't contain key word runtime") + } + _, err = r.Stream(context.Background(), "start") + if err == nil || !strings.Contains(err.Error(), "runtime") { + t.Fatal("test any type runtime error fail, error is nil or error doesn't contain key word runtime") + } + + // test branch any type + // success + g = NewGraph[string, string]() + err = g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node1", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node2", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node3", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node3", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch("node1", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + return "node2", nil + }, map[string]bool{"node2": true, "node3": true})) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node3", END) + if err != nil { + t.Fatal(err) + } + rr, err := g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + ret, err := rr.Invoke(context.Background(), "start") + if err != nil { + t.Fatal(err) + } + if ret != "startnode1node2" { + t.Fatal("test branch any type fail, result is unexpected") + } + streamResult, err = rr.Stream(context.Background(), "start") + if err != nil { + t.Fatal(err) + } + ret, err = concatStreamReader(streamResult) + if err != nil { + t.Fatal(err) + } + if ret != "startnode1node2" { + t.Fatal("test branch any type fail, result is unexpected") + } + // fail + g = NewGraph[string, string]() + err = g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return 1 /*error type*/, nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node2", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node3", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node3", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch("node1", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + return "node2", nil + }, map[string]bool{"node2": true, "node3": true})) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node3", END) + if err != nil { + t.Fatal(err) + } + rr, err = g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + _, err = rr.Invoke(context.Background(), "start") + if err == nil || !strings.Contains(err.Error(), "runtime") { + t.Fatal("test branch any type fail, haven't report runtime error") + } + _, err = rr.Stream(context.Background(), "start") + if err == nil || !strings.Contains(err.Error(), "runtime") { + t.Fatal("test branch any type fail, haven't report runtime error") + } +} + +func TestValidateMultiAnyValueBranch(t *testing.T) { + // success + g := NewGraph[string, map[string]any]() + err := g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node1", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node2": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node3", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node3": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node4", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node4": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node5", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node5": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch("node1", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + return "node2", nil + }, map[string]bool{"node2": true, "node3": true})) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch("node1", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + return "node4", nil + }, map[string]bool{"node4": true, "node5": true})) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node3", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node4", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node5", END) + if err != nil { + t.Fatal(err) + } + rr, err := g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + ret, err := rr.Invoke(context.Background(), "start") + if err != nil { + t.Fatal(err) + } + if !ret["node2"].(bool) || !ret["node4"].(bool) { + t.Fatal("test branch any type fail, result is unexpected") + } + streamResult, err := rr.Stream(context.Background(), "start") + if err != nil { + t.Fatal(err) + } + ret, err = concatStreamReader(streamResult) + if err != nil { + t.Fatal(err) + } + if !ret["node2"].(bool) || !ret["node4"].(bool) { + t.Fatal("test branch any type fail, result is unexpected") + } + + // fail + g = NewGraph[string, map[string]any]() + err = g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node1", nil })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node2": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node3", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node3": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node4", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node4": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node5", InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{"node5": true}, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch("node1", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + return "node2", nil + }, map[string]bool{"node2": true, "node3": true})) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch("node1", NewGraphBranch(func(ctx context.Context, in int /*error type*/) (endNode string, err error) { + return "node4", nil + }, map[string]bool{"node4": true, "node5": true})) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node3", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node4", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node5", END) + if err != nil { + t.Fatal(err) + } + rr, err = g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + _, err = rr.Invoke(context.Background(), "start") + if err == nil || !strings.Contains(err.Error(), "runtime") { + t.Fatal("test multi branch any type fail, haven't report runtime error") + } + _, err = rr.Stream(context.Background(), "start") + if err == nil || !strings.Contains(err.Error(), "runtime") { + t.Fatal("test multi branch any type fail, haven't report runtime error") + } +} + +func TestAnyTypeWithKey(t *testing.T) { + g := NewGraph[any, map[string]any]() + err := g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node1", nil }), WithInputKey("node1")) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output any, err error) { return input + "node2", nil }), WithOutputKey("node2")) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node1", "node2") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + r, err := g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + result, err := r.Invoke(context.Background(), map[string]any{"node1": "start"}) + if err != nil { + t.Fatal(err) + } + if result["node2"] != "startnode1node2" { + t.Fatal("test any type with key fail, result is unexpected") + } + + streamResult, err := r.Stream(context.Background(), map[string]any{"node1": "start"}) + if err != nil { + t.Fatal(err) + } + ret, err := concatStreamReader(streamResult) + if err != nil { + t.Fatal(err) + } + if ret["node2"] != "startnode1node2" { + t.Fatal("test any type with key fail, result is unexpected") + } +} + +func TestInputKey(t *testing.T) { + g := NewGraph[map[string]any, map[string]any]() + err := g.AddChatTemplateNode("1", prompt.FromMessages(schema.FString, schema.UserMessage("{var1}")), WithOutputKey("1"), WithInputKey("1")) + if err != nil { + t.Fatal(err) + } + err = g.AddChatTemplateNode("2", prompt.FromMessages(schema.FString, schema.UserMessage("{var2}")), WithOutputKey("2"), WithInputKey("2")) + if err != nil { + t.Fatal(err) + } + err = g.AddChatTemplateNode("3", prompt.FromMessages(schema.FString, schema.UserMessage("{var3}")), WithOutputKey("3"), WithInputKey("3")) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "2") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "3") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("2", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("3", END) + if err != nil { + t.Fatal(err) + } + r, err := g.Compile(context.Background(), WithMaxRunSteps(100)) + if err != nil { + t.Fatal(err) + } + + ctx := context.Background() + result, err := r.Invoke(ctx, map[string]any{ + "1": map[string]any{"var1": "a"}, + "2": map[string]any{"var2": "b"}, + "3": map[string]any{"var3": "c"}, + }) + if err != nil { + t.Fatal(err) + } + if result["1"].([]*schema.Message)[0].Content != "a" || + result["2"].([]*schema.Message)[0].Content != "b" || + result["3"].([]*schema.Message)[0].Content != "c" { + t.Fatal("invoke different") + } + + sr, sw := schema.Pipe[map[string]any](10) + sw.Send(map[string]any{"1": map[string]any{"var1": "a"}}, nil) + sw.Send(map[string]any{"2": map[string]any{"var2": "b"}}, nil) + sw.Send(map[string]any{"3": map[string]any{"var3": "c"}}, nil) + sw.Close() + + streamResult, err := r.Transform(ctx, sr) + if err != nil { + t.Fatal(err) + } + defer streamResult.Close() + + result = make(map[string]any) + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + for k, v := range chunk { + result[k] = v + } + } + if result["1"].([]*schema.Message)[0].Content != "a" || + result["2"].([]*schema.Message)[0].Content != "b" || + result["3"].([]*schema.Message)[0].Content != "c" { + t.Fatal("transform different") + } +} + +func TestTransferTask(t *testing.T) { + in := [][]string{ + { + "1", + "2", + }, + { + "3", + "4", + "5", + "6", + }, + { + "5", + "6", + "7", + }, + { + "7", + "8", + }, + { + "8", + }, + } + invertedEdges := map[string][]string{ + "1": {"3", "4"}, + "2": {"5", "6"}, + "3": {"5"}, + "4": {"6"}, + "5": {"7"}, + "7": {"8"}, + } + in = transferTask(in, invertedEdges) + + if !reflect.DeepEqual( + [][]string{ + { + "1", + }, + { + "3", + "2", + }, + { + "5", + }, + { + "7", + "4", + }, + { + "8", + "6", + }, + }, in) { + t.Fatal("not equal") + } +} + +func TestPregelEnd(t *testing.T) { + g := NewGraph[string, string]() + err := g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "node1", nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddLambdaNode("node2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "node2", nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node1", END) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node1", "node2") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node2", END) + if err != nil { + t.Fatal(err) + } + runner, err := g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + + out, err := runner.Invoke(context.Background(), "") + if err != nil { + t.Fatal(err) + } + if out != "node1" { + t.Fatal("graph output is unexpected") + } +} + +type cb struct { + gInfo *GraphInfo +} + +func (c *cb) OnFinish(ctx context.Context, info *GraphInfo) { + c.gInfo = info +} + +func TestGraphCompileCallback(t *testing.T) { + t.Run("graph compile callback", func(t *testing.T) { + type s struct{} + + g := NewGraph[map[string]any, map[string]any](WithGenLocalState(func(ctx context.Context) *s { return &s{} })) + + lambda := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "node1", nil + }) + lambdaOpts := []GraphAddNodeOpt{WithNodeName("lambda_1"), WithInputKey("input_key")} + err := g.AddLambdaNode("node1", lambda, lambdaOpts...) + assert.NoError(t, err) + + err = g.AddPassthroughNode("pass1") + assert.NoError(t, err) + err = g.AddPassthroughNode("pass2") + assert.NoError(t, err) + + condition := func(ctx context.Context, input string) (string, error) { + return input, nil + } + + branch := NewGraphBranch(condition, map[string]bool{"pass1": true, "pass2": true}) + err = g.AddBranch("node1", branch) + assert.NoError(t, err) + + err = g.AddEdge(START, "node1") + assert.NoError(t, err) + + lambda2 := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "node2", nil + }) + lambdaOpts2 := []GraphAddNodeOpt{WithNodeName("lambda_2")} + subSubGraph := NewGraph[string, string]() + err = subSubGraph.AddLambdaNode("sub1", lambda2, lambdaOpts2...) + assert.NoError(t, err) + err = subSubGraph.AddEdge(START, "sub1") + assert.NoError(t, err) + err = subSubGraph.AddEdge("sub1", END) + assert.NoError(t, err) + + subGraph := NewGraph[string, string]() + var ssGraphCompileOpts []GraphCompileOption + ssGraphOpts := []GraphAddNodeOpt{WithGraphCompileOptions(ssGraphCompileOpts...)} + err = subGraph.AddGraphNode("sub_sub_1", subSubGraph, ssGraphOpts...) + assert.NoError(t, err) + err = subGraph.AddEdge(START, "sub_sub_1") + assert.NoError(t, err) + err = subGraph.AddEdge("sub_sub_1", END) + assert.NoError(t, err) + + subGraphCompileOpts := []GraphCompileOption{WithMaxRunSteps(2), WithGraphName("sub_graph")} + subGraphOpts := []GraphAddNodeOpt{WithGraphCompileOptions(subGraphCompileOpts...)} + err = g.AddGraphNode("sub_graph", subGraph, subGraphOpts...) + assert.NoError(t, err) + + err = g.AddEdge("pass1", "sub_graph") + assert.NoError(t, err) + err = g.AddEdge("pass2", "sub_graph") + assert.NoError(t, err) + + lambda3 := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "node3", nil + }) + lambdaOpts3 := []GraphAddNodeOpt{WithNodeName("lambda_3"), WithOutputKey("lambda_3")} + err = g.AddLambdaNode("node3", lambda3, lambdaOpts3...) + assert.NoError(t, err) + + lambda4 := InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "node4", nil + }) + lambdaOpts4 := []GraphAddNodeOpt{WithNodeName("lambda_4"), WithOutputKey("lambda_4")} + err = g.AddLambdaNode("node4", lambda4, lambdaOpts4...) + assert.NoError(t, err) + + err = g.AddEdge("sub_graph", "node3") + assert.NoError(t, err) + err = g.AddEdge("sub_graph", "node4") + assert.NoError(t, err) + err = g.AddEdge("node3", END) + assert.NoError(t, err) + err = g.AddEdge("node4", END) + assert.NoError(t, err) + + c := &cb{} + opt := []GraphCompileOption{WithGraphCompileCallbacks(c), WithGraphName("top_level")} + _, err = g.Compile(context.Background(), opt...) + assert.NoError(t, err) + expected := &GraphInfo{ + CompileOptions: opt, + Nodes: map[string]GraphNodeInfo{ + "node1": { + Component: ComponentOfLambda, + Instance: lambda, + GraphAddNodeOpts: lambdaOpts, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "lambda_1", + InputKey: "input_key", + }, + "pass1": { + Component: ComponentOfPassthrough, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "", + }, + "pass2": { + Component: ComponentOfPassthrough, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "", + }, + "sub_graph": { + Component: ComponentOfGraph, + Instance: subGraph, + GraphAddNodeOpts: subGraphOpts, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "", + GraphInfo: &GraphInfo{ + CompileOptions: subGraphCompileOpts, + Nodes: map[string]GraphNodeInfo{ + "sub_sub_1": { + Component: ComponentOfGraph, + Instance: subSubGraph, + GraphAddNodeOpts: ssGraphOpts, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "", + GraphInfo: &GraphInfo{ + CompileOptions: ssGraphCompileOpts, + Nodes: map[string]GraphNodeInfo{ + "sub1": { + Component: ComponentOfLambda, + Instance: lambda2, + GraphAddNodeOpts: lambdaOpts2, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "lambda_2", + }, + }, + Edges: map[string][]string{ + START: {"sub1"}, + "sub1": {END}, + }, + DataEdges: map[string][]string{ + START: {"sub1"}, + "sub1": {END}, + }, + Branches: map[string][]GraphBranch{}, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + }, + }, + }, + Edges: map[string][]string{ + START: {"sub_sub_1"}, + "sub_sub_1": {END}, + }, + DataEdges: map[string][]string{ + START: {"sub_sub_1"}, + "sub_sub_1": {END}, + }, + Branches: map[string][]GraphBranch{}, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "sub_graph", + }, + }, + "node3": { + Component: ComponentOfLambda, + Instance: lambda3, + GraphAddNodeOpts: lambdaOpts3, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "lambda_3", + OutputKey: "lambda_3", + }, + "node4": { + Component: ComponentOfLambda, + Instance: lambda4, + GraphAddNodeOpts: lambdaOpts4, + InputType: reflect.TypeOf(""), + OutputType: reflect.TypeOf(""), + Name: "lambda_4", + OutputKey: "lambda_4", + }, + }, + Edges: map[string][]string{ + START: {"node1"}, + "pass1": {"sub_graph"}, + "pass2": {"sub_graph"}, + "sub_graph": {"node3", "node4"}, + "node3": {END}, + "node4": {END}, + }, + DataEdges: map[string][]string{ + START: {"node1"}, + "pass1": {"sub_graph"}, + "pass2": {"sub_graph"}, + "sub_graph": {"node3", "node4"}, + "node3": {END}, + "node4": {END}, + }, + Branches: map[string][]GraphBranch{ + "node1": {*branch}, + }, + InputType: reflect.TypeOf(map[string]any{}), + OutputType: reflect.TypeOf(map[string]any{}), + Name: "top_level", + } + + stateFn := c.gInfo.GenStateFn + assert.NotNil(t, stateFn) + assert.Equal(t, &s{}, stateFn(context.Background())) + + assert.Equal(t, 1, len(c.gInfo.NewGraphOptions)) + c.gInfo.NewGraphOptions = nil + + c.gInfo.GenStateFn = nil + + actualCompileOptions := newGraphCompileOptions(c.gInfo.CompileOptions...) + expectedCompileOptions := newGraphCompileOptions(expected.CompileOptions...) + assert.Equal(t, len(expectedCompileOptions.callbacks), len(actualCompileOptions.callbacks)) + assert.Same(t, expectedCompileOptions.callbacks[0], actualCompileOptions.callbacks[0]) + actualCompileOptions.callbacks = nil + actualCompileOptions.origOpts = nil + expectedCompileOptions.callbacks = nil + expectedCompileOptions.origOpts = nil + assert.Equal(t, expectedCompileOptions, actualCompileOptions) + + c.gInfo.CompileOptions = nil + expected.CompileOptions = nil + assert.Equal(t, expected.Branches["node1"][0].endNodes, c.gInfo.Branches["node1"][0].endNodes) + assert.Equal(t, expected.Branches["node1"][0].inputType, c.gInfo.Branches["node1"][0].inputType) + + expected.Branches["node1"] = []GraphBranch{} + c.gInfo.Branches["node1"] = []GraphBranch{} + assert.Equal(t, expected, c.gInfo) + }) +} + +func TestCheckAddEdge(t *testing.T) { + g := NewGraph[string, string]() + err := g.AddPassthroughNode("1") + if err != nil { + t.Fatal(err) + } + err = g.AddPassthroughNode("2") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", "2") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", "2") + if err == nil { + t.Fatal("add edge repeatedly haven't report error") + } +} + +func TestStartWithEnd(t *testing.T) { + g := NewGraph[string, string]() + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + return END, nil + }, map[string]bool{"1": true, END: true})) + if err != nil { + t.Fatal(err) + } + r, err := g.Compile(context.Background()) + if err != nil { + t.Fatal(err) + } + sr, sw := schema.Pipe[string](1) + sw.Send("test", nil) + sw.Close() + result, err := r.Transform(context.Background(), sr) + if err != nil { + t.Fatal(err) + } + for { + chunk, err := result.Recv() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + if chunk != "test" { + t.Fatal("result is out of expect") + } + } +} + +func TestToString(t *testing.T) { + ps := runTypePregel.String() + assert.Equal(t, "Pregel", ps) + + ds := runTypeDAG + assert.Equal(t, "DAG", ds.String()) +} + +func TestInputKeyError(t *testing.T) { + g := NewGraph[map[string]any, string]() + err := g.AddLambdaNode("node1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithInputKey("node1")) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "node1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("node1", END) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + r, err := g.Compile(ctx) + if err != nil { + t.Fatal(err) + } + // invoke + _, err = r.Invoke(ctx, map[string]any{"unknown": "123"}) + if err == nil || !strings.Contains(err.Error(), "cannot find input key: node1") { + t.Fatal("cannot report input key error correctly") + } + + // transform + sr, sw := schema.Pipe[map[string]any](1) + sw.Send(map[string]any{"unknown": "123"}, nil) + sw.Close() + _, err = r.Transform(ctx, sr) + if err == nil || !strings.Contains(err.Error(), "stream reader is empty, concat fail") { + t.Fatal("cannot report input key error correctly") + } +} + +func TestContextCancel(t *testing.T) { + ctx := context.Background() + g := NewGraph[string, string]() + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", END) + if err != nil { + t.Fatal(err) + } + r, err := g.Compile(ctx) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(ctx) + cancel() + _, err = r.Invoke(ctx, "test") + if !strings.Contains(err.Error(), "context has been canceled") { + t.Fatal("graph have not returned canceled error") + } +} + +func TestDAGStart(t *testing.T) { + g := NewGraph[map[string]any, map[string]any]() + err := g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) { + return map[string]any{"1": "1"}, nil + })) + assert.NoError(t, err) + err = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) { + return input, nil + })) + assert.NoError(t, err) + err = g.AddEdge(START, "1") + assert.NoError(t, err) + err = g.AddEdge("1", "2") + assert.NoError(t, err) + err = g.AddEdge(START, "2") + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + r, err := g.Compile(context.Background(), WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + result, err := r.Invoke(context.Background(), map[string]any{"start": "start"}) + assert.NoError(t, err) + assert.Equal(t, map[string]any{"start": "start", "1": "1"}, result) +} + +func concatLambda(s string) *Lambda { + return InvokableLambda(func(ctx context.Context, input string) (output string, err error) { return input + s, nil }) +} +func mapLambda(k, v string) *Lambda { + return InvokableLambda(func(ctx context.Context, input map[string]string) (output map[string]string, err error) { + return map[string]string{ + k: v, + }, nil + }) +} + +func TestBaseDAGBranch(t *testing.T) { + g := NewGraph[string, string]() + + err := g.AddLambdaNode("1", concatLambda("1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", concatLambda("2")) + assert.NoError(t, err) + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + if len(in) > 3 { + return "2", nil + } + return "1", nil + }, map[string]bool{"1": true, "2": true})) + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + result, err := r.Invoke(ctx, "hi") + assert.NoError(t, err) + assert.Equal(t, "hi1", result) +} + +func TestMultiDAGBranch(t *testing.T) { + g := NewGraph[map[string]string, map[string]string]() + + err := g.AddLambdaNode("1", mapLambda("1", "1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", mapLambda("2", "2")) + assert.NoError(t, err) + err = g.AddLambdaNode("3", mapLambda("3", "3")) + assert.NoError(t, err) + err = g.AddLambdaNode("4", mapLambda("4", "4")) + assert.NoError(t, err) + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) { + if len(in["input"]) > 3 { + return "2", nil + } + return "1", nil + }, map[string]bool{"1": true, "2": true})) + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) { + if len(in["input"]) > 3 { + return "4", nil + } + return "3", nil + }, map[string]bool{"3": true, "4": true})) + assert.NoError(t, err) + + err = g.AddEdge("1", END) + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + err = g.AddEdge("4", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + result, err := r.Invoke(ctx, map[string]string{"input": "hi"}) + assert.NoError(t, err) + assert.Equal(t, map[string]string{ + "1": "1", + "3": "3", + }, result) +} + +func TestCrossDAGBranch(t *testing.T) { + g := NewGraph[map[string]string, map[string]string]() + + err := g.AddLambdaNode("1", mapLambda("1", "1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", mapLambda("2", "2")) + assert.NoError(t, err) + err = g.AddLambdaNode("3", mapLambda("3", "3")) + assert.NoError(t, err) + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) { + if len(in["input"]) > 3 { + return "2", nil + } + return "1", nil + }, map[string]bool{"1": true, "2": true})) + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in map[string]string) (endNode string, err error) { + if len(in["input"]) > 3 { + return "3", nil + } + return "2", nil + }, map[string]bool{"2": true, "3": true})) + assert.NoError(t, err) + + err = g.AddEdge("1", END) + assert.NoError(t, err) + err = g.AddEdge("2", END) + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + result, err := r.Invoke(ctx, map[string]string{"input": "hi"}) + assert.NoError(t, err) + assert.Equal(t, map[string]string{ + "1": "1", + "2": "2", + }, result) +} + +func TestNestedDAGBranch(t *testing.T) { + g := NewGraph[string, string]() + + err := g.AddLambdaNode("1", concatLambda("1")) + assert.NoError(t, err) + err = g.AddLambdaNode("2", concatLambda("2")) + assert.NoError(t, err) + err = g.AddLambdaNode("3", concatLambda("3")) + assert.NoError(t, err) + err = g.AddLambdaNode("4", concatLambda("4")) + assert.NoError(t, err) + err = g.AddBranch(START, NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + if len(in) > 3 { + return "2", nil + } + return "1", nil + }, map[string]bool{"1": true, "2": true})) + err = g.AddBranch("2", NewGraphBranch(func(ctx context.Context, in string) (endNode string, err error) { + if len(in) > 10 { + return "4", nil + } + return "3", nil + }, map[string]bool{"3": true, "4": true})) + assert.NoError(t, err) + err = g.AddEdge("1", END) + assert.NoError(t, err) + err = g.AddEdge("3", END) + assert.NoError(t, err) + err = g.AddEdge("4", END) + assert.NoError(t, err) + + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + result, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, "hello23", result) + result, err = r.Invoke(ctx, "hi") + assert.NoError(t, err) + assert.Equal(t, "hi1", result) + result, err = r.Invoke(ctx, "hellohello") + assert.NoError(t, err) + assert.Equal(t, "hellohello24", result) +} + +func TestHandlerTypeValidate(t *testing.T) { + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state string) { + return "" + })) + // passthrough pre fail + err := g.AddPassthroughNode("1", WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + })) + assert.ErrorContains(t, err, "passthrough node[1]'s pre handler type isn't any") + g.buildError = nil + // passthrough pre fail with input key + err = g.AddPassthroughNode("1", WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + }), WithInputKey("input")) + assert.ErrorContains(t, err, "node[1]'s pre handler type[string] is different from its input type[map[string]interface {}]") + g.buildError = nil + // passthrough post fail + err = g.AddPassthroughNode("1", WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + })) + assert.ErrorContains(t, err, "passthrough node[1]'s post handler type isn't any") + g.buildError = nil + // passthrough post fail with input key + err = g.AddPassthroughNode("1", WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + }), WithInputKey("input")) + assert.ErrorContains(t, err, "passthrough node[1]'s post handler type isn't any") + g.buildError = nil + // passthrough pre success + err = g.AddPassthroughNode("1", WithStatePreHandler(func(ctx context.Context, in any, state string) (any, error) { + return "", nil + })) + assert.NoError(t, err) + // passthrough pre success with input key + err = g.AddPassthroughNode("2", WithStatePreHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) { + return nil, nil + }), WithInputKey("input")) + assert.NoError(t, err) + // passthrough post success + err = g.AddPassthroughNode("3", WithStatePostHandler(func(ctx context.Context, in any, state string) (any, error) { + return "", nil + })) + assert.NoError(t, err) + // passthrough post success with output key + err = g.AddPassthroughNode("4", WithStatePostHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) { + return nil, nil + }), WithOutputKey("output")) + assert.NoError(t, err) + // common node pre fail + err = g.AddLambdaNode("5", InvokableLambda(func(ctx context.Context, input int) (output int, err error) { + return 0, nil + }), WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + })) + assert.ErrorContains(t, err, "node[5]'s pre handler type[string] is different from its input type[int]") + g.buildError = nil + // common node post fail + err = g.AddLambdaNode("5", InvokableLambda(func(ctx context.Context, input int) (output int, err error) { + return 0, nil + }), WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + })) + assert.ErrorContains(t, err, "node[5]'s post handler type[string] is different from its output type[int]") + g.buildError = nil + // common node pre success + err = g.AddLambdaNode("5", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", nil + }), WithStatePreHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + })) + assert.NoError(t, err) + // common node post success + err = g.AddLambdaNode("6", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", nil + }), WithStatePostHandler(func(ctx context.Context, in string, state string) (string, error) { + return "", nil + })) + assert.NoError(t, err) + // pre state fail + err = g.AddLambdaNode("7", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", nil + }), WithStatePreHandler(func(ctx context.Context, in string, state int) (string, error) { + return "", nil + })) + assert.ErrorContains(t, err, "node[7]'s pre handler state type[int] is different from graph[string]") + g.buildError = nil + // post state fail + err = g.AddLambdaNode("7", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", nil + }), WithStatePostHandler(func(ctx context.Context, in string, state int) (string, error) { + return "", nil + })) + assert.ErrorContains(t, err, "node[7]'s post handler state type[int] is different from graph[string]") + g.buildError = nil + // common pre success with input key + err = g.AddLambdaNode("7", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", nil + }), WithStatePreHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) { + return nil, nil + }), WithInputKey("input")) + assert.NoError(t, err) + // common post success with output key + err = g.AddLambdaNode("8", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return "", nil + }), WithStatePostHandler(func(ctx context.Context, in map[string]any, state string) (map[string]any, error) { + return nil, nil + }), WithOutputKey("output")) + assert.NoError(t, err) +} + +func TestSetFanInMergeConfig_RealStreamNode(t *testing.T) { + for _, triggerMode := range []NodeTriggerMode{AnyPredecessor, AllPredecessor} { + t.Run(string(triggerMode), func(t *testing.T) { + g := NewGraph[int, map[string]any]() + + // Add two stream nodes that output streams of int slices + err := g.AddLambdaNode("s1", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) { + sr, sw := schema.Pipe[int](2) + sw.Send(input+1, nil) + sw.Send(input+2, nil) + sw.Close() + return sr, nil + }), WithOutputKey("s1")) + assert.NoError(t, err) + err = g.AddLambdaNode("s2", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) { + sr, sw := schema.Pipe[int](2) + sw.Send(input+10, nil) + sw.Send(input+20, nil) + sw.Close() + return sr, nil + }), WithOutputKey("s2")) + assert.NoError(t, err) + + // Connect edges: START -> s1, START -> s2, s1 -> END, s2 -> END + err = g.AddEdge(START, "s1") + assert.NoError(t, err) + err = g.AddEdge(START, "s2") + assert.NoError(t, err) + err = g.AddEdge("s1", END) + assert.NoError(t, err) + err = g.AddEdge("s2", END) + assert.NoError(t, err) + + r, err := g.Compile(context.Background(), WithNodeTriggerMode(triggerMode), + WithFanInMergeConfig(map[string]FanInMergeConfig{END: {StreamMergeWithSourceEOF: true}})) + assert.NoError(t, err) + + // Run the graph in stream mode and check for SourceEOF events + sr, err := r.Stream(context.Background(), 1) + assert.NoError(t, err) + + merged := make(map[string]map[int]bool) + var sourceEOFCount int + sourceNames := make(map[string]bool) + for { + m, e := sr.Recv() + if e != nil { + if name, ok := schema.GetSourceName(e); ok { + sourceEOFCount++ + sourceNames[name] = true + continue + } + if e == io.EOF { + break + } + assert.NoError(t, e) + } + + for k, v := range m { + if merged[k] == nil { + merged[k] = make(map[int]bool) + } + + merged[k][v.(int)] = true + } + } + + // The merged map should contain both results + assert.Equal(t, map[string]map[int]bool{"s1": {2: true, 3: true}, "s2": {11: true, 21: true}}, merged) + assert.Equal(t, 2, sourceEOFCount, "should receive SourceEOF for each input stream when StreamMergeWithSourceEOF is true") + assert.True(t, sourceNames["s1"], "should receive SourceEOF from s1") + assert.True(t, sourceNames["s2"], "should receive SourceEOF from s2") + }) + } +} + +func TestFindLoops(t *testing.T) { + tests := []struct { + name string + startNodes []string + chanCalls map[string]*chanCall + expected [][]string + }{ + { + name: "Graph without cycles", + startNodes: []string{"A"}, + chanCalls: map[string]*chanCall{ + "A": { + controls: []string{"B", "C"}, + }, + "B": { + controls: []string{"D"}, + }, + "C": { + controls: []string{"E"}, + }, + "D": { + controls: []string{}, + }, + "E": { + controls: []string{}, + }, + }, + expected: [][]string{}, + }, + { + name: "Graph with self-loop", + startNodes: []string{"A"}, + chanCalls: map[string]*chanCall{ + "A": { + controls: []string{"A", "B"}, + }, + "B": { + controls: []string{}, + }, + }, + expected: [][]string{{"A", "A"}}, + }, + { + name: "Graph with simple cycle", + startNodes: []string{"A", "B", "C"}, + chanCalls: map[string]*chanCall{ + "A": { + controls: []string{"B"}, + }, + "B": { + controls: []string{"C"}, + }, + "C": { + controls: []string{"A"}, + }, + }, + expected: [][]string{{"A", "B", "C", "A"}}, + }, + { + name: "Graph with multiple cycles", + startNodes: []string{"A", "B", "C", "D", "E", "F"}, + chanCalls: map[string]*chanCall{ + "A": { + controls: []string{"B", "D"}, + }, + "B": { + controls: []string{"C"}, + }, + "C": { + controls: []string{"B"}, + }, + "D": { + controls: []string{"E"}, + }, + "E": { + controls: []string{"F"}, + }, + "F": { + controls: []string{"D"}, + }, + }, + expected: [][]string{{"B", "C", "B"}, {"D", "E", "F", "D"}}, + }, + { + name: "Graph with branch cycle", + startNodes: []string{"A", "C"}, + chanCalls: map[string]*chanCall{ + "A": { + controls: []string{"B"}, + writeToBranches: []*GraphBranch{ + { + endNodes: map[string]bool{ + "C": true, + }, + }, + }, + }, + "B": { + controls: []string{}, + }, + "C": { + controls: []string{"A"}, + }, + }, + expected: [][]string{{"A", "C", "A"}}, + }, + { + name: "Empty graph", + startNodes: []string{}, + chanCalls: map[string]*chanCall{}, + expected: [][]string{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + loops := findLoops(tt.startNodes, tt.chanCalls) + + assert.Equal(t, len(tt.expected), len(loops)) + + if len(tt.expected) > 0 { + normalizedExpected := normalizeLoops(tt.expected) + normalizedActual := normalizeLoops(loops) + assert.Equal(t, normalizedExpected, normalizedActual) + } + }) + } +} +func normalizeLoops(loops [][]string) []string { + result := make([]string, 0, len(loops)) + + for _, loop := range loops { + if len(loop) == 0 { + continue + } + + normalizedLoop := make([]string, len(loop)) + copy(normalizedLoop, loop) + if normalizedLoop[0] != normalizedLoop[len(normalizedLoop)-1] { + normalizedLoop = append(normalizedLoop, normalizedLoop[0]) + } + + minIdx := 0 + for i := 1; i < len(normalizedLoop)-1; i++ { + if normalizedLoop[i] < normalizedLoop[minIdx] { + minIdx = i + } + } + + canonicalLoop := "" + for i := 0; i < len(normalizedLoop)-1; i++ { + idx := (minIdx + i) % (len(normalizedLoop) - 1) + canonicalLoop += normalizedLoop[idx] + "," + } + canonicalLoop += normalizedLoop[minIdx] + + result = append(result, canonicalLoop) + } + + sort.Strings(result) + return result +} + +func TestPrintTasks(t *testing.T) { + var ts []*task + assert.Equal(t, "[]", printTask(ts)) + ts = []*task{{nodeKey: "1"}} + assert.Equal(t, "[1]", printTask(ts)) + ts = []*task{{nodeKey: "1"}, {nodeKey: "2"}, {nodeKey: "3"}} + assert.Equal(t, "[1, 2, 3]", printTask(ts)) +} + +func TestSkipBranch(t *testing.T) { + g := NewGraph[string, string]() + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + _ = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + _ = g.AddEdge(START, "1") + _ = g.AddBranch("1", NewGraphMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) { + return map[string]bool{}, nil + }, map[string]bool{"2": true})) + _ = g.AddEdge("2", END) + + ctx := context.Background() + r, err := g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + _, err = r.Invoke(ctx, "input") + assert.ErrorContains(t, err, "[GraphRunError] no tasks to execute, last completed nodes: [1]") + + g = NewGraph[string, string]() + _ = g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + _ = g.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })) + _ = g.AddEdge(START, "1") + _ = g.AddBranch("1", NewGraphMultiBranch(func(ctx context.Context, in string) (endNode map[string]bool, err error) { + return map[string]bool{}, nil + }, map[string]bool{"2": true})) + _ = g.AddEdge("2", END) + _ = g.AddEdge(START, "2") + r, err = g.Compile(ctx, WithNodeTriggerMode(AllPredecessor)) + assert.NoError(t, err) + result, err := r.Invoke(ctx, "input") + assert.NoError(t, err) + assert.Equal(t, "input", result) +} + +func TestGetStateInGraphCallback(t *testing.T) { + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (s *state) { + return &state{} + })) + assert.NoError(t, g.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }))) + assert.NoError(t, g.AddEdge(START, "1")) + assert.NoError(t, g.AddEdge("1", END)) + + ctx := context.Background() + r, err := g.Compile(ctx) + assert.NoError(t, err) + + _, err = r.Invoke(ctx, "input", WithCallbacks(&testGraphStateCallbackHandler{t: t})) + assert.NoError(t, err) +} + +type state struct { + A string +} + +type testGraphStateCallbackHandler struct { + t *testing.T +} + +func (t *testGraphStateCallbackHandler) OnStart(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + assert.NoError(t.t, ProcessState[*state](ctx, func(ctx context.Context, s *state) error { + s.A = "test" + return nil + })) + return ctx +} + +func (t *testGraphStateCallbackHandler) OnEnd(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + return ctx +} + +func (t *testGraphStateCallbackHandler) OnError(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + return ctx +} + +func (t *testGraphStateCallbackHandler) OnStartWithStreamInput(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + return ctx +} + +func (t *testGraphStateCallbackHandler) OnEndWithStreamOutput(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + return ctx +} + +func TestUniqueSlice(t *testing.T) { + assert.Equal(t, []string{"a", "b", "c"}, uniqueSlice([]string{"a", "b", "a", "c", "b"})) + assert.Equal(t, []string{}, uniqueSlice([]string{})) +} diff --git a/compose/interrupt.go b/compose/interrupt.go new file mode 100644 index 0000000..98a5eee --- /dev/null +++ b/compose/interrupt.go @@ -0,0 +1,363 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + + "github.com/cloudwego/eino/internal/core" + "github.com/cloudwego/eino/schema" +) + +// WithInterruptBeforeNodes instructs to interrupt before the given nodes. +func WithInterruptBeforeNodes(nodes []string) GraphCompileOption { + return func(options *graphCompileOptions) { + options.interruptBeforeNodes = nodes + } +} + +// WithInterruptAfterNodes instructs to interrupt after the given nodes. +func WithInterruptAfterNodes(nodes []string) GraphCompileOption { + return func(options *graphCompileOptions) { + options.interruptAfterNodes = nodes + } +} + +// Deprecated: prefer Interrupt/StatefulInterrupt and CompositeInterrupt. +// If you need to pass the legacy error into CompositeInterrupt, wrap it using WrapInterruptAndRerunIfNeeded first. +var InterruptAndRerun = deprecatedInterruptAndRerun +var deprecatedInterruptAndRerun = errors.New("interrupt and rerun") + +// NewInterruptAndRerunErr creates a legacy interrupt-and-rerun error. +// Deprecated: prefer Interrupt(ctx, info) or StatefulInterrupt(ctx, info, state). +// If passing into CompositeInterrupt, wrap using WrapInterruptAndRerunIfNeeded first. +func NewInterruptAndRerunErr(extra any) error { + return deprecatedInterruptAndRerunErr(extra) +} +func deprecatedInterruptAndRerunErr(extra any) error { + return &core.InterruptSignal{InterruptInfo: core.InterruptInfo{ + Info: extra, + IsRootCause: true, + }} +} + +type wrappedInterruptAndRerun struct { + ps Address + inner error +} + +func (w *wrappedInterruptAndRerun) Error() string { + return fmt.Sprintf("interrupt and rerun at address %s: %s", w.ps.String(), w.inner.Error()) +} + +func (w *wrappedInterruptAndRerun) Unwrap() error { + return w.inner +} + +// WrapInterruptAndRerunIfNeeded wraps the deprecated old interrupt errors, with the current execution address. +// If the error is returned by either Interrupt, StatefulInterrupt or CompositeInterrupt, +// it will be returned as-is without wrapping +func WrapInterruptAndRerunIfNeeded(ctx context.Context, step AddressSegment, err error) error { + addr := GetCurrentAddress(ctx) + newAddr := append(append([]AddressSegment{}, addr...), step) + if errors.Is(err, deprecatedInterruptAndRerun) { + return &wrappedInterruptAndRerun{ + ps: newAddr, + inner: err, + } + } + + ire := &core.InterruptSignal{} + if errors.As(err, &ire) { + if ire.Address == nil { + return &wrappedInterruptAndRerun{ + ps: newAddr, + inner: err, + } + } + return ire + } + + return fmt.Errorf("failed to wrap error as addressed InterruptAndRerun: %w", err) +} + +// Interrupt creates a special error that signals the execution engine to interrupt +// the current run at the component's specific address and save a checkpoint. +// +// This is the standard way for a single, non-composite component to signal a resumable interruption. +// +// - ctx: The context of the running component, used to retrieve the current execution address. +// - info: User-facing information about the interrupt. This is not persisted but is exposed to the +// calling application via the InterruptCtx to provide context (e.g., a reason for the pause). +func Interrupt(ctx context.Context, info any) error { + is, err := core.Interrupt(ctx, info, nil, nil) + if err != nil { + return err + } + + return is +} + +// StatefulInterrupt creates a special error that signals the execution engine to interrupt +// the current run at the component's specific address and save a checkpoint. +// +// This is the standard way for a single, non-composite component to signal a resumable interruption. +// +// - ctx: The context of the running component, used to retrieve the current execution address. +// - info: User-facing information about the interrupt. This is not persisted but is exposed to the +// calling application via the InterruptCtx to provide context (e.g., a reason for the pause). +// - state: The internal state that the interrupting component needs to persist to be able to resume +// its work later. This state is saved in the checkpoint and will be provided back to the component +// upon resumption via GetInterruptState. +func StatefulInterrupt(ctx context.Context, info any, state any) error { + is, err := core.Interrupt(ctx, info, state, nil) + if err != nil { + return err + } + + return is +} + +// CompositeInterrupt creates a special error that signals a composite interruption. +// It is designed for "composite" nodes (like ToolsNode) that manage multiple, independent, +// interruptible sub-processes. It bundles multiple sub-interrupt errors into a single error +// that the engine can deconstruct into a flat list of resumable points. +// +// This function is robust and can handle several types of errors from sub-processes: +// +// - A `Interrupt` or `StatefulInterrupt` error from a simple component. +// +// - A nested `CompositeInterrupt` error from another composite component. +// +// - An error containing `InterruptInfo` returned by a `Runnable` (e.g., a Graph within a lambda node). +// +// - An error returned by \'WrapInterruptAndRerunIfNeeded\' for the legacy old interrupt and rerun error, +// and for the error returned by the deprecated old interrupt errors. +// +// Parameters: +// +// - ctx: The context of the running composite node. +// +// - info: User-facing information for the composite node itself. Can be nil. +// This info will be attached to InterruptInfo.RerunNodeExtra. +// Provided mainly for compatibility purpose as the composite node itself +// is not an interrupt point with interrupt ID, +// which means it lacks enough reason to give a user-facing info. +// +// - state: The state for the composite node itself. Can be nil. +// This could be useful when the composite node needs to restore state, +// such as its input (e.g. ToolsNode). +// +// - errs: a list of errors emitted by sub-processes. +// +// NOTE: if the error you passed in is the deprecated old interrupt and rerun err, or an error returned by +// the deprecated old interrupt function, you must wrap it using WrapInterruptAndRerunIfNeeded first +// before passing them into this function. +func CompositeInterrupt(ctx context.Context, info any, state any, errs ...error) error { + if len(errs) == 0 { + return StatefulInterrupt(ctx, info, state) + } + + var cErrs []*core.InterruptSignal + for _, err := range errs { + wrapped := &wrappedInterruptAndRerun{} + if errors.As(err, &wrapped) { + inner := wrapped.Unwrap() + if errors.Is(inner, deprecatedInterruptAndRerun) { + id := uuid.NewString() + cErrs = append(cErrs, &core.InterruptSignal{ + ID: id, + Address: wrapped.ps, + InterruptInfo: core.InterruptInfo{ + Info: nil, + IsRootCause: true, + }, + }) + continue + } + + ire := &core.InterruptSignal{} + if errors.As(err, &ire) { + id := uuid.NewString() + cErrs = append(cErrs, &core.InterruptSignal{ + ID: id, + Address: wrapped.ps, + InterruptInfo: core.InterruptInfo{ + Info: ire.InterruptInfo.Info, + IsRootCause: ire.InterruptInfo.IsRootCause, + }, + InterruptState: core.InterruptState{ + State: ire.InterruptState.State, + }, + }) + } + + continue + } + + ire := &core.InterruptSignal{} + if errors.As(err, &ire) { + cErrs = append(cErrs, ire) + continue + } + + ie := &interruptError{} + if errors.As(err, &ie) { + is := core.FromInterruptContexts(ie.Info.InterruptContexts) + cErrs = append(cErrs, is) + continue + } + + return fmt.Errorf("composite interrupt but one of the sub error is not interrupt and rerun error: %w", err) + } + + is, err := core.Interrupt(ctx, info, state, cErrs) + if err != nil { + return err + } + return is +} + +// IsInterruptRerunError reports whether the error represents an interrupt-and-rerun +// and returns any attached info. +func IsInterruptRerunError(err error) (any, bool) { + info, _, ok := isInterruptRerunError(err) + return info, ok +} + +func isInterruptRerunError(err error) (info any, state any, ok bool) { + if errors.Is(err, deprecatedInterruptAndRerun) { + return nil, nil, true + } + ire := &core.InterruptSignal{} + if errors.As(err, &ire) { + return ire.Info, ire.State, true + } + return nil, nil, false +} + +// InterruptInfo aggregates interrupt metadata for composite or nested runs. +type InterruptInfo struct { + State any + BeforeNodes []string + AfterNodes []string + RerunNodes []string + RerunNodesExtra map[string]any + SubGraphs map[string]*InterruptInfo + InterruptContexts []*InterruptCtx +} + +func init() { + schema.RegisterName[*InterruptInfo]("_eino_compose_interrupt_info") +} + +// AddressSegmentType defines the type of a segment in an execution address. +type AddressSegmentType = core.AddressSegmentType + +const ( + // AddressSegmentNode represents a segment of an address that corresponds to a graph node. + AddressSegmentNode AddressSegmentType = "node" + // AddressSegmentTool represents a segment of an address that corresponds to a specific tool call within a ToolsNode. + AddressSegmentTool AddressSegmentType = "tool" + // AddressSegmentRunnable represents a segment of an address that corresponds to an instance of the Runnable interface. + // Currently the possible Runnable types are: Graph, Workflow and Chain. + // Note that for sub-graphs added through AddGraphNode to another graph is not a Runnable. + // So a AddressSegmentRunnable indicates a standalone Root level Graph, + // or a Root level Graph inside a node such as Lambda node. + AddressSegmentRunnable AddressSegmentType = "runnable" +) + +// Address represents a full, hierarchical address to a point in the execution structure. +type Address = core.Address + +// AddressSegment represents a single segment in the hierarchical address of an execution point. +// A sequence of AddressSegments uniquely identifies a location within a potentially nested structure. +type AddressSegment = core.AddressSegment + +// InterruptCtx provides a complete, user-facing context for a single, resumable interrupt point. +type InterruptCtx = core.InterruptCtx + +// ExtractInterruptInfo extracts InterruptInfo from an error if present. +func ExtractInterruptInfo(err error) (info *InterruptInfo, existed bool) { + if err == nil { + return nil, false + } + var iE *interruptError + if errors.As(err, &iE) { + return iE.Info, true + } + var sIE *subGraphInterruptError + if errors.As(err, &sIE) { + return sIE.Info, true + } + return nil, false +} + +type interruptError struct { + Info *InterruptInfo +} + +func (e *interruptError) Error() string { + return fmt.Sprintf("interrupt happened, info: %+v", e.Info) +} + +func (e *interruptError) GetInterruptContexts() []*InterruptCtx { + if e.Info == nil { + return nil + } + return e.Info.InterruptContexts +} + +func isSubGraphInterrupt(err error) *subGraphInterruptError { + if err == nil { + return nil + } + var iE *subGraphInterruptError + if errors.As(err, &iE) { + return iE + } + return nil +} + +type subGraphInterruptError struct { + Info *InterruptInfo + CheckPoint *checkpoint + + signal *core.InterruptSignal +} + +func (e *subGraphInterruptError) Error() string { + return fmt.Sprintf("interrupt happened, info: %+v", e.Info) +} + +func isInterruptError(err error) bool { + if _, ok := ExtractInterruptInfo(err); ok { + return true + } + if info := isSubGraphInterrupt(err); info != nil { + return true + } + if _, ok := IsInterruptRerunError(err); ok { + return true + } + + return false +} diff --git a/compose/introspect.go b/compose/introspect.go new file mode 100644 index 0000000..f15d2f7 --- /dev/null +++ b/compose/introspect.go @@ -0,0 +1,57 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "reflect" + + "github.com/cloudwego/eino/components" +) + +// GraphNodeInfo the info which end users pass in when they are adding nodes to graph. +type GraphNodeInfo struct { + Component components.Component + Instance any + GraphAddNodeOpts []GraphAddNodeOpt + InputType, OutputType reflect.Type // mainly for lambda, whose input and output types cannot be inferred by component type + Name string + InputKey, OutputKey string + GraphInfo *GraphInfo + Mappings []*FieldMapping +} + +// GraphInfo the info which end users pass in when they are compiling a graph. +// it is used in compile callback for user to get the node info and instance. +// you may need all details info of the graph for observation. +type GraphInfo struct { + CompileOptions []GraphCompileOption + Nodes map[string]GraphNodeInfo // node key -> node info + Edges map[string][]string // edge start node key -> edge end node key, control edges + DataEdges map[string][]string + Branches map[string][]GraphBranch // branch start node key -> branch + InputType, OutputType reflect.Type + Name string + + NewGraphOptions []NewGraphOption + GenStateFn func(context.Context) any +} + +// GraphCompileCallback is the callback which will be called when graph compilation finishes. +type GraphCompileCallback interface { + OnFinish(ctx context.Context, info *GraphInfo) +} diff --git a/compose/pregel.go b/compose/pregel.go new file mode 100644 index 0000000..583dec9 --- /dev/null +++ b/compose/pregel.go @@ -0,0 +1,95 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import "fmt" + +func pregelChannelBuilder(_ []string, _ []string, _ func() any, _ func() streamReader) channel { + return &pregelChannel{Values: make(map[string]any)} +} + +type pregelChannel struct { + Values map[string]any + + mergeConfig FanInMergeConfig +} + +func (ch *pregelChannel) setMergeConfig(cfg FanInMergeConfig) { + ch.mergeConfig.StreamMergeWithSourceEOF = cfg.StreamMergeWithSourceEOF +} + +func (ch *pregelChannel) load(c channel) error { + dc, ok := c.(*pregelChannel) + if !ok { + return fmt.Errorf("load pregel channel fail, got %T, want *pregelChannel", c) + } + ch.Values = dc.Values + return nil +} + +func (ch *pregelChannel) convertValues(fn func(map[string]any) error) error { + return fn(ch.Values) +} + +func (ch *pregelChannel) reportValues(ins map[string]any) error { + for k, v := range ins { + ch.Values[k] = v + } + return nil +} + +func (ch *pregelChannel) get(isStream bool, name string, edgeHandler *edgeHandlerManager) ( + any, bool, error) { + if len(ch.Values) == 0 { + return nil, false, nil + } + defer func() { ch.Values = map[string]any{} }() + values := make([]any, len(ch.Values)) + names := make([]string, len(ch.Values)) + i := 0 + for k, v := range ch.Values { + resolvedV, err := edgeHandler.handle(k, name, v, isStream) + if err != nil { + return nil, false, err + } + values[i] = resolvedV + names[i] = k + i++ + } + + if len(values) == 1 { + return values[0], true, nil + } + + // merge + mergeOpts := &mergeOptions{ + streamMergeWithSourceEOF: ch.mergeConfig.StreamMergeWithSourceEOF, + names: names, + } + v, err := mergeValues(values, mergeOpts) + if err != nil { + return nil, false, err + } + return v, true, nil +} + +func (ch *pregelChannel) reportSkip(_ []string) bool { + return false +} +func (ch *pregelChannel) reportDependencies(_ []string) { + return +} diff --git a/compose/resume.go b/compose/resume.go new file mode 100644 index 0000000..a4513fc --- /dev/null +++ b/compose/resume.go @@ -0,0 +1,156 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + + "github.com/cloudwego/eino/internal/core" +) + +// GetInterruptState provides a type-safe way to check for and retrieve the persisted state from a previous interruption. +// It is the primary function a component should use to understand its past state. +// +// It returns three values: +// - wasInterrupted (bool): True if the node was part of a previous interruption, regardless of whether state was provided. +// - state (T): The typed state object, if it was provided and matches type `T`. +// - hasState (bool): True if state was provided during the original interrupt and successfully cast to type `T`. +func GetInterruptState[T any](ctx context.Context) (wasInterrupted bool, hasState bool, state T) { + return core.GetInterruptState[T](ctx) +} + +// GetResumeContext checks if the current component is the target of a resume operation +// and retrieves any data provided by the user for that resumption. +// +// This function is typically called *after* a component has already determined it is in a +// resumed state by calling GetInterruptState. +// +// It returns three values: +// - isResumeFlow: A boolean that is true if the current component's address was explicitly targeted +// by a call to Resume() or ResumeWithData(). +// - hasData: A boolean that is true if data was provided for this component (i.e., not nil). +// - data: The typed data provided by the user. +// +// ### How to Use This Function: A Decision Framework +// +// The correct usage pattern depends on the application's desired resume strategy. +// +// #### Strategy 1: Implicit "Resume All" +// In some use cases, any resume operation implies that *all* interrupted points should proceed. +// For example, if an application's UI only provides a single "Continue" button for a set of +// interruptions. In this model, a component can often just use `GetInterruptState` to see if +// `wasInterrupted` is true and then proceed with its logic, as it can assume it is an intended target. +// It may still call `GetResumeContext` to check for optional data, but the `isResumeFlow` flag is less critical. +// +// #### Strategy 2: Explicit "Targeted Resume" (Most Common) +// For applications with multiple, distinct interrupt points that must be resumed independently, it is +// crucial to differentiate which point is being resumed. This is the primary use case for the `isResumeFlow` flag. +// - If `isResumeFlow` is `true`: Your component is the explicit target. You should consume +// the `data` (if any) and complete your work. +// - If `isResumeFlow` is `false`: Another component is the target. You MUST re-interrupt +// (e.g., by returning `StatefulInterrupt(...)`) to preserve your state and allow the +// resume signal to propagate. +// +// ### Guidance for Composite Components +// +// Composite components (like `Graph` or other `Runnable`s that contain sub-processes) have a dual role: +// 1. Check for Self-Targeting: A composite component can itself be the target of a resume +// operation, for instance, to modify its internal state. It may call `GetResumeContext` +// to check for data targeted at its own address. +// 2. Act as a Conduit: After checking for itself, its primary role is to re-execute its children, +// allowing the resume context to flow down to them. It must not consume a resume signal +// intended for one of its descendants. +func GetResumeContext[T any](ctx context.Context) (isResumeFlow bool, hasData bool, data T) { + return core.GetResumeContext[T](ctx) +} + +// GetCurrentAddress returns the hierarchical address of the currently executing component. +// The address is a sequence of segments, each identifying a structural part of the execution +// like an agent, a graph node, or a tool call. This can be useful for logging or debugging. +func GetCurrentAddress(ctx context.Context) Address { + return core.GetCurrentAddress(ctx) +} + +// Resume prepares a context for an "Explicit Targeted Resume" operation by targeting one or more +// components without providing data. It is a convenience wrapper around BatchResumeWithData. +// +// This is useful when the act of resuming is itself the signal, and no extra data is needed. +// The components at the provided addresses (interrupt IDs) will receive `isResumeFlow = true` +// when they call `GetResumeContext`. +func Resume(ctx context.Context, interruptIDs ...string) context.Context { + resumeData := make(map[string]any, len(interruptIDs)) + for _, addr := range interruptIDs { + resumeData[addr] = nil + } + return BatchResumeWithData(ctx, resumeData) +} + +// ResumeWithData prepares a context to resume a single, specific component with data. +// It is the primary function for the "Explicit Targeted Resume" strategy when data is required. +// It is a convenience wrapper around BatchResumeWithData. +// The `interruptID` parameter is the unique interrupt ID of the target component. +func ResumeWithData(ctx context.Context, interruptID string, data any) context.Context { + return BatchResumeWithData(ctx, map[string]any{interruptID: data}) +} + +// BatchResumeWithData is the core function for preparing a resume context. It injects a map +// of resume targets and their corresponding data into the context. +// +// The `resumeData` map should contain the interrupt IDs (which are the string form of addresses) of the +// components to be resumed as keys. The value can be the resume data for that component, or `nil` +// if no data is needed (equivalent to using `Resume`). +// +// This function is the foundation for the "Explicit Targeted Resume" strategy. Components whose interrupt IDs +// are present as keys in the map will receive `isResumeFlow = true` when they call `GetResumeContext`. +func BatchResumeWithData(ctx context.Context, resumeData map[string]any) context.Context { + return core.BatchResumeWithData(ctx, resumeData) +} + +func getNodePath(ctx context.Context) (*NodePath, bool) { + currentAddress := GetCurrentAddress(ctx) + if len(currentAddress) == 0 { + return nil, false + } + + nodePath := make([]string, 0, len(currentAddress)) + for _, p := range currentAddress { + if p.Type == AddressSegmentRunnable { + nodePath = []string{} + continue + } + + nodePath = append(nodePath, p.ID) + } + + return NewNodePath(nodePath...), len(nodePath) > 0 +} + +// AppendAddressSegment creates a new execution context for a sub-component (e.g., a graph node or a tool call). +// +// It extends the current context's address with a new segment and populates the new context with the +// appropriate interrupt state and resume data for that specific sub-address. +// +// - ctx: The parent context, typically the one passed into the component's Invoke/Stream method. +// - segType: The type of the new address segment (e.g., "node", "tool"). +// - segID: The unique ID for the new address segment. +func AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string) context.Context { + return core.AppendAddressSegment(ctx, segType, segID, "") +} + +func appendToolAddressSegment(ctx context.Context, segID string, subID string) context.Context { + return core.AppendAddressSegment(ctx, AddressSegmentTool, segID, subID) +} diff --git a/compose/resume_test.go b/compose/resume_test.go new file mode 100644 index 0000000..da6056d --- /dev/null +++ b/compose/resume_test.go @@ -0,0 +1,1208 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "encoding/json" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +type myInterruptState struct { + OriginalInput string +} + +type myResumeData struct { + Message string +} + +type resumeTestState struct { + OnStartCalledOnResume bool `json:"on_start_called_on_resume"` + Counter int `json:"counter"` +} + +func init() { + schema.Register[resumeTestState]() +} + +func TestInterruptStateAndResumeForRootGraph(t *testing.T) { + // create a graph with a lambda node + // this lambda node will interrupt with a typed state and an info for end-user + // verify the info thrown by the lambda node + // resume with a structured resume data + // within the lambda node, getRunCtx and verify the state and resume data + g := NewGraph[string, string]() + + lambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + wasInterrupted, hasState, state := GetInterruptState[*myInterruptState](ctx) + if !wasInterrupted { + // First run: interrupt with state + return "", StatefulInterrupt(ctx, + map[string]any{"reason": "scheduled maintenance"}, + &myInterruptState{OriginalInput: input}, + ) + } + + // This is a resumed run. + assert.True(t, hasState) + assert.Equal(t, "initial input", state.OriginalInput) + + isResume, hasData, data := GetResumeContext[*myResumeData](ctx) + assert.True(t, isResume) + assert.True(t, hasData) + assert.Equal(t, "let's continue", data.Message) + + return "Resumed successfully with input: " + state.OriginalInput, nil + }) + + _ = g.AddLambdaNode("lambda", lambda) + _ = g.AddEdge(START, "lambda") + _ = g.AddEdge("lambda", END) + + graph, err := g.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()), WithGraphName("root")) + assert.NoError(t, err) + + // First invocation, which should be interrupted + checkPointID := "test-checkpoint-1" + _, err = graph.Invoke(context.Background(), "initial input", WithCheckPointID(checkPointID)) + + // Verify the interrupt error and extracted info + assert.Error(t, err) + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + assert.NotNil(t, interruptInfo) + + interruptContexts := interruptInfo.InterruptContexts + assert.Equal(t, 1, len(interruptContexts)) + assert.Equal(t, "runnable:root;node:lambda", interruptContexts[0].Address.String()) + assert.Equal(t, map[string]any{"reason": "scheduled maintenance"}, interruptContexts[0].Info) + + // Prepare resume data + ctx := ResumeWithData(context.Background(), interruptContexts[0].ID, + &myResumeData{Message: "let's continue"}) + + // Resume execution + output, err := graph.Invoke(ctx, "", WithCheckPointID(checkPointID)) + + // Verify the final result + assert.NoError(t, err) + assert.Equal(t, "Resumed successfully with input: initial input", output) +} + +func TestProcessStateInOnStartDuringResume(t *testing.T) { + graphOnStartCallCount := 0 + processStateErrorOnResume := error(nil) + + cb := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + if info.Name == "test-process-state-onstart" { + graphOnStartCallCount++ + err := ProcessState[*resumeTestState](ctx, func(ctx context.Context, s *resumeTestState) error { + s.Counter++ + return nil + }) + if graphOnStartCallCount > 1 { + processStateErrorOnResume = err + } + } + return ctx + }). + Build() + + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) *resumeTestState { + return &resumeTestState{} + })) + + lambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + wasInterrupted, _, _ := GetInterruptState[*myInterruptState](ctx) + if !wasInterrupted { + return "", StatefulInterrupt(ctx, + map[string]any{"reason": "test interrupt"}, + &myInterruptState{OriginalInput: input}, + ) + } + + var stateCounter int + err := ProcessState[*resumeTestState](ctx, func(ctx context.Context, s *resumeTestState) error { + stateCounter = s.Counter + return nil + }) + assert.NoError(t, err) + assert.Equal(t, 2, stateCounter, "Counter should be 2 (first run OnStart + resume OnStart)") + + return "success", nil + }) + + _ = g.AddLambdaNode("lambda", lambda) + _ = g.AddEdge(START, "lambda") + _ = g.AddEdge("lambda", END) + + graph, err := g.Compile(context.Background(), + WithCheckPointStore(newInMemoryStore()), + WithGraphName("test-process-state-onstart"), + ) + assert.NoError(t, err) + + checkPointID := "test-checkpoint-process-state" + _, err = graph.Invoke(context.Background(), "test input", WithCheckPointID(checkPointID), WithCallbacks(cb)) + + assert.Error(t, err, "First invocation should return an error") + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt, "Should be an interrupt error") + assert.NotNil(t, interruptInfo) + assert.Equal(t, 1, graphOnStartCallCount, "Graph OnStart should be called once on first run") + + ctx := ResumeWithData(context.Background(), interruptInfo.InterruptContexts[0].ID, &myResumeData{}) + + output, err := graph.Invoke(ctx, "", WithCheckPointID(checkPointID), WithCallbacks(cb)) + assert.NoError(t, err) + assert.Equal(t, "success", output) + assert.Equal(t, 2, graphOnStartCallCount, "Graph OnStart should be called twice (first run + resume)") + assert.NoError(t, processStateErrorOnResume, "ProcessState should work in OnStart during resume") +} + +func TestInterruptStateAndResumeForSubGraph(t *testing.T) { + // create a graph + // create a another graph with a lambda node, as this graph as a sub-graph of the previous graph + // this lambda node will interrupt with a typed state and an info for end-user + // verify the info thrown by the lambda node + // resume with a structured resume data + // within the lambda node, getRunCtx and verify the state and resume data + subGraph := NewGraph[string, string]() + + lambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + wasInterrupted, hasState, state := GetInterruptState[*myInterruptState](ctx) + if !wasInterrupted { + // First run: interrupt with state + return "", StatefulInterrupt(ctx, + map[string]any{"reason": "sub-graph maintenance"}, + &myInterruptState{OriginalInput: input}, + ) + } + + // Second (resumed) run + assert.True(t, hasState) + assert.Equal(t, "main input", state.OriginalInput) + + isResume, hasData, data := GetResumeContext[*myResumeData](ctx) + assert.True(t, isResume) + assert.True(t, hasData) + assert.Equal(t, "let's continue sub-graph", data.Message) + + return "Sub-graph resumed successfully", nil + }) + + _ = subGraph.AddLambdaNode("inner_lambda", lambda) + _ = subGraph.AddEdge(START, "inner_lambda") + _ = subGraph.AddEdge("inner_lambda", END) + + // Create the main graph + mainGraph := NewGraph[string, string]() + _ = mainGraph.AddGraphNode("sub_graph_node", subGraph) + _ = mainGraph.AddEdge(START, "sub_graph_node") + _ = mainGraph.AddEdge("sub_graph_node", END) + + compiledMainGraph, err := mainGraph.Compile(context.Background(), WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + // First invocation, which should be interrupted + checkPointID := "test-subgraph-checkpoint-1" + _, err = compiledMainGraph.Invoke(context.Background(), "main input", WithCheckPointID(checkPointID)) + + // Verify the interrupt error and extracted info + assert.Error(t, err) + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + assert.NotNil(t, interruptInfo) + + interruptContexts := interruptInfo.InterruptContexts + assert.Equal(t, 1, len(interruptContexts)) + assert.Equal(t, "runnable:;node:sub_graph_node;node:inner_lambda", interruptContexts[0].Address.String()) + assert.Equal(t, map[string]any{"reason": "sub-graph maintenance"}, interruptContexts[0].Info) + + // Prepare resume data + ctx := ResumeWithData(context.Background(), interruptContexts[0].ID, + &myResumeData{Message: "let's continue sub-graph"}) + + // Resume execution + output, err := compiledMainGraph.Invoke(ctx, "", WithCheckPointID(checkPointID)) + + // Verify the final result + assert.NoError(t, err) + assert.Equal(t, "Sub-graph resumed successfully", output) +} + +func TestInterruptStateAndResumeForToolInNestedSubGraph(t *testing.T) { + // create a ROOT graph. + // create a sub graph A, add A to ROOT graph using AddGraphNode. + // create a sub-sub graph B, add B to A using AddGraphNode. + // within sub-sub graph B, add a ChatModelNode, which is a Mock chat model that implements the ToolCallingChatModel + // interface. + // add a Mock InvokableTool to this mock chat model. + // within sub-sub graph B, also add a ToolsNode that will execute this Mock InvokableTool. + // this tool will interrupt with a typed state and an info for end-user + // verify the info thrown by the tool. + // resume with a structured resume data. + // within the Tool, getRunCtx and verify the state and resume data + ctrl := gomock.NewController(t) + + // 1. Define the interrupting tool + mockTool := &mockInterruptingTool{tt: t} + + // 2. Define the sub-sub-graph (B) + subSubGraphB := NewGraph[[]*schema.Message, []*schema.Message]() + + // Mock Chat Model that calls the tool + mockChatModel := mockModel.NewMockToolCallingChatModel(ctrl) + mockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "tool_call_123", Function: schema.FunctionCall{Name: "interrupt_tool", Arguments: `{"input": "test"}`}}, + }, + }, nil).AnyTimes() + mockChatModel.EXPECT().WithTools(gomock.Any()).Return(mockChatModel, nil).AnyTimes() + + toolsNode, err := NewToolNode(context.Background(), &ToolsNodeConfig{Tools: []tool.BaseTool{mockTool}}) + assert.NoError(t, err) + + _ = subSubGraphB.AddChatModelNode("model", mockChatModel) + _ = subSubGraphB.AddToolsNode("tools", toolsNode) + _ = subSubGraphB.AddEdge(START, "model") + _ = subSubGraphB.AddEdge("model", "tools") + _ = subSubGraphB.AddEdge("tools", END) + + // 3. Define sub-graph (A) + subGraphA := NewGraph[[]*schema.Message, []*schema.Message]() + _ = subGraphA.AddGraphNode("sub_graph_b", subSubGraphB) + _ = subGraphA.AddEdge(START, "sub_graph_b") + _ = subGraphA.AddEdge("sub_graph_b", END) + + // 4. Define root graph + rootGraph := NewGraph[[]*schema.Message, []*schema.Message]() + _ = rootGraph.AddGraphNode("sub_graph_a", subGraphA) + _ = rootGraph.AddEdge(START, "sub_graph_a") + _ = rootGraph.AddEdge("sub_graph_a", END) + + // 5. Compile and run + compiledRootGraph, err := rootGraph.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()), + WithGraphName("root")) + assert.NoError(t, err) + + // First invocation - should interrupt + checkPointID := "test-nested-tool-interrupt" + initialInput := []*schema.Message{schema.UserMessage("hello")} + _, err = compiledRootGraph.Invoke(context.Background(), initialInput, WithCheckPointID(checkPointID)) + + // 6. Verify the interrupt + assert.Error(t, err) + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + assert.NotNil(t, interruptInfo) + + interruptContexts := interruptInfo.InterruptContexts + assert.Len(t, interruptContexts, 1) // Only the root cause is returned + + // Verify the root cause context + rootCause := interruptContexts[0] + expectedPath := "runnable:root;node:sub_graph_a;node:sub_graph_b;node:tools;tool:interrupt_tool:tool_call_123" + assert.Equal(t, expectedPath, rootCause.Address.String()) + assert.True(t, rootCause.IsRootCause) + assert.Equal(t, map[string]any{"reason": "tool maintenance"}, rootCause.Info) + + // Verify the parent via the Parent field + assert.NotNil(t, rootCause.Parent) + assert.Equal(t, "runnable:root;node:sub_graph_a;node:sub_graph_b;node:tools", rootCause.Parent.Address.String()) + assert.False(t, rootCause.Parent.IsRootCause) + + // 7. Resume execution + ctx := ResumeWithData(context.Background(), rootCause.ID, &myResumeData{Message: "let's continue tool"}) + output, err := compiledRootGraph.Invoke(ctx, initialInput, WithCheckPointID(checkPointID)) + + // 8. Verify final result + assert.NoError(t, err) + assert.NotNil(t, output) + assert.Len(t, output, 1) + assert.Equal(t, "Tool resumed successfully", output[0].Content) +} + +const PathSegmentTypeProcess AddressSegmentType = "process" + +// processState is the state for a single sub-process in the batch test. +type processState struct { + Step int +} + +// batchState is the composite state for the whole batch lambda. +type batchState struct { + ProcessStates map[string]*processState + Results map[string]string +} + +type processResumeData struct { + Instruction string +} + +func init() { + schema.RegisterName[*myInterruptState]("my_interrupt_state") + schema.RegisterName[*batchState]("batch_state") + schema.RegisterName[*processState]("process_state") +} + +func TestMultipleInterruptsAndResumes(t *testing.T) { + // define a new lambda node that act as a 'batch' node + // it kick starts 3 parallel processes, each will interrupt on first run, while preserving their own state. + // each of the process should have their own user-facing interrupt info. + // define a new AddressSegmentType for these sub processes. + // the lambda should use StatefulInterrupt to interrupt and preserve the state, + // which is a specific struct type that implements the CompositeInterruptState interface. + // there should also be a specific struct that that implements the CompositeInterruptInfo interface, + // which helps the end-user to fetch the nested interrupt info. + // put this lambda node within a graph and invoke the graph. + // simulate the user getting the flat list of 3 interrupt points using GetInterruptContexts + // the user then decides to resume two of the three interrupt points + // the first resume has resume data, while the second resume does not.(ResumeWithData vs. Resume) + // verify the resume data and state for the resumed interrupt points. + processIDs := []string{"p0", "p1", "p2"} + + // This is the logic for a single "process" + runProcess := func(ctx context.Context, id string) (string, error) { + // Check if this specific process was interrupted before + wasInterrupted, hasState, pState := GetInterruptState[*processState](ctx) + if !wasInterrupted { + // First run for this process, interrupt it. + return "", StatefulInterrupt(ctx, + map[string]any{"reason": "process " + id + " needs input"}, + &processState{Step: 1}, + ) + } + + assert.True(t, hasState) + assert.Equal(t, 1, pState.Step) + + // Check if we are being resumed + isResume, hasData, pData := GetResumeContext[*processResumeData](ctx) + if !isResume { + // Not being resumed, so interrupt again. + return "", StatefulInterrupt(ctx, + map[string]any{"reason": "process " + id + " still needs input"}, + pState, + ) + } + + // We are being resumed. + if hasData { + // Resumed with data + return "process " + id + " done with instruction: " + pData.Instruction, nil + } + // Resumed without data + return "process " + id + " done", nil + } + + // This is the main "batch" lambda that orchestrates the processes + batchLambda := InvokableLambda(func(ctx context.Context, _ string) (map[string]string, error) { + // Restore the state of the batch node itself + _, _, persistedBatchState := GetInterruptState[*batchState](ctx) + if persistedBatchState == nil { + persistedBatchState = &batchState{ + Results: make(map[string]string), + } + } + + var errs []error + + for _, id := range processIDs { + // If this process already completed in a previous run, skip it. + if _, done := persistedBatchState.Results[id]; done { + continue + } + + // Create a sub-context for each process + subCtx := AppendAddressSegment(ctx, PathSegmentTypeProcess, id) + res, err := runProcess(subCtx, id) + + if err != nil { + _, ok := IsInterruptRerunError(err) + assert.True(t, ok) + errs = append(errs, err) + } else { + // Process completed, save its result to the state for the next run. + persistedBatchState.Results[id] = res + } + } + + if len(errs) > 0 { + return nil, CompositeInterrupt(ctx, nil, persistedBatchState, errs...) + } + + return persistedBatchState.Results, nil + }) + + g := NewGraph[string, map[string]string]() + _ = g.AddLambdaNode("batch", batchLambda) + _ = g.AddEdge(START, "batch") + _ = g.AddEdge("batch", END) + + graph, err := g.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()), + WithGraphName("root")) + assert.NoError(t, err) + + // --- 1. First invocation, all 3 processes should interrupt --- + checkPointID := "multi-interrupt-test" + _, err = graph.Invoke(context.Background(), "", WithCheckPointID(checkPointID)) + + assert.Error(t, err) + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + interruptContexts := interruptInfo.InterruptContexts + assert.Len(t, interruptContexts, 3) // Only the 3 root causes + + found := make(map[string]bool) + addrToID := make(map[string]string) + var parentCtx *InterruptCtx + for _, iCtx := range interruptContexts { + addrStr := iCtx.Address.String() + found[addrStr] = true + addrToID[addrStr] = iCtx.ID + assert.True(t, iCtx.IsRootCause) + assert.Equal(t, map[string]any{"reason": "process " + iCtx.Address[2].ID + " needs input"}, iCtx.Info) + // Check that all share the same parent + assert.NotNil(t, iCtx.Parent) + if parentCtx == nil { + parentCtx = iCtx.Parent + assert.Equal(t, "runnable:root;node:batch", parentCtx.Address.String()) + assert.False(t, parentCtx.IsRootCause) + } else { + assert.Same(t, parentCtx, iCtx.Parent) + } + } + assert.True(t, found["runnable:root;node:batch;process:p0"]) + assert.True(t, found["runnable:root;node:batch;process:p1"]) + assert.True(t, found["runnable:root;node:batch;process:p2"]) + + // --- 2. Second invocation, resume 2 of 3 processes --- + // Resume p0 with data, and p2 without data. p1 remains interrupted. + resumeCtx := ResumeWithData(context.Background(), addrToID["runnable:root;node:batch;process:p0"], &processResumeData{Instruction: "do it"}) + resumeCtx = Resume(resumeCtx, addrToID["runnable:root;node:batch;process:p2"]) + + _, err = graph.Invoke(resumeCtx, "", WithCheckPointID(checkPointID)) + + // Expect an interrupt again, but only for p1 + assert.Error(t, err) + interruptInfo2, isInterrupt2 := ExtractInterruptInfo(err) + assert.True(t, isInterrupt2) + interruptContexts2 := interruptInfo2.InterruptContexts + assert.Len(t, interruptContexts2, 1) // Only p1 is left + rootCause2 := interruptContexts2[0] + assert.Equal(t, "runnable:root;node:batch;process:p1", rootCause2.Address.String()) + assert.NotNil(t, rootCause2.Parent) + assert.Equal(t, "runnable:root;node:batch", rootCause2.Parent.Address.String()) + + // --- 3. Third invocation, resume the last process --- + finalResumeCtx := Resume(context.Background(), rootCause2.ID) + finalOutput, err := graph.Invoke(finalResumeCtx, "", WithCheckPointID(checkPointID)) + + assert.NoError(t, err) + assert.Equal(t, "process p0 done with instruction: do it", finalOutput["p0"]) + assert.Equal(t, "process p1 done", finalOutput["p1"]) + assert.Equal(t, "process p2 done", finalOutput["p2"]) +} + +// toolsNodeResumeTargetCallback captures isResumeTarget for ToolsNode during OnStart +type toolsNodeResumeTargetCallback struct { + mu sync.Mutex + isResumeTargetLog []bool +} + +func (c *toolsNodeResumeTargetCallback) OnStart(ctx context.Context, info *callbacks.RunInfo, _ callbacks.CallbackInput) context.Context { + if info.Component == ComponentOfToolsNode { + isResumeTarget, _, _ := GetResumeContext[any](ctx) + c.mu.Lock() + c.isResumeTargetLog = append(c.isResumeTargetLog, isResumeTarget) + c.mu.Unlock() + } + return ctx +} + +func (c *toolsNodeResumeTargetCallback) OnEnd(ctx context.Context, _ *callbacks.RunInfo, _ callbacks.CallbackOutput) context.Context { + return ctx +} + +func (c *toolsNodeResumeTargetCallback) OnError(ctx context.Context, _ *callbacks.RunInfo, _ error) context.Context { + return ctx +} + +func (c *toolsNodeResumeTargetCallback) OnStartWithStreamInput(ctx context.Context, _ *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + input.Close() + return ctx +} + +func (c *toolsNodeResumeTargetCallback) OnEndWithStreamOutput(ctx context.Context, _ *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + output.Close() + return ctx +} + +// mockReentryTool is a helper for the reentry test +type mockReentryTool struct { + t *testing.T + mu sync.Mutex + isResumeTargetByRunID map[string]bool +} + +func (t *mockReentryTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "reentry_tool", + Desc: "A tool that can be re-entered in a resumed graph.", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{"input": {Type: schema.String}}), + }, nil +} + +func (t *mockReentryTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + wasInterrupted, hasState, _ := tool.GetInterruptState[any](ctx) + isResume, hasData, data := tool.GetResumeContext[*myResumeData](ctx) + + callID := GetToolCallID(ctx) + + t.mu.Lock() + if t.isResumeTargetByRunID != nil { + t.isResumeTargetByRunID[callID] = isResume + } + t.mu.Unlock() + + // Special handling for the re-entrant call to make assertions explicit. + if callID == "call_3" { + if !isResume { + // This is the first run of the re-entrant call. Its context must be clean. + // This is the core assertion for this test. + assert.False(t.t, wasInterrupted, "re-entrant call 'call_3' should not have been interrupted on its first run") + assert.False(t.t, hasState, "re-entrant call 'call_3' should not have state on its first run") + // Now, interrupt it as part of the test flow. + return "", tool.StatefulInterrupt(ctx, nil, "some state for "+callID) + } + // This is the resumed run of the re-entrant call. + assert.True(t.t, wasInterrupted, "resumed call 'call_3' must have been interrupted") + assert.True(t.t, hasData, "resumed call 'call_3' should have data") + return "Resumed " + data.Message, nil + } + + // Standard logic for the initial calls (call_1, call_2) + if !wasInterrupted { + // First run for call_1 and call_2, should interrupt. + return "", tool.StatefulInterrupt(ctx, nil, "some state for "+callID) + } + + // From here, wasInterrupted is true for call_1 and call_2. + if isResume { + // The user is explicitly resuming this call. + assert.True(t.t, hasData, "call %s should have resume data", callID) + return "Resumed " + data.Message, nil + } + + // The tool was interrupted before, but is not being resumed now. Re-interrupt. + return "", tool.StatefulInterrupt(ctx, nil, "some state for "+callID) +} + +func TestReentryForResumedTools(t *testing.T) { + // create a 'ReAct' style graph with a ChatModel node and a ToolsNode. + // within the ToolsNode there is an interruptible tool that will emit interrupt on first run. + // During the first invocation of the graph, there should be two tool calls (of the same tool) that interrupt. + // The user chooses to resume one of the interrupted tool call in second invocation, + // and this time, the resumed tool call should be successful, while the other should interrupt immediately again. + // The user then chooses to resume the other interrupted tool call in third invocation, + // and this time, the ChatModel decides to call the tool again, + // and this time the tool's runCtx should think it was not interrupted nor resumed. + ctrl := gomock.NewController(t) + + // 1. Define the interrupting tool and callback + reentryTool := &mockReentryTool{t: t, isResumeTargetByRunID: make(map[string]bool)} + toolsNodeCB := &toolsNodeResumeTargetCallback{} + + // 2. Define the graph + g := NewGraph[[]*schema.Message, *schema.Message]() + + // Mock Chat Model that drives the ReAct loop + mockChatModel := mockModel.NewMockToolCallingChatModel(ctrl) + toolsNode, err := NewToolNode(context.Background(), &ToolsNodeConfig{Tools: []tool.BaseTool{reentryTool}}) + assert.NoError(t, err) + + // Expectation for the 1st invocation: model returns two tool calls + mockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "reentry_tool", Arguments: `{"input": "a"}`}}, + {ID: "call_2", Function: schema.FunctionCall{Name: "reentry_tool", Arguments: `{"input": "b"}`}}, + }, + }, nil).Times(1) + + // Expectation for the 2nd invocation (after resuming call_1): model does nothing, graph continues + // Expectation for the 3rd invocation (after resuming call_2): model calls the tool again + mockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, msgs []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_3", Function: schema.FunctionCall{Name: "reentry_tool", Arguments: `{"input": "c"}`}}, + }, + }, nil + }).Times(1) + + // Expectation for the final invocation: model returns final answer + mockChatModel.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).Return(&schema.Message{ + Role: schema.Assistant, + Content: "all done", + }, nil).Times(1) + + _ = g.AddChatModelNode("model", mockChatModel) + _ = g.AddToolsNode("tools", toolsNode) + _ = g.AddEdge(START, "model") + + // Add the crucial branch to decide whether to call tools or end. + modelBranch := func(ctx context.Context, msg *schema.Message) (string, error) { + if len(msg.ToolCalls) > 0 { + return "tools", nil + } + return END, nil + } + err = g.AddBranch("model", NewGraphBranch(modelBranch, map[string]bool{"tools": true, END: true})) + assert.NoError(t, err) + + _ = g.AddEdge("tools", "model") // Loop back for ReAct style + + // 3. Compile and run + graph, err := g.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()), + WithGraphName("root")) + assert.NoError(t, err) + checkPointID := "reentry-test" + + // --- 1. First invocation: call_1 and call_2 should interrupt --- + _, err = graph.Invoke(context.Background(), []*schema.Message{schema.UserMessage("start")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB)) + assert.Error(t, err) + interruptInfo1, _ := ExtractInterruptInfo(err) + interrupts1 := interruptInfo1.InterruptContexts + assert.Len(t, interrupts1, 2) // Only the two tool calls + found1 := make(map[string]bool) + addrToID1 := make(map[string]string) + for _, iCtx := range interrupts1 { + addrStr := iCtx.Address.String() + found1[addrStr] = true + addrToID1[addrStr] = iCtx.ID + assert.True(t, iCtx.IsRootCause) + assert.NotNil(t, iCtx.Parent) + assert.Equal(t, "runnable:root;node:tools", iCtx.Parent.Address.String()) + } + assert.True(t, found1["runnable:root;node:tools;tool:reentry_tool:call_1"]) + assert.True(t, found1["runnable:root;node:tools;tool:reentry_tool:call_2"]) + + // First invocation: neither call_1 nor call_2 should be resume targets + assert.False(t, reentryTool.isResumeTargetByRunID["call_1"], "first run: call_1 should not be resume target") + assert.False(t, reentryTool.isResumeTargetByRunID["call_2"], "first run: call_2 should not be resume target") + + // First invocation: ToolsNode should NOT be a resume target + assert.Len(t, toolsNodeCB.isResumeTargetLog, 1, "ToolsNode OnStart should be called once in first invocation") + assert.False(t, toolsNodeCB.isResumeTargetLog[0], "first run: ToolsNode should NOT be resume target") + + // Clear for next invocation + reentryTool.isResumeTargetByRunID = make(map[string]bool) + toolsNodeCB.isResumeTargetLog = nil + + // --- 2. Second invocation: resume call_1, expect call_2 to interrupt again --- + resumeCtx2 := ResumeWithData(context.Background(), addrToID1["runnable:root;node:tools;tool:reentry_tool:call_1"], + &myResumeData{Message: "resume call 1"}) + _, err = graph.Invoke(resumeCtx2, []*schema.Message{schema.UserMessage("start")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB)) + assert.Error(t, err) + interruptInfo2, _ := ExtractInterruptInfo(err) + interrupts2 := interruptInfo2.InterruptContexts + assert.Len(t, interrupts2, 1) // Only call_2 + rootCause2 := interrupts2[0] + assert.Equal(t, "runnable:root;node:tools;tool:reentry_tool:call_2", rootCause2.Address.String()) + assert.NotNil(t, rootCause2.Parent) + assert.Equal(t, "runnable:root;node:tools", rootCause2.Parent.Address.String()) + + // Second invocation: call_1 is resumed, call_2 is NOT resumed (re-interrupts) + assert.True(t, reentryTool.isResumeTargetByRunID["call_1"], "second run: call_1 should be resume target") + assert.False(t, reentryTool.isResumeTargetByRunID["call_2"], "second run: call_2 should NOT be resume target (it re-interrupts)") + + // Second invocation: ToolsNode SHOULD be a resume target (because call_1 child is being resumed) + assert.Len(t, toolsNodeCB.isResumeTargetLog, 1, "ToolsNode OnStart should be called once in second invocation") + assert.True(t, toolsNodeCB.isResumeTargetLog[0], "second run: ToolsNode SHOULD be resume target (child call_1 is being resumed)") + + // Clear for next invocation + reentryTool.isResumeTargetByRunID = make(map[string]bool) + toolsNodeCB.isResumeTargetLog = nil + + // --- 3. Third invocation: resume call_2, model makes a new call (call_3) which should interrupt --- + resumeCtx3 := ResumeWithData(context.Background(), rootCause2.ID, &myResumeData{Message: "resume call 2"}) + _, err = graph.Invoke(resumeCtx3, []*schema.Message{schema.UserMessage("start")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB)) + assert.Error(t, err) + interruptInfo3, _ := ExtractInterruptInfo(err) + interrupts3 := interruptInfo3.InterruptContexts + assert.Len(t, interrupts3, 1) // Only call_3 + rootCause3 := interrupts3[0] + assert.Equal(t, "runnable:root;node:tools;tool:reentry_tool:call_3", rootCause3.Address.String()) // Note: this is the new call_3 + assert.NotNil(t, rootCause3.Parent) + assert.Equal(t, "runnable:root;node:tools", rootCause3.Parent.Address.String()) + + // Third invocation: call_2 is resumed, call_3 is new (not resumed) + assert.True(t, reentryTool.isResumeTargetByRunID["call_2"], "third run: call_2 should be resume target") + assert.False(t, reentryTool.isResumeTargetByRunID["call_3"], "third run: call_3 should NOT be resume target (it's new)") + + // Third invocation: ToolsNode is called twice (once for call_2 resume, once for call_3 new) + // First call: ToolsNode SHOULD be resume target (call_2 is being resumed) + // Second call: ToolsNode should NOT be resume target (call_3 is new, no children to resume) + assert.Len(t, toolsNodeCB.isResumeTargetLog, 2, "ToolsNode OnStart should be called twice in third invocation") + assert.True(t, toolsNodeCB.isResumeTargetLog[0], "third run first ToolsNode call: SHOULD be resume target (child call_2 is being resumed)") + assert.False(t, toolsNodeCB.isResumeTargetLog[1], "third run second ToolsNode call: should NOT be resume target (call_3 is new)") + + // Clear for next invocation + reentryTool.isResumeTargetByRunID = make(map[string]bool) + toolsNodeCB.isResumeTargetLog = nil + + // --- 4. Final invocation: resume call_3, expect final answer --- + resumeCtx4 := ResumeWithData(context.Background(), rootCause3.ID, + &myResumeData{Message: "resume call 3"}) + output, err := graph.Invoke(resumeCtx4, []*schema.Message{schema.UserMessage("start")}, WithCheckPointID(checkPointID), WithCallbacks(toolsNodeCB)) + assert.NoError(t, err) + assert.Equal(t, "all done", output.Content) + + // Fourth invocation: call_3 is resumed + assert.True(t, reentryTool.isResumeTargetByRunID["call_3"], "fourth run: call_3 should be resume target") + + // Fourth invocation: ToolsNode SHOULD be resume target (call_3 is being resumed) + assert.Len(t, toolsNodeCB.isResumeTargetLog, 1, "ToolsNode OnStart should be called once in fourth invocation") + assert.True(t, toolsNodeCB.isResumeTargetLog[0], "fourth run: ToolsNode SHOULD be resume target (child call_3 is being resumed)") +} + +// mockInterruptingTool is a helper for the nested tool interrupt test +type mockInterruptingTool struct { + tt *testing.T +} + +func (t *mockInterruptingTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "interrupt_tool", + Desc: "A tool that interrupts execution.", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "input": {Type: schema.String, Desc: "Some input", Required: true}, + }), + }, nil +} + +func (t *mockInterruptingTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + var args map[string]string + _ = json.Unmarshal([]byte(argumentsInJSON), &args) + + wasInterrupted, hasState, state := tool.GetInterruptState[*myInterruptState](ctx) + if !wasInterrupted { + // First run: interrupt + return "", tool.StatefulInterrupt(ctx, + map[string]any{"reason": "tool maintenance"}, + &myInterruptState{OriginalInput: args["input"]}, + ) + } + + // Second (resumed) run + assert.True(t.tt, hasState) + assert.Equal(t.tt, "test", state.OriginalInput) + + isResume, hasData, data := tool.GetResumeContext[*myResumeData](ctx) + assert.True(t.tt, isResume) + assert.True(t.tt, hasData) + assert.Equal(t.tt, "let's continue tool", data.Message) + + return "Tool resumed successfully", nil +} + +func TestGraphInterruptWithinLambda(t *testing.T) { + // this test case aims to verify behaviors when a standalone graph is within a lambda, + // which in turn is within the root graph. + // the expected behavior is: + // - internal graph will naturally append to the Address + // - internal graph interrupts, where the Address includes steps for both the root graph and the internal graph + // - lambda extracts InterruptInfo, then GetInterruptContexts + // - lambda then acts as a composite node, uses CompositeInterrupt to pass up the + // internal interrupt points + // - the root graph interrupts + // - end-user extracts the interrupt ID and related info + // - end-user uses ResumeWithData to resume the ID + // - lambda node resumes, invokes the inner graph as usual + // - the internal graph resumes the interrupted node + // To implement this test, within the internal graph you can define another lambda node that can interrupt resume. + + // 1. Define the innermost lambda that actually interrupts + interruptingLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + wasInterrupted, hasState, state := GetInterruptState[*myInterruptState](ctx) + if !wasInterrupted { + return "", StatefulInterrupt(ctx, "inner interrupt", &myInterruptState{OriginalInput: input}) + } + + assert.True(t, hasState) + assert.Equal(t, "top level input", state.OriginalInput) + + isResume, hasData, data := GetResumeContext[*myResumeData](ctx) + assert.True(t, isResume) + assert.True(t, hasData) + assert.Equal(t, "resume inner", data.Message) + + return "inner lambda resumed successfully", nil + }) + + // 2. Define the internal graph that contains the interrupting lambda + innerGraph := NewGraph[string, string]() + _ = innerGraph.AddLambdaNode("inner_lambda", interruptingLambda) + _ = innerGraph.AddEdge(START, "inner_lambda") + _ = innerGraph.AddEdge("inner_lambda", END) + // Give the inner graph a name so it can create its "runnable" addr step. + compiledInnerGraph, err := innerGraph.Compile(context.Background(), WithGraphName("inner"), WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + // 3. Define the outer lambda that acts as a composite node + compositeLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + // The lambda invokes the inner graph. If the inner graph interrupts, this lambda + // must act as a proper composite node and wrap the error. + output, err := compiledInnerGraph.Invoke(ctx, input, WithCheckPointID("inner-cp")) + if err != nil { + _, isInterrupt := ExtractInterruptInfo(err) + if !isInterrupt { + return "", err // Not an interrupt, just fail + } + + // The composite interrupt itself can be stateless, as it's just a wrapper. + // It signals to the framework to look inside the subErrs and correctly + // prepend the current addr to the paths of the inner interrupts. + return "", CompositeInterrupt(ctx, "composite interrupt from lambda", nil, err) + } + return output, nil + }) + + // 4. Define the root graph + rootGraph := NewGraph[string, string]() + _ = rootGraph.AddLambdaNode("composite_lambda", compositeLambda) + _ = rootGraph.AddEdge(START, "composite_lambda") + _ = rootGraph.AddEdge("composite_lambda", END) + // Give the root graph a name for its "runnable" addr step. + compiledRootGraph, err := rootGraph.Compile(context.Background(), WithGraphName("root"), WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + // 5. First invocation - should interrupt + checkPointID := "graph-in-lambda-test" + _, err = compiledRootGraph.Invoke(context.Background(), "top level input", WithCheckPointID(checkPointID)) + + // 6. Verify the interrupt + assert.Error(t, err) + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + interruptContexts := interruptInfo.InterruptContexts + assert.Len(t, interruptContexts, 1) // Only the root cause is returned + + // The addr is now fully qualified, including the runnable steps from both graphs. + rootCause := interruptContexts[0] + expectedPath := "runnable:root;node:composite_lambda;runnable:inner;node:inner_lambda" + assert.Equal(t, expectedPath, rootCause.Address.String()) + assert.Equal(t, "inner interrupt", rootCause.Info) + assert.True(t, rootCause.IsRootCause) + + // Check parent hierarchy + assert.NotNil(t, rootCause.Parent) + assert.Equal(t, "runnable:root;node:composite_lambda;runnable:inner", rootCause.Parent.Address.String()) + assert.Nil(t, rootCause.Parent.Info) // The inner runnable doesn't have its own info + assert.False(t, rootCause.Parent.IsRootCause) + + // Check grandparent + assert.NotNil(t, rootCause.Parent.Parent) + assert.Equal(t, "runnable:root;node:composite_lambda", rootCause.Parent.Parent.Address.String()) + assert.Equal(t, "composite interrupt from lambda", rootCause.Parent.Parent.Info) + assert.False(t, rootCause.Parent.Parent.IsRootCause) + + // 7. Resume execution using the complete, fully-qualified ID + resumeCtx := ResumeWithData(context.Background(), rootCause.ID, &myResumeData{Message: "resume inner"}) + finalOutput, err := compiledRootGraph.Invoke(resumeCtx, "top level input", WithCheckPointID(checkPointID)) + + // 8. Verify final result + assert.NoError(t, err) + assert.Equal(t, "inner lambda resumed successfully", finalOutput) +} + +func TestLegacyInterrupt(t *testing.T) { + // this test case aims to test the behavior of the deprecated InterruptAndRerun, + // NewInterruptAndRerunErr within CompositeInterrupt. + // Define two sub-processes(functions), one interrupts with InterruptAndRerun, + // the other interrupts with NewInterruptAndRerunErr. + // create a lambda as a composite node, within the lambda invokes the two sub-processes. + // create the graph, add lambda node and invoke it. + // after verifying the interrupt points, just invokes again without explicit resume. + // verify the same interrupt IDs again. + // then finally Resume() the graph. + + // 1. Define the sub-processes that use legacy and modern interrupts + subProcess1 := func(ctx context.Context) (string, error) { + isResume, _, data := GetResumeContext[string](ctx) + if isResume { + return data, nil + } + return "", deprecatedInterruptAndRerun + } + subProcess2 := func(ctx context.Context) (string, error) { + isResume, _, data := GetResumeContext[string](ctx) + if isResume { + return data, nil + } + return "", deprecatedInterruptAndRerunErr("legacy info") + } + subProcess3 := func(ctx context.Context) (string, error) { + isResume, _, data := GetResumeContext[string](ctx) + if isResume { + return data, nil + } + // Use the modern, addr-aware interrupt function + return "", Interrupt(ctx, "modern info") + } + + // 2. Define the composite lambda + compositeLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + // If the lambda itself is being resumed, it means the whole process is done. + isResume, _, data := GetResumeContext[string](ctx) + + // Run sub-processes and collect their errors + var ( + errs []error + outStr string + ) + + const PathStepCustom AddressSegmentType = "custom" + subCtx1 := AppendAddressSegment(ctx, PathStepCustom, "1") + out1, err1 := subProcess1(subCtx1) + if err1 != nil { + // Wrap the legacy error to give it a addr + wrappedErr := WrapInterruptAndRerunIfNeeded(ctx, AddressSegment{Type: PathStepCustom, ID: "1"}, err1) + errs = append(errs, wrappedErr) + } else { + outStr += out1 + } + subCtx2 := AppendAddressSegment(ctx, PathStepCustom, "2") + out2, err2 := subProcess2(subCtx2) + if err2 != nil { + // Wrap the legacy error to give it a addr + wrappedErr := WrapInterruptAndRerunIfNeeded(ctx, AddressSegment{Type: PathStepCustom, ID: "2"}, err2) + errs = append(errs, wrappedErr) + } else { + outStr += out2 + } + subCtx3 := AppendAddressSegment(ctx, PathStepCustom, "3") + out3, err3 := subProcess3(subCtx3) + if err3 != nil { + // The error from Interrupt() is already addr-aware. WrapInterruptAndRerunIfNeeded + // should handle this gracefully and return the error as-is. + wrappedErr := WrapInterruptAndRerunIfNeeded(ctx, AddressSegment{Type: PathStepCustom, ID: "3"}, err3) + errs = append(errs, wrappedErr) + } else { + outStr += out3 + } + + if len(errs) > 0 { + // Return a composite interrupt containing the wrapped legacy errors + return "", CompositeInterrupt(ctx, "legacy composite", nil, errs...) + } + + if isResume { + outStr = outStr + " " + data + } + + return outStr, nil + }) + + // 3. Create and compile the graph + rootGraph := NewGraph[string, string]() + _ = rootGraph.AddLambdaNode("legacy_composite", compositeLambda) + _ = rootGraph.AddEdge(START, "legacy_composite") + _ = rootGraph.AddEdge("legacy_composite", END) + compiledGraph, err := rootGraph.Compile(context.Background(), WithGraphName("root"), WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + // 4. First invocation - should interrupt + checkPointID := "legacy-interrupt-test" + _, err = compiledGraph.Invoke(context.Background(), "input", WithCheckPointID(checkPointID)) + + // 5. Verify the three interrupt points + assert.Error(t, err) + info, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + assert.Len(t, info.InterruptContexts, 3) // Only the 3 root causes + + found := make(map[string]any) + addrToID := make(map[string]string) + var parentCtx *InterruptCtx + for _, iCtx := range info.InterruptContexts { + addrStr := iCtx.Address.String() + found[addrStr] = iCtx.Info + addrToID[addrStr] = iCtx.ID + assert.True(t, iCtx.IsRootCause) + // Check parent + assert.NotNil(t, iCtx.Parent) + if parentCtx == nil { + parentCtx = iCtx.Parent + assert.Equal(t, "runnable:root;node:legacy_composite", parentCtx.Address.String()) + assert.Equal(t, "legacy composite", parentCtx.Info) + assert.False(t, parentCtx.IsRootCause) + } else { + assert.Same(t, parentCtx, iCtx.Parent) + } + } + expectedID1 := "runnable:root;node:legacy_composite;custom:1" + expectedID2 := "runnable:root;node:legacy_composite;custom:2" + expectedID3 := "runnable:root;node:legacy_composite;custom:3" + assert.Contains(t, found, expectedID1) + assert.Nil(t, found[expectedID1]) // From InterruptAndRerun + assert.Contains(t, found, expectedID2) + assert.Equal(t, "legacy info", found[expectedID2]) // From NewInterruptAndRerunErr + assert.Contains(t, found, expectedID3) + assert.Equal(t, "modern info", found[expectedID3]) // From Interrupt + + // 6. Second invocation (re-run without resume) - should yield the same interrupts + _, err = compiledGraph.Invoke(context.Background(), "input", WithCheckPointID(checkPointID)) + assert.Error(t, err) + info2, isInterrupt2 := ExtractInterruptInfo(err) + assert.True(t, isInterrupt2) + assert.Len(t, info2.InterruptContexts, 3, "Should have the same number of interrupts on re-run") + + // 7. Third invocation - Resume all three interrupt points with specific data + resumeData := map[string]any{ + addrToID[expectedID1]: "output1", + addrToID[expectedID2]: "output2", + addrToID[expectedID3]: "output3", + } + resumeCtx := BatchResumeWithData(context.Background(), resumeData) + // TODO: The legacy interrupt wrapping does not currently work correctly with BatchResumeWithData. + // The graph re-interrupts instead of completing. This should be fixed in the core framework. + _, err = compiledGraph.Invoke(resumeCtx, "input", WithCheckPointID(checkPointID)) + assert.Error(t, err) +} + +type wrapperToolForTest struct { + compiledGraph Runnable[string, string] + isResumeTargetLog []bool +} + +func (w *wrapperToolForTest) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "wrapperTool", + Desc: "A tool that wraps a nested graph", + }, nil +} + +func (w *wrapperToolForTest) InvokableRun(ctx context.Context, input string, opts ...tool.Option) (string, error) { + isResumeTarget, _, _ := tool.GetResumeContext[any](ctx) + w.isResumeTargetLog = append(w.isResumeTargetLog, isResumeTarget) + + result, err := w.compiledGraph.Invoke(ctx, input) + if err != nil { + if _, ok := ExtractInterruptInfo(err); ok { + return "", tool.CompositeInterrupt(ctx, "wrapper tool interrupt", nil, err) + } + return "", err + } + return result, nil +} + +func TestToolCompositeInterruptWithNestedGraphInterrupt(t *testing.T) { + ctx := context.Background() + + var innerNodeIsResumeTarget bool + subSubGraph := NewGraph[string, string]() + err := subSubGraph.AddLambdaNode("interruptNode", InvokableLambda(func(ctx context.Context, input string) (string, error) { + wasInterrupted, _, _ := GetInterruptState[any](ctx) + if !wasInterrupted { + return "", Interrupt(ctx, "sub-sub graph interrupt info") + } + isResumeTarget, _, _ := GetResumeContext[any](ctx) + innerNodeIsResumeTarget = isResumeTarget + return "resumed successfully", nil + })) + assert.NoError(t, err) + assert.NoError(t, subSubGraph.AddEdge(START, "interruptNode")) + assert.NoError(t, subSubGraph.AddEdge("interruptNode", END)) + + nestedGraph := NewGraph[string, string]() + err = nestedGraph.AddGraphNode("subSubGraph", subSubGraph) + assert.NoError(t, err) + assert.NoError(t, nestedGraph.AddEdge(START, "subSubGraph")) + assert.NoError(t, nestedGraph.AddEdge("subSubGraph", END)) + compiledNestedGraph, err := nestedGraph.Compile(ctx) + assert.NoError(t, err) + + wrapperTool := &wrapperToolForTest{compiledGraph: compiledNestedGraph.(Runnable[string, string])} + + toolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{Tools: []tool.BaseTool{wrapperTool}}) + assert.NoError(t, err) + + outerGraph := NewGraph[*schema.Message, []*schema.Message]() + err = outerGraph.AddToolsNode("tools", toolsNode) + assert.NoError(t, err) + assert.NoError(t, outerGraph.AddEdge(START, "tools")) + assert.NoError(t, outerGraph.AddEdge("tools", END)) + + compiledOuterGraph, err := outerGraph.Compile(ctx, WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + checkpointID := "test-wrapper-tool-resume" + inputMsg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + {ID: "call_1", Function: schema.FunctionCall{Name: "wrapperTool", Arguments: `"test"`}}, + }, + } + + _, err = compiledOuterGraph.Invoke(ctx, inputMsg, WithCheckPointID(checkpointID)) + assert.Error(t, err) + + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok, "should be an interrupt error") + assert.NotNil(t, info) + assert.NotEmpty(t, info.InterruptContexts) + + rootCause := info.InterruptContexts[0] + assert.Equal(t, "sub-sub graph interrupt info", rootCause.Info) + assert.True(t, rootCause.IsRootCause) + + var wrapperToolParent *InterruptCtx + for p := rootCause.Parent; p != nil; p = p.Parent { + if p.Info == "wrapper tool interrupt" { + wrapperToolParent = p + break + } + } + assert.NotNil(t, wrapperToolParent, "should have parent from wrapper tool with info 'wrapper tool interrupt'") + + assert.Len(t, wrapperTool.isResumeTargetLog, 1) + assert.False(t, wrapperTool.isResumeTargetLog[0], "first invocation: wrapper tool should not be resume target") + + resumeCtx := Resume(ctx, rootCause.ID) + _, err = compiledOuterGraph.Invoke(resumeCtx, inputMsg, WithCheckPointID(checkpointID)) + assert.NoError(t, err) + + assert.True(t, innerNodeIsResumeTarget, "inner node should be resume target") + + assert.Len(t, wrapperTool.isResumeTargetLog, 2) + assert.True(t, wrapperTool.isResumeTargetLog[1], "second invocation: wrapper tool should be resume target because its child is targeted") +} diff --git a/compose/runnable.go b/compose/runnable.go new file mode 100644 index 0000000..731484b --- /dev/null +++ b/compose/runnable.go @@ -0,0 +1,539 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +// Runnable is the interface for an executable object. Graph, Chain can be compiled into Runnable. +// runnable is the core conception of eino, we do downgrade compatibility for four data flow patterns, +// and can automatically connect components that only implement one or more methods. +// eg, if a component only implements Stream() method, you can still call Invoke() to convert stream output to invoke output. +type Runnable[I, O any] interface { + Invoke(ctx context.Context, input I, opts ...Option) (output O, err error) + Stream(ctx context.Context, input I, opts ...Option) (output *schema.StreamReader[O], err error) + Collect(ctx context.Context, input *schema.StreamReader[I], opts ...Option) (output O, err error) + Transform(ctx context.Context, input *schema.StreamReader[I], opts ...Option) (output *schema.StreamReader[O], err error) +} + +type invoke func(ctx context.Context, input any, opts ...any) (output any, err error) +type transform func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) + +// composableRunnable the wrapper for all executable object directly provided by the user. +// one instance corresponds to one instance of the executable object. +// all information comes from executable object without any other dimensions of information. +// for the graphNode, ChainBranch, StatePreHandler, StatePostHandler etc. +type composableRunnable struct { + i invoke + t transform + + inputType reflect.Type + outputType reflect.Type + optionType reflect.Type + + *genericHelper + + isPassthrough bool + + meta *executorMeta + + // only available when in Graph node + // if composableRunnable not in Graph node, this field would be nil + nodeInfo *nodeInfo +} + +func runnableLambda[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption], c Collect[I, O, TOption], + t Transform[I, O, TOption], enableCallback bool) *composableRunnable { + rp := newRunnablePacker(i, s, c, t, enableCallback) + + return rp.toComposableRunnable() +} + +func isNilAssignableType(t reflect.Type) bool { + switch t.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice: + return true + default: + return false + } +} + +type runnablePacker[I, O, TOption any] struct { + i Invoke[I, O, TOption] + s Stream[I, O, TOption] + c Collect[I, O, TOption] + t Transform[I, O, TOption] +} + +func (rp *runnablePacker[I, O, TOption]) wrapRunnableCtx(ctxWrapper func(ctx context.Context, opts ...TOption) context.Context) { + i, s, c, t := rp.i, rp.s, rp.c, rp.t + rp.i = func(ctx context.Context, input I, opts ...TOption) (output O, err error) { + ctx = ctxWrapper(ctx, opts...) + return i(ctx, input, opts...) + } + rp.s = func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) { + ctx = ctxWrapper(ctx, opts...) + return s(ctx, input, opts...) + } + rp.c = func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) { + ctx = ctxWrapper(ctx, opts...) + return c(ctx, input, opts...) + } + + rp.t = func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output *schema.StreamReader[O], err error) { + ctx = ctxWrapper(ctx, opts...) + return t(ctx, input, opts...) + } +} + +func (rp *runnablePacker[I, O, TOption]) toComposableRunnable() *composableRunnable { + inputType := generic.TypeOf[I]() + outputType := generic.TypeOf[O]() + optionType := generic.TypeOf[TOption]() + c := &composableRunnable{ + genericHelper: newGenericHelper[I, O](), + inputType: inputType, + outputType: outputType, + optionType: optionType, + } + + i := func(ctx context.Context, input any, opts ...any) (output any, err error) { + in, ok := input.(I) + if !ok { + // When a nil is passed as an 'any' type, its original type information is lost, + // becoming an untyped nil. This would cause type assertions to fail for nil-able + // input types, so recreate the zero value of I when nil is valid for I. + if input == nil && isNilAssignableType(reflect.TypeOf((*I)(nil)).Elem()) { + var i I + in = i + } else { + panic(newUnexpectedInputTypeErr(inputType, reflect.TypeOf(input))) + } + } + + tos, err := convertOption[TOption](opts...) + if err != nil { + return nil, err + } + return rp.Invoke(ctx, in, tos...) + } + + t := func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) { + in, ok := unpackStreamReader[I](input) + if !ok { + panic(newUnexpectedInputTypeErr(reflect.TypeOf(in), input.getType())) + } + + tos, err := convertOption[TOption](opts...) + if err != nil { + return nil, err + } + + out, err := rp.Transform(ctx, in, tos...) + if err != nil { + return nil, err + } + + return packStreamReader(out), nil + } + + c.i = i + c.t = t + + return c +} + +// Invoke works like `ping => pong`. +func (rp *runnablePacker[I, O, TOption]) Invoke(ctx context.Context, + input I, opts ...TOption) (output O, err error) { + return rp.i(ctx, input, opts...) +} + +// Stream works like `ping => stream output`. +func (rp *runnablePacker[I, O, TOption]) Stream(ctx context.Context, + input I, opts ...TOption) (output *schema.StreamReader[O], err error) { + + return rp.s(ctx, input, opts...) +} + +// Collect works like `stream input => pong`. +func (rp *runnablePacker[I, O, TOption]) Collect(ctx context.Context, + input *schema.StreamReader[I], opts ...TOption) (output O, err error) { + return rp.c(ctx, input, opts...) +} + +// Transform works like `stream input => stream output`. +func (rp *runnablePacker[I, O, TOption]) Transform(ctx context.Context, + input *schema.StreamReader[I], opts ...TOption) (output *schema.StreamReader[O], err error) { + return rp.t(ctx, input, opts...) +} + +func defaultImplConcatStreamReader[T any]( + sr *schema.StreamReader[T]) (T, error) { + + c, err := concatStreamReader(sr) + if err != nil { + var t T + return t, fmt.Errorf("concat stream reader fail: %w", err) + } + + return c, nil +} + +func invokeByStream[I, O, TOption any](s Stream[I, O, TOption]) Invoke[I, O, TOption] { + return func(ctx context.Context, input I, opts ...TOption) (output O, err error) { + sr, err := s(ctx, input, opts...) + if err != nil { + return output, err + } + + return defaultImplConcatStreamReader(sr) + } +} + +func invokeByCollect[I, O, TOption any](c Collect[I, O, TOption]) Invoke[I, O, TOption] { + return func(ctx context.Context, input I, opts ...TOption) (output O, err error) { + sr := schema.StreamReaderFromArray([]I{input}) + + return c(ctx, sr, opts...) + } +} + +func invokeByTransform[I, O, TOption any](t Transform[I, O, TOption]) Invoke[I, O, TOption] { + return func(ctx context.Context, input I, opts ...TOption) (output O, err error) { + srInput := schema.StreamReaderFromArray([]I{input}) + + srOutput, err := t(ctx, srInput, opts...) + if err != nil { + return output, err + } + + return defaultImplConcatStreamReader(srOutput) + } +} + +func streamByTransform[I, O, TOption any](t Transform[I, O, TOption]) Stream[I, O, TOption] { + return func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) { + srInput := schema.StreamReaderFromArray([]I{input}) + + return t(ctx, srInput, opts...) + } +} + +func streamByInvoke[I, O, TOption any](i Invoke[I, O, TOption]) Stream[I, O, TOption] { + return func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) { + out, err := i(ctx, input, opts...) + if err != nil { + return nil, err + } + + return schema.StreamReaderFromArray([]O{out}), nil + } +} + +func streamByCollect[I, O, TOption any](c Collect[I, O, TOption]) Stream[I, O, TOption] { + return func(ctx context.Context, input I, opts ...TOption) (output *schema.StreamReader[O], err error) { + srInput := schema.StreamReaderFromArray([]I{input}) + out, err := c(ctx, srInput, opts...) + if err != nil { + return nil, err + } + + return schema.StreamReaderFromArray([]O{out}), nil + } +} + +func collectByTransform[I, O, TOption any](t Transform[I, O, TOption]) Collect[I, O, TOption] { + return func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) { + srOutput, err := t(ctx, input, opts...) + if err != nil { + return output, err + } + + return defaultImplConcatStreamReader(srOutput) + } +} + +func collectByInvoke[I, O, TOption any](i Invoke[I, O, TOption]) Collect[I, O, TOption] { + return func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) { + in, err := defaultImplConcatStreamReader(input) + if err != nil { + return output, err + } + + return i(ctx, in, opts...) + } +} + +func collectByStream[I, O, TOption any](s Stream[I, O, TOption]) Collect[I, O, TOption] { + return func(ctx context.Context, input *schema.StreamReader[I], opts ...TOption) (output O, err error) { + in, err := defaultImplConcatStreamReader(input) + if err != nil { + return output, err + } + + srOutput, err := s(ctx, in, opts...) + if err != nil { + return output, err + } + + return defaultImplConcatStreamReader(srOutput) + } +} + +func transformByStream[I, O, TOption any](s Stream[I, O, TOption]) Transform[I, O, TOption] { + return func(ctx context.Context, input *schema.StreamReader[I], + opts ...TOption) (output *schema.StreamReader[O], err error) { + in, err := defaultImplConcatStreamReader(input) + if err != nil { + return output, err + } + + return s(ctx, in, opts...) + } +} + +func transformByCollect[I, O, TOption any](c Collect[I, O, TOption]) Transform[I, O, TOption] { + return func(ctx context.Context, input *schema.StreamReader[I], + opts ...TOption) (output *schema.StreamReader[O], err error) { + out, err := c(ctx, input, opts...) + if err != nil { + return output, err + } + + return schema.StreamReaderFromArray([]O{out}), nil + } +} + +func transformByInvoke[I, O, TOption any](i Invoke[I, O, TOption]) Transform[I, O, TOption] { + return func(ctx context.Context, input *schema.StreamReader[I], + opts ...TOption) (output *schema.StreamReader[O], err error) { + in, err := defaultImplConcatStreamReader(input) + if err != nil { + return output, err + } + + out, err := i(ctx, in, opts...) + if err != nil { + return output, err + } + + return schema.StreamReaderFromArray([]O{out}), nil + } +} + +func newRunnablePacker[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption], + c Collect[I, O, TOption], t Transform[I, O, TOption], enableCallback bool) *runnablePacker[I, O, TOption] { + + r := &runnablePacker[I, O, TOption]{} + + if enableCallback { + if i != nil { + i = invokeWithCallbacks(i) + } + + if s != nil { + s = streamWithCallbacks(s) + } + + if c != nil { + c = collectWithCallbacks(c) + } + + if t != nil { + t = transformWithCallbacks(t) + } + } + + if i != nil { + r.i = i + } else if s != nil { + r.i = invokeByStream(s) + } else if c != nil { + r.i = invokeByCollect(c) + } else { + r.i = invokeByTransform(t) + } + + if s != nil { + r.s = s + } else if t != nil { + r.s = streamByTransform(t) + } else if i != nil { + r.s = streamByInvoke(i) + } else { + r.s = streamByCollect(c) + } + + if c != nil { + r.c = c + } else if t != nil { + r.c = collectByTransform(t) + } else if i != nil { + r.c = collectByInvoke(i) + } else { + r.c = collectByStream(s) + } + + if t != nil { + r.t = t + } else if s != nil { + r.t = transformByStream(s) + } else if c != nil { + r.t = transformByCollect(c) + } else { + r.t = transformByInvoke(i) + } + + return r +} + +func toGenericRunnable[I, O any](cr *composableRunnable, ctxWrapper func(ctx context.Context, opts ...Option) context.Context) ( + *runnablePacker[I, O, Option], error) { + i := func(ctx context.Context, input I, opts ...Option) (output O, err error) { + out, err := cr.i(ctx, input, toAnyList(opts)...) + if err != nil { + return output, err + } + + to, ok := out.(O) + if !ok { + // When a nil is passed as an 'any' type, its original type information is lost, + // becoming an untyped nil. This would cause type assertions to fail. + // So if the output is nil and the target type O is an interface, we need to explicitly create a nil of type O. + if out == nil && generic.TypeOf[O]().Kind() == reflect.Interface { + var o O + to = o + } else { + panic(newUnexpectedInputTypeErr(generic.TypeOf[O](), reflect.TypeOf(out))) + } + } + return to, nil + } + + t := func(ctx context.Context, input *schema.StreamReader[I], + opts ...Option) (output *schema.StreamReader[O], err error) { + in := packStreamReader(input) + out, err := cr.t(ctx, in, toAnyList(opts)...) + + if err != nil { + return nil, err + } + + output, ok := unpackStreamReader[O](out) + if !ok { + panic("impossible") + } + + return output, nil + } + + r := newRunnablePacker(i, nil, nil, t, false) + r.wrapRunnableCtx(ctxWrapper) + + return r, nil +} + +func inputKeyedComposableRunnable(key string, r *composableRunnable) *composableRunnable { + wrapper := *r + wrapper.genericHelper = wrapper.genericHelper.forMapInput() + i := r.i + wrapper.i = func(ctx context.Context, input any, opts ...any) (output any, err error) { + v, ok := input.(map[string]any)[key] + if !ok { + return nil, fmt.Errorf("cannot find input key: %s", key) + } + out, err := i(ctx, v, opts...) + if err != nil { + return nil, err + } + + return out, nil + } + + t := r.t + wrapper.t = func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) { + nInput, ok := r.inputStreamFilter(key, input) + if !ok { + return nil, fmt.Errorf("inputStreamFilter failed, key= %s, node name= %s, err= %w", key, r.nodeInfo.name, err) + } + out, err := t(ctx, nInput, opts...) + if err != nil { + return nil, err + } + + return out, nil + } + + wrapper.inputType = generic.TypeOf[map[string]any]() + return &wrapper +} + +func outputKeyedComposableRunnable(key string, r *composableRunnable) *composableRunnable { + wrapper := *r + wrapper.genericHelper = wrapper.genericHelper.forMapOutput() + i := r.i + wrapper.i = func(ctx context.Context, input any, opts ...any) (output any, err error) { + out, err := i(ctx, input, opts...) + if err != nil { + return nil, err + } + + return map[string]any{key: out}, nil + } + + t := r.t + wrapper.t = func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) { + out, err := t(ctx, input, opts...) + if err != nil { + return nil, err + } + + return out.withKey(key), nil + } + + wrapper.outputType = generic.TypeOf[map[string]any]() + + return &wrapper +} + +// composablePassthrough special runnable that passthrough input to output +func composablePassthrough() *composableRunnable { + r := &composableRunnable{isPassthrough: true, nodeInfo: &nodeInfo{}} + + r.i = func(ctx context.Context, input any, opts ...any) (output any, err error) { + return input, nil + } + + r.t = func(ctx context.Context, input streamReader, opts ...any) (output streamReader, err error) { + return input, nil + } + + r.meta = &executorMeta{ + component: ComponentOfPassthrough, + isComponentCallbackEnabled: false, + componentImplType: "Passthrough", + } + + return r +} diff --git a/compose/runnable_test.go b/compose/runnable_test.go new file mode 100644 index 0000000..95744b7 --- /dev/null +++ b/compose/runnable_test.go @@ -0,0 +1,210 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestRunnableLambda(t *testing.T) { + ctx := context.Background() + + t.Run("invoke_to_runnable", func(t *testing.T) { + rl := runnableLambda( + func(ctx context.Context, input int, opts ...Option) (output string, err error) { + return strconv.Itoa(input) + "+" + opts[0].options[0].(string), nil + }, + nil, nil, nil, false) + + ctxWrapper := func(ctx context.Context, opts ...Option) context.Context { + return ctx + } + gr, err := toGenericRunnable[int, string](rl, ctxWrapper) + assert.NoError(t, err) + out, err := gr.Invoke(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err := gr.Stream(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sri, swi := schema.Pipe[int](1) + _ = swi.Send(10, nil) + swi.Close() + sriArr := sri.Copy(2) + + out, err = gr.Collect(ctx, sriArr[0], WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + }) + + t.Run("stream_to_runnable", func(t *testing.T) { + rl := runnableLambda(nil, + func(ctx context.Context, input int, opts ...Option) (output *schema.StreamReader[string], err error) { + sro, swo := schema.Pipe[string](3) + _ = swo.Send(strconv.Itoa(input), nil) + _ = swo.Send("+", nil) + _ = swo.Send(opts[0].options[0].(string), nil) + swo.Close() + return sro, nil + }, nil, nil, false) + + ctxWrapper := func(ctx context.Context, opts ...Option) context.Context { + return ctx + } + gr, err := toGenericRunnable[int, string](rl, ctxWrapper) + assert.NoError(t, err) + out, err := gr.Invoke(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err := gr.Stream(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sri, swi := schema.Pipe[int](1) + _ = swi.Send(10, nil) + swi.Close() + sriArr := sri.Copy(2) + + out, err = gr.Collect(ctx, sriArr[0], WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + }) + + t.Run("transform_to_runnable", func(t *testing.T) { + rl := runnableLambda( + nil, nil, nil, + func(ctx context.Context, input *schema.StreamReader[int], opts ...Option) (output *schema.StreamReader[string], err error) { + + in, e := input.Recv() + if errors.Is(e, io.EOF) { + return nil, fmt.Errorf("unpected EOF") + } + input.Close() + + sro, swo := schema.Pipe[string](3) + _ = swo.Send(strconv.Itoa(in), nil) + _ = swo.Send("+", nil) + _ = swo.Send(opts[0].options[0].(string), nil) + swo.Close() + return sro, nil + }, + false) + + ctxWrapper := func(ctx context.Context, opts ...Option) context.Context { + return ctx + } + gr, err := toGenericRunnable[int, string](rl, ctxWrapper) + assert.NoError(t, err) + out, err := gr.Invoke(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err := gr.Stream(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sri, swi := schema.Pipe[int](1) + _ = swi.Send(10, nil) + swi.Close() + sriArr := sri.Copy(2) + + out, err = gr.Collect(ctx, sriArr[0], WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + }) + + t.Run("collect_to_runnable", func(t *testing.T) { + rl := runnableLambda(nil, nil, + func(ctx context.Context, input *schema.StreamReader[int], opts ...Option) (output string, err error) { + in, e := input.Recv() + if errors.Is(e, io.EOF) { + return "", fmt.Errorf("unpected EOF") + } + input.Close() + + return strconv.Itoa(in) + "+" + opts[0].options[0].(string), nil + }, + nil, false) + + ctxWrapper := func(ctx context.Context, opts ...Option) context.Context { + return ctx + } + + gr, err := toGenericRunnable[int, string](rl, ctxWrapper) + assert.NoError(t, err) + out, err := gr.Invoke(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err := gr.Stream(ctx, 10, WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sri, swi := schema.Pipe[int](1) + _ = swi.Send(10, nil) + swi.Close() + sriArr := sri.Copy(2) + + out, err = gr.Collect(ctx, sriArr[0], WithLambdaOption("100")) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + + sr, err = gr.Transform(ctx, sriArr[1], WithLambdaOption("100")) + assert.NoError(t, err) + out, err = concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "10+100", out) + }) +} diff --git a/compose/state.go b/compose/state.go new file mode 100644 index 0000000..f350601 --- /dev/null +++ b/compose/state.go @@ -0,0 +1,196 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + "sync" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +// GenLocalState is a function that generates the state. +type GenLocalState[S any] func(ctx context.Context) (state S) + +type stateKey struct{} + +type internalState struct { + state any + mu sync.Mutex + parent *internalState +} + +// StatePreHandler is a function called before the node is executed. +// Notice: if user called Stream but with StatePreHandler, the StatePreHandler will read all stream chunks and merge them into a single object. +type StatePreHandler[I, S any] func(ctx context.Context, in I, state S) (I, error) + +// StatePostHandler is a function called after the node is executed. +// Notice: if user called Stream but with StatePostHandler, the StatePostHandler will read all stream chunks and merge them into a single object. +type StatePostHandler[O, S any] func(ctx context.Context, out O, state S) (O, error) + +// StreamStatePreHandler is a function that is called before the node is executed with stream input and output. +type StreamStatePreHandler[I, S any] func(ctx context.Context, in *schema.StreamReader[I], state S) (*schema.StreamReader[I], error) + +// StreamStatePostHandler is a function that is called after the node is executed with stream input and output. +type StreamStatePostHandler[O, S any] func(ctx context.Context, out *schema.StreamReader[O], state S) (*schema.StreamReader[O], error) + +func convertPreHandler[I, S any](handler StatePreHandler[I, S]) *composableRunnable { + rf := func(ctx context.Context, in I, opts ...any) (I, error) { + cState, pMu, err := getState[S](ctx) + if err != nil { + return in, err + } + pMu.Lock() + defer pMu.Unlock() + + return handler(ctx, in, cState) + } + + return runnableLambda[I, I](rf, nil, nil, nil, false) +} + +func convertPostHandler[O, S any](handler StatePostHandler[O, S]) *composableRunnable { + rf := func(ctx context.Context, out O, opts ...any) (O, error) { + cState, pMu, err := getState[S](ctx) + if err != nil { + return out, err + } + pMu.Lock() + defer pMu.Unlock() + + return handler(ctx, out, cState) + } + + return runnableLambda[O, O](rf, nil, nil, nil, false) +} + +func streamConvertPreHandler[I, S any](handler StreamStatePreHandler[I, S]) *composableRunnable { + rf := func(ctx context.Context, in *schema.StreamReader[I], opts ...any) (*schema.StreamReader[I], error) { + cState, pMu, err := getState[S](ctx) + if err != nil { + return in, err + } + pMu.Lock() + defer pMu.Unlock() + + return handler(ctx, in, cState) + } + + return runnableLambda[I, I](nil, nil, nil, rf, false) +} + +func streamConvertPostHandler[O, S any](handler StreamStatePostHandler[O, S]) *composableRunnable { + rf := func(ctx context.Context, out *schema.StreamReader[O], opts ...any) (*schema.StreamReader[O], error) { + cState, pMu, err := getState[S](ctx) + if err != nil { + return out, err + } + pMu.Lock() + defer pMu.Unlock() + + return handler(ctx, out, cState) + } + + return runnableLambda[O, O](nil, nil, nil, rf, false) +} + +// ProcessState processes the state from the context in a concurrency-safe way. +// This is the recommended way to access and modify state in custom nodes. +// The provided function handler will be executed with exclusive access to the state (protected by mutex). +// +// State Lookup Behavior: +// - If the requested state type exists in the current graph, it will be returned +// - If not found in current graph, ProcessState will search in parent graph states (for nested graphs) +// - This enables nested graphs to access state from their parent graphs +// - Follows lexical scoping: inner state of the same type shadows outer state +// +// Concurrency Safety: +// - ProcessState automatically locks the mutex of the state being accessed (current or parent level) +// - Each state level has its own mutex, allowing concurrent access to different levels +// - The lock is held for the entire duration of the handler function +// +// Note: This method will report an error if the state type doesn't match or state is not found in the context chain. +// +// Example - Basic usage in a single graph: +// +// lambdaFunc := func(ctx context.Context, in string, opts ...any) (string, error) { +// err := compose.ProcessState[*MyState](ctx, func(ctx context.Context, state *MyState) error { +// // Safely modify state +// state.Count++ +// return nil +// }) +// if err != nil { +// return "", err +// } +// return in, nil +// } +// +// Example - Nested graph accessing parent state: +// +// // In an inner graph node +// innerNode := func(ctx context.Context, input string) (string, error) { +// // Access parent graph's state +// err := compose.ProcessState[*OuterState](ctx, func(ctx context.Context, s *OuterState) error { +// s.Counter++ // Safely modify parent state +// return nil +// }) +// if err != nil { +// return "", err +// } +// +// // Also access inner graph's own state +// err = compose.ProcessState[*InnerState](ctx, func(ctx context.Context, s *InnerState) error { +// s.Data = "processed" +// return nil +// }) +// return input, nil +// } +func ProcessState[S any](ctx context.Context, handler func(context.Context, S) error) error { + s, pMu, err := getState[S](ctx) + if err != nil { + return fmt.Errorf("get state from context fail: %w", err) + } + pMu.Lock() + defer pMu.Unlock() + return handler(ctx, s) +} + +func getState[S any](ctx context.Context) (S, *sync.Mutex, error) { + state := ctx.Value(stateKey{}) + + if state == nil { + var s S + return s, nil, fmt.Errorf("have not set state") + } + + interState := state.(*internalState) + + for interState != nil { + if cState, ok := interState.state.(S); ok { + return cState, &interState.mu, nil + } + interState = interState.parent + } + + var s S + return s, nil, fmt.Errorf("cannot find state with type: %v in states chain, "+ + "current state type: %v", + generic.TypeOf[S](), reflect.TypeOf(state.(*internalState).state)) +} diff --git a/compose/state_test.go b/compose/state_test.go new file mode 100644 index 0000000..247b4b6 --- /dev/null +++ b/compose/state_test.go @@ -0,0 +1,765 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "io" + "strings" + "sync" + "testing" + "time" + "unicode/utf8" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +type midStr string + +func TestStateGraphWithEdge(t *testing.T) { + + ctx := context.Background() + + const ( + nodeOfL1 = "invokable" + nodeOfL2 = "streamable" + nodeOfL3 = "transformable" + ) + + type testState struct { + ms []string + } + + gen := func(ctx context.Context) *testState { + return &testState{} + } + + sg := NewGraph[string, string](WithGenLocalState(gen)) + + l1 := InvokableLambda(func(ctx context.Context, in string) (out midStr, err error) { + return midStr("InvokableLambda: " + in), nil + }) + + l1StateToInput := func(ctx context.Context, in string, state *testState) (string, error) { + state.ms = append(state.ms, in) + return in, nil + } + + l1StateToOutput := func(ctx context.Context, out midStr, state *testState) (midStr, error) { + state.ms = append(state.ms, string(out)) + return out, nil + } + + err := sg.AddLambdaNode(nodeOfL1, l1, + WithStatePreHandler(l1StateToInput), WithStatePostHandler(l1StateToOutput)) + assert.NoError(t, err) + + l2 := StreamableLambda(func(ctx context.Context, input midStr) (output *schema.StreamReader[string], err error) { + outStr := "StreamableLambda: " + string(input) + + sr, sw := schema.Pipe[string](utf8.RuneCountInString(outStr)) + + go func() { + for _, field := range strings.Fields(outStr) { + sw.Send(field+" ", nil) + } + sw.Close() + }() + + return sr, nil + }) + + l2StateToOutput := func(ctx context.Context, out string, state *testState) (string, error) { + state.ms = append(state.ms, out) + return out, nil + } + + err = sg.AddLambdaNode(nodeOfL2, l2, WithStatePostHandler(l2StateToOutput)) + assert.NoError(t, err) + + l3 := TransformableLambda(func(ctx context.Context, input *schema.StreamReader[string]) ( + output *schema.StreamReader[string], err error) { + + prefix := "TransformableLambda: " + sr, sw := schema.Pipe[string](20) + + go func() { + for _, field := range strings.Fields(prefix) { + sw.Send(field+" ", nil) + } + defer input.Close() + + for { + chunk, err := input.Recv() + if err != nil { + if err == io.EOF { + break + } + // TODO: how to trace this kind of error in the goroutine of processing stream + sw.Send(chunk, err) + break + } + + sw.Send(chunk, nil) + + } + sw.Close() + }() + + return sr, nil + }) + + l3StateToOutput := func(ctx context.Context, out string, state *testState) (string, error) { + state.ms = append(state.ms, out) + assert.Len(t, state.ms, 4) + return out, nil + } + + err = sg.AddLambdaNode(nodeOfL3, l3, WithStatePostHandler(l3StateToOutput)) + assert.NoError(t, err) + + err = sg.AddEdge(START, nodeOfL1) + assert.NoError(t, err) + + err = sg.AddEdge(nodeOfL1, nodeOfL2) + assert.NoError(t, err) + + err = sg.AddEdge(nodeOfL2, nodeOfL3) + assert.NoError(t, err) + + err = sg.AddEdge(nodeOfL3, END) + assert.NoError(t, err) + + run, err := sg.Compile(ctx) + assert.NoError(t, err) + + out, err := run.Invoke(ctx, "how are you") + assert.NoError(t, err) + assert.Equal(t, "TransformableLambda: StreamableLambda: InvokableLambda: how are you ", out) + + stream, err := run.Stream(ctx, "how are you") + assert.NoError(t, err) + out, err = concatStreamReader(stream) + assert.NoError(t, err) + assert.Equal(t, "TransformableLambda: StreamableLambda: InvokableLambda: how are you ", out) + + sr, sw := schema.Pipe[string](1) + sw.Send("how are you", nil) + sw.Close() + + stream, err = run.Transform(ctx, sr) + assert.NoError(t, err) + out, err = concatStreamReader(stream) + assert.NoError(t, err) + assert.Equal(t, "TransformableLambda: StreamableLambda: InvokableLambda: how are you ", out) +} + +func TestStateGraphUtils(t *testing.T) { + t.Run("getState_success", func(t *testing.T) { + type testStruct struct { + UserID int64 + } + + ctx := context.Background() + + ctx = context.WithValue(ctx, stateKey{}, &internalState{ + state: &testStruct{UserID: 10}, + }) + + var userID int64 + err := ProcessState[*testStruct](ctx, func(_ context.Context, state *testStruct) error { + userID = state.UserID + return nil + }) + assert.NoError(t, err) + assert.Equal(t, int64(10), userID) + }) + + t.Run("getState_nil", func(t *testing.T) { + type testStruct struct { + UserID int64 + } + + ctx := context.Background() + ctx = context.WithValue(ctx, stateKey{}, &internalState{}) + + err := ProcessState[*testStruct](ctx, func(_ context.Context, state *testStruct) error { + return nil + }) + assert.ErrorContains(t, err, "cannot find state with type: *compose.testStruct in states chain, "+ + "current state type: ") + }) + + t.Run("getState_type_error", func(t *testing.T) { + type testStruct struct { + UserID int64 + } + + ctx := context.Background() + ctx = context.WithValue(ctx, stateKey{}, &internalState{ + state: &testStruct{UserID: 10}, + }) + + err := ProcessState[string](ctx, func(_ context.Context, state string) error { + return nil + }) + assert.ErrorContains(t, err, "cannot find state with type: string in states chain, "+ + "current state type: *compose.testStruct") + + }) +} + +func TestStateChain(t *testing.T) { + ctx := context.Background() + type testState struct { + Field1 string + Field2 string + } + sc := NewChain[string, string](WithGenLocalState(func(ctx context.Context) (state *testState) { + return &testState{} + })) + + r, err := sc.AppendLambda(InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + err = ProcessState[*testState](ctx, func(_ context.Context, state *testState) error { + state.Field1 = "node1" + return nil + }) + if err != nil { + return "", err + } + return input, nil + }), WithStatePostHandler(func(ctx context.Context, out string, state *testState) (string, error) { + state.Field2 = "node2" + return out, nil + })). + AppendLambda(InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + }), WithStatePreHandler(func(ctx context.Context, in string, state *testState) (string, error) { + return in + state.Field1 + state.Field2, nil + })).Compile(ctx) + if err != nil { + t.Fatal(err) + } + result, err := r.Invoke(ctx, "start") + if err != nil { + t.Fatal(err) + } + if result != "startnode1node2" { + t.Fatal("result is unexpected") + } +} + +func TestStreamState(t *testing.T) { + type testState struct { + Field1 string + } + ctx := context.Background() + s := &testState{Field1: "1"} + g := NewGraph[string, string](WithGenLocalState(func(ctx context.Context) (state *testState) { return s })) + err := g.AddLambdaNode("1", TransformableLambda(func(ctx context.Context, input *schema.StreamReader[string]) (output *schema.StreamReader[string], err error) { + return input, nil + }), WithStreamStatePreHandler(func(ctx context.Context, in *schema.StreamReader[string], state *testState) (*schema.StreamReader[string], error) { + sr, sw := schema.Pipe[string](5) + for i := 0; i < 5; i++ { + sw.Send(state.Field1, nil) + } + sw.Close() + return sr, nil + }), WithStreamStatePostHandler(func(ctx context.Context, in *schema.StreamReader[string], state *testState) (*schema.StreamReader[string], error) { + ss := in.Copy(2) + for { + chunk, err := ss[0].Recv() + if err == io.EOF { + return ss[1], nil + } + if err != nil { + return nil, err + } + state.Field1 += chunk + } + })) + if err != nil { + t.Fatal(err) + } + err = g.AddEdge(START, "1") + if err != nil { + t.Fatal(err) + } + err = g.AddEdge("1", END) + if err != nil { + t.Fatal(err) + } + r, err := g.Compile(ctx) + if err != nil { + t.Fatal(err) + } + sr, _ := schema.Pipe[string](1) + streamResult, err := r.Transform(ctx, sr) + if err != nil { + t.Fatal(err) + } + if s.Field1 != "111111" { + t.Fatal("state is unexpected") + } + for i := 0; i < 5; i++ { + chunk, err := streamResult.Recv() + if err != nil { + t.Fatal(err) + } + if chunk != "1" { + t.Fatal("result is unexpected") + } + } + _, err = streamResult.Recv() + if err != io.EOF { + t.Fatal("result is unexpected") + } +} + +// Nested Graph State Tests + +type NestedOuterState struct { + Value string + Counter int +} + +type NestedInnerState struct { + Value string +} + +func init() { + schema.RegisterName[*NestedOuterState]("NestedOuterState") + schema.RegisterName[*NestedInnerState]("NestedInnerState") +} + +func TestNestedGraphStateAccess(t *testing.T) { + // Test that inner graph can access outer graph's state + genOuterState := func(ctx context.Context) *NestedOuterState { + return &NestedOuterState{Value: "outer", Counter: 0} + } + + genInnerState := func(ctx context.Context) *NestedInnerState { + return &NestedInnerState{Value: "inner"} + } + + innerNode := func(ctx context.Context, input string) (string, error) { + // Access both inner and outer state + var outerValue string + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + outerValue = s.Value + return nil + }) + if err != nil { + return "", err + } + + var innerValue string + err = ProcessState(ctx, func(ctx context.Context, s *NestedInnerState) error { + innerValue = s.Value + return nil + }) + if err != nil { + return "", err + } + + return fmt.Sprintf("%s_inner=%s_outer=%s", input, innerValue, outerValue), nil + } + + innerGraph := NewGraph[string, string](WithGenLocalState(genInnerState)) + _ = innerGraph.AddLambdaNode("inner_node", InvokableLambda(innerNode)) + _ = innerGraph.AddEdge(START, "inner_node") + _ = innerGraph.AddEdge("inner_node", END) + + outerGraph := NewGraph[string, string](WithGenLocalState(genOuterState)) + _ = outerGraph.AddGraphNode("inner_graph", innerGraph) + _ = outerGraph.AddEdge(START, "inner_graph") + _ = outerGraph.AddEdge("inner_graph", END) + + r, err := outerGraph.Compile(context.Background()) + assert.NoError(t, err) + + out, err := r.Invoke(context.Background(), "start") + assert.NoError(t, err) + assert.Equal(t, "start_inner=inner_outer=outer", out) +} + +func TestNestedGraphStateShadowing(t *testing.T) { + // Test that inner state shadows outer state of the same type (lexical scoping) + type CommonState struct { + Value string + } + + genOuterState := func(ctx context.Context) *CommonState { + return &CommonState{Value: "outer"} + } + + genInnerState := func(ctx context.Context) *CommonState { + return &CommonState{Value: "inner"} + } + + innerNode := func(ctx context.Context, input string) (string, error) { + var value string + err := ProcessState(ctx, func(ctx context.Context, s *CommonState) error { + // Should see "inner" because inner state shadows outer state + value = s.Value + return nil + }) + if err != nil { + return "", err + } + return input + "_" + value, nil + } + + innerGraph := NewGraph[string, string](WithGenLocalState(genInnerState)) + _ = innerGraph.AddLambdaNode("inner_node", InvokableLambda(innerNode)) + _ = innerGraph.AddEdge(START, "inner_node") + _ = innerGraph.AddEdge("inner_node", END) + + outerGraph := NewGraph[string, string](WithGenLocalState(genOuterState)) + _ = outerGraph.AddGraphNode("inner_graph", innerGraph) + _ = outerGraph.AddEdge(START, "inner_graph") + _ = outerGraph.AddEdge("inner_graph", END) + + r, err := outerGraph.Compile(context.Background()) + assert.NoError(t, err) + + out, err := r.Invoke(context.Background(), "start") + assert.NoError(t, err) + assert.Equal(t, "start_inner", out) +} + +func TestNestedGraphStateAfterResume(t *testing.T) { + // Test that state parent linking works correctly after resume + // when the outer state is restored from checkpoint (new instance) + genOuterState := func(ctx context.Context) *NestedOuterState { + return &NestedOuterState{Value: "outer", Counter: 0} + } + + genInnerState := func(ctx context.Context) *NestedInnerState { + return &NestedInnerState{Value: "inner"} + } + + // Node that modifies outer state + outerNode := func(ctx context.Context, input string) (string, error) { + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + s.Counter = 42 + return nil + }) + if err != nil { + return "", err + } + return input, nil + } + + // Inner node that reads outer state + innerNode := func(ctx context.Context, input string) (string, error) { + var outerCounter int + var outerValue string + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + // Should see the modified counter value from the restored state + outerCounter = s.Counter + outerValue = s.Value + return nil + }) + if err != nil { + return "", err + } + return fmt.Sprintf("%s_counter=%d_value=%s", input, outerCounter, outerValue), nil + } + + innerGraph := NewGraph[string, string](WithGenLocalState(genInnerState)) + _ = innerGraph.AddLambdaNode("inner_node", InvokableLambda(innerNode)) + _ = innerGraph.AddEdge(START, "inner_node") + _ = innerGraph.AddEdge("inner_node", END) + + outerGraph := NewGraph[string, string](WithGenLocalState(genOuterState)) + _ = outerGraph.AddLambdaNode("outer_node", InvokableLambda(outerNode)) + _ = outerGraph.AddGraphNode("inner_graph", innerGraph, WithGraphCompileOptions(WithInterruptBeforeNodes([]string{"inner_node"}))) + _ = outerGraph.AddEdge(START, "outer_node") + _ = outerGraph.AddEdge("outer_node", "inner_graph") + _ = outerGraph.AddEdge("inner_graph", END) + + store := newInMemoryStore() + r, err := outerGraph.Compile(context.Background(), WithCheckPointStore(store)) + assert.NoError(t, err) + + // First run - should interrupt after modifying outer state + _, err = r.Invoke(context.Background(), "start", WithCheckPointID("state_resume_test")) + assert.Error(t, err) + + // Resume - outer state should be restored with Counter=42 + // Inner graph should link to this restored outer state + out, err := r.Invoke(context.Background(), "start", WithCheckPointID("state_resume_test")) + assert.NoError(t, err) + assert.Equal(t, "start_counter=42_value=outer", out) +} + +func TestLambdaNestedGraphStateAccess(t *testing.T) { + // Test that inner graph invoked from a lambda can access outer graph's state + // This tests the case: outer graph -> lambda node -> inner graph (using CompositeInterrupt) + genOuterState := func(ctx context.Context) *NestedOuterState { + return &NestedOuterState{Value: "outer", Counter: 100} + } + + genInnerState := func(ctx context.Context) *NestedInnerState { + return &NestedInnerState{Value: "inner"} + } + + // Inner node that accesses outer state + innerNode := func(ctx context.Context, input string) (string, error) { + var outerValue string + var outerCounter int + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + outerValue = s.Value + outerCounter = s.Counter + return nil + }) + if err != nil { + return "", err + } + + var innerValue string + err = ProcessState(ctx, func(ctx context.Context, s *NestedInnerState) error { + innerValue = s.Value + return nil + }) + if err != nil { + return "", err + } + + return fmt.Sprintf("%s_inner=%s_outer=%s_%d", input, innerValue, outerValue, outerCounter), nil + } + + // Build inner graph + innerGraph := NewGraph[string, string](WithGenLocalState(genInnerState)) + _ = innerGraph.AddLambdaNode("inner_node", InvokableLambda(innerNode)) + _ = innerGraph.AddEdge(START, "inner_node") + _ = innerGraph.AddEdge("inner_node", END) + + // Compile inner graph as a standalone runnable + innerRunnable, err := innerGraph.Compile(context.Background()) + assert.NoError(t, err) + + // Lambda that invokes the inner graph + lambdaNode := InvokableLambda(func(ctx context.Context, input string) (string, error) { + // Simply invoke the inner graph - state context is passed through + return innerRunnable.Invoke(ctx, input) + }) + + // Build outer graph + outerGraph := NewGraph[string, string](WithGenLocalState(genOuterState)) + _ = outerGraph.AddLambdaNode("lambda_with_graph", lambdaNode) + _ = outerGraph.AddEdge(START, "lambda_with_graph") + _ = outerGraph.AddEdge("lambda_with_graph", END) + + r, err := outerGraph.Compile(context.Background()) + assert.NoError(t, err) + + out, err := r.Invoke(context.Background(), "start") + assert.NoError(t, err) + assert.Equal(t, "start_inner=inner_outer=outer_100", out) +} + +func TestLambdaNestedGraphStateAfterResume(t *testing.T) { + // Test that state parent linking works correctly after resume + // in the lambda-nested case (outer graph -> lambda -> inner graph) + genOuterState := func(ctx context.Context) *NestedOuterState { + return &NestedOuterState{Value: "outer", Counter: 0} + } + + genInnerState := func(ctx context.Context) *NestedInnerState { + return &NestedInnerState{Value: "inner"} + } + + // Outer node that modifies state + outerNode := func(ctx context.Context, input string) (string, error) { + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + s.Counter = 99 + return nil + }) + if err != nil { + return "", err + } + return input, nil + } + + // Inner lambda that interrupts on first run, reads outer state on resume + innerLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + wasInterrupted, _, _ := GetInterruptState[*NestedInnerState](ctx) + if !wasInterrupted { + // First run: interrupt + return "", StatefulInterrupt(ctx, "inner interrupt", &NestedInnerState{Value: "inner"}) + } + + // Resumed: read outer state + var outerCounter int + var outerValue string + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + // Should see the modified counter from the restored state + outerCounter = s.Counter + outerValue = s.Value + return nil + }) + if err != nil { + return "", err + } + + return fmt.Sprintf("%s_counter=%d_value=%s", input, outerCounter, outerValue), nil + }) + + // Build inner graph + innerGraph := NewGraph[string, string](WithGenLocalState(genInnerState)) + _ = innerGraph.AddLambdaNode("inner_lambda", innerLambda) + _ = innerGraph.AddEdge(START, "inner_lambda") + _ = innerGraph.AddEdge("inner_lambda", END) + + // Compile inner graph as standalone runnable with checkpoint support + innerRunnable, err := innerGraph.Compile(context.Background(), + WithGraphName("inner"), + WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + // Composite lambda that invokes the inner graph and handles interrupts + compositeLambda := InvokableLambda(func(ctx context.Context, input string) (string, error) { + output, err := innerRunnable.Invoke(ctx, input, WithCheckPointID("inner-cp")) + if err != nil { + _, isInterrupt := ExtractInterruptInfo(err) + if !isInterrupt { + return "", err + } + // Wrap the interrupt using CompositeInterrupt + return "", CompositeInterrupt(ctx, "composite interrupt", nil, err) + } + return output, nil + }) + + // Build outer graph + outerGraph := NewGraph[string, string](WithGenLocalState(genOuterState)) + _ = outerGraph.AddLambdaNode("outer_node", InvokableLambda(outerNode)) + _ = outerGraph.AddLambdaNode("composite_lambda", compositeLambda) + _ = outerGraph.AddEdge(START, "outer_node") + _ = outerGraph.AddEdge("outer_node", "composite_lambda") + _ = outerGraph.AddEdge("composite_lambda", END) + + // Compile outer graph + outerRunnable, err := outerGraph.Compile(context.Background(), + WithGraphName("root"), + WithCheckPointStore(newInMemoryStore())) + assert.NoError(t, err) + + // First run - should interrupt after modifying outer state + checkPointID := "lambda_state_resume_test" + _, err = outerRunnable.Invoke(context.Background(), "start", WithCheckPointID(checkPointID)) + assert.Error(t, err) + + interruptInfo, isInterrupt := ExtractInterruptInfo(err) + assert.True(t, isInterrupt) + + // Resume - outer state should be restored with Counter=99 + // Inner lambda should link to this restored outer state + ctx := ResumeWithData(context.Background(), interruptInfo.InterruptContexts[0].ID, nil) + out, err := outerRunnable.Invoke(ctx, "start", WithCheckPointID(checkPointID)) + assert.NoError(t, err) + + // Verify the inner lambda saw the modified counter from the restored outer state + assert.Contains(t, out, "counter=99") + assert.Contains(t, out, "value=outer") +} + +func TestNestedGraphStateConcurrency(t *testing.T) { + // Test that concurrent access to parent and child states uses correct locks + // This verifies that ProcessState properly locks the parent state's mutex when accessing it + genOuterState := func(ctx context.Context) *NestedOuterState { + return &NestedOuterState{Value: "outer", Counter: 0} + } + + genInnerState := func(ctx context.Context) *NestedInnerState { + return &NestedInnerState{Value: "inner"} + } + + // Inner node that concurrently modifies both outer and inner state + innerNode := func(ctx context.Context, input string) (string, error) { + var wg sync.WaitGroup + errors := make(chan error, 20) + + // Launch 10 goroutines that modify outer state + // If locks don't work correctly, we'll see race conditions + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + err := ProcessState(ctx, func(ctx context.Context, s *NestedOuterState) error { + // ProcessState should hold the parent's lock during this entire function + current := s.Counter + time.Sleep(time.Millisecond) // Simulate work + s.Counter = current + 1 + return nil + }) + if err != nil { + errors <- err + } + }() + } + + // Launch 10 goroutines that modify inner state + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + err := ProcessState(ctx, func(ctx context.Context, s *NestedInnerState) error { + // This uses the inner state's own lock + return nil + }) + if err != nil { + errors <- err + } + }() + } + + wg.Wait() + close(errors) + + // Check for errors + for err := range errors { + return "", err + } + + return input, nil + } + + innerGraph := NewGraph[string, string](WithGenLocalState(genInnerState)) + _ = innerGraph.AddLambdaNode("inner_node", InvokableLambda(innerNode)) + _ = innerGraph.AddEdge(START, "inner_node") + _ = innerGraph.AddEdge("inner_node", END) + + outerGraph := NewGraph[string, string](WithGenLocalState(genOuterState)) + _ = outerGraph.AddGraphNode("inner_graph", innerGraph) + _ = outerGraph.AddEdge(START, "inner_graph") + _ = outerGraph.AddEdge("inner_graph", END) + + r, err := outerGraph.Compile(context.Background()) + assert.NoError(t, err) + + _, err = r.Invoke(context.Background(), "start") + assert.NoError(t, err) + + // Note: This test is primarily validated by running with -race flag + // If locks don't work correctly, the race detector will catch it +} diff --git a/compose/stream_concat.go b/compose/stream_concat.go new file mode 100644 index 0000000..4feeba8 --- /dev/null +++ b/compose/stream_concat.go @@ -0,0 +1,88 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "errors" + "io" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/schema" +) + +// RegisterStreamChunkConcatFunc registers a function to concat stream chunks. +// It's required when you want to concat stream chunks of a specific type. +// for example you call Invoke() but node only implements Stream(). +// call at process init +// not thread safe +// eg. +// +// type testStruct struct { +// field1 string +// field2 int +// } +// compose.RegisterStreamChunkConcatFunc(func(items []testStruct) (testStruct, error) { +// return testStruct{ +// field1: items[1].field1, // may implement inplace logic by your scenario +// field2: items[0].field2 + items[1].field2, +// }, nil +// }) +func RegisterStreamChunkConcatFunc[T any](fn func([]T) (T, error)) { + internal.RegisterStreamChunkConcatFunc(fn) +} + +var emptyStreamConcatErr = errors.New("stream reader is empty, concat fail") + +func concatStreamReader[T any](sr *schema.StreamReader[T]) (T, error) { + defer sr.Close() + + var items []T + + for { + chunk, err := sr.Recv() + if err != nil { + if err == io.EOF { + break + } + + if _, ok := schema.GetSourceName(err); ok { + continue + } + + var t T + return t, newStreamReadError(err) + } + + items = append(items, chunk) + } + + if len(items) == 0 { + var t T + return t, emptyStreamConcatErr + } + + if len(items) == 1 { + return items[0], nil + } + + res, err := internal.ConcatItems(items) + if err != nil { + var t T + return t, err + } + return res, nil +} diff --git a/compose/stream_concat_test.go b/compose/stream_concat_test.go new file mode 100644 index 0000000..ce0210c --- /dev/null +++ b/compose/stream_concat_test.go @@ -0,0 +1,218 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "errors" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/schema" +) + +type tStreamConcatItemForTest struct { + s string +} + +func concatTStreamForTest(items []tStreamConcatItemForTest) (tStreamConcatItemForTest, error) { + var s string + for _, item := range items { + s += item.s + } + + return tStreamConcatItemForTest{s: s}, nil +} + +func concatIntForTest(items []int) (int, error) { + var i int + for _, item := range items { + i += item + } + + return i, nil +} + +type tConcatErrForTest struct{} + +func concatTStreamError(_ []tConcatErrForTest) (tConcatErrForTest, error) { + return tConcatErrForTest{}, errors.New("test error") +} + +func TestConcatRegistry(t *testing.T) { + + RegisterStreamChunkConcatFunc(concatTStreamForTest) + + sr, sw := schema.Pipe[tStreamConcatItemForTest](10) + go func() { + for i := 0; i < 10; i++ { + sw.Send(tStreamConcatItemForTest{s: strconv.Itoa(i)}, nil) + } + sw.Close() + }() + + lastVal, err := concatStreamReader(sr) + assert.Nil(t, err) + + assert.Equal(t, "0123456789", lastVal.s) +} + +func TestStringConcat(t *testing.T) { + sr, sw := schema.Pipe[string](10) + go func() { + for i := 0; i < 10; i++ { + sw.Send(strconv.Itoa(i), nil) + } + + sw.Close() + }() + + lastVal, err := concatStreamReader(sr) + assert.Nil(t, err) + + assert.Equal(t, "0123456789", lastVal) +} + +func TestMessageConcat(t *testing.T) { + sr, sw := schema.Pipe[*schema.Message](10) + go func() { + for i := 0; i < 10; i++ { + content := schema.UserMessage(strconv.Itoa(i)) + if i%4 == 0 { + content.Extra = map[string]any{ + "key_1": strconv.Itoa(i), + strconv.Itoa(i): strconv.Itoa(i), + } + } + sw.Send(content, nil) + } + sw.Close() + }() + + lastVal, err := concatStreamReader(sr) + assert.Nil(t, err) + assert.Equal(t, "0123456789", lastVal.Content) + assert.Len(t, lastVal.Extra, 4) + assert.Equal(t, map[string]any{ + "key_1": "048", + "0": "0", + "4": "4", + "8": "8", + }, lastVal.Extra) + +} + +func TestMapConcat(t *testing.T) { + RegisterStreamChunkConcatFunc(concatTStreamForTest) + RegisterStreamChunkConcatFunc(concatIntForTest) + + t.Run("simple map", func(t *testing.T) { + sr, sw := schema.Pipe[map[string]any](10) + + go func() { + for i := 0; i < 10; i++ { + sw.Send(map[string]any{ + "string": strconv.Itoa(i), + "custom_concat": tStreamConcatItemForTest{s: strconv.Itoa(9 - i)}, + "count": i, + }, nil) + } + sw.Close() + }() + + lastVal, err := concatStreamReader(sr) + assert.Nil(t, err) + + assert.Equal(t, "0123456789", lastVal["string"]) + assert.Equal(t, "9876543210", lastVal["custom_concat"].(tStreamConcatItemForTest).s) + assert.Equal(t, 45, lastVal["count"]) + + }) + + t.Run("complex map", func(t *testing.T) { + sr, sw := schema.Pipe[map[string]any](10) + + go func() { + for i := 0; i < 10; i++ { + // 嵌套 map, 仅允许第一层做类型合并，第二层直接覆盖 + sw.Send(map[string]any{ // 嵌套 map + "string": strconv.Itoa(i), + "deep_map": map[string]any{ + "message": &schema.Message{ + Content: strconv.Itoa(i), + }, + "custom_concat_deep": tStreamConcatItemForTest{s: strconv.Itoa(9 - i)}, + "count": i, + }, + "custom_concat": tStreamConcatItemForTest{s: strconv.Itoa(9 - i)}, + "count": i, + }, nil) + } + sw.Close() + }() + + lastVal, err := concatStreamReader(sr) + assert.Nil(t, err) + + assert.Equal(t, "0123456789", lastVal["string"]) + assert.Equal(t, 45, lastVal["count"]) + assert.Equal(t, "0123456789", lastVal["deep_map"].(map[string]any)["message"].(*schema.Message).Content) + assert.Equal(t, "9876543210", lastVal["deep_map"].(map[string]any)["custom_concat_deep"].(tStreamConcatItemForTest).s) + assert.Equal(t, 45, lastVal["deep_map"].(map[string]any)["count"]) + }) +} + +func TestConcatError(t *testing.T) { + t.Run("map type not equal", func(t *testing.T) { + a := map[string]any{ + "str": "string_01", + "x": "string_in_a", + } + + b := map[string]any{ + "str": "string_02", + "x": 123, + } + _, err := internal.ConcatItems([]map[string]any{a, b}) + assert.NotNil(t, err) + }) + + t.Run("merge error", func(t *testing.T) { + RegisterStreamChunkConcatFunc(concatTStreamError) + + _, err := internal.ConcatItems([]tConcatErrForTest{{}, {}}) + assert.NotNil(t, err) + }) +} + +func TestConcatSliceValue(t *testing.T) { + type testStruct struct { + A string + } + + s := []testStruct{{}, {A: "123"}, {}} + result, err := internal.ConcatItems(s) + assert.Nil(t, err) + assert.Equal(t, testStruct{A: "123"}, result) + + s = []testStruct{{}, {}, {}} + result, err = internal.ConcatItems(s) + assert.Nil(t, err) + assert.Equal(t, testStruct{}, result) +} diff --git a/compose/stream_reader.go b/compose/stream_reader.go new file mode 100644 index 0000000..07acbb3 --- /dev/null +++ b/compose/stream_reader.go @@ -0,0 +1,129 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "reflect" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +type streamReader interface { + copy(n int) []streamReader + getType() reflect.Type + getChunkType() reflect.Type + merge([]streamReader) streamReader + withKey(string) streamReader + close() + toAnyStreamReader() *schema.StreamReader[any] + mergeWithNames([]streamReader, []string) streamReader +} + +type streamReaderPacker[T any] struct { + sr *schema.StreamReader[T] +} + +func (srp streamReaderPacker[T]) close() { + srp.sr.Close() +} + +func (srp streamReaderPacker[T]) copy(n int) []streamReader { + ret := make([]streamReader, n) + srs := srp.sr.Copy(n) + + for i := 0; i < n; i++ { + ret[i] = streamReaderPacker[T]{srs[i]} + } + + return ret +} + +func (srp streamReaderPacker[T]) getType() reflect.Type { + return reflect.TypeOf(srp.sr) +} + +func (srp streamReaderPacker[T]) getChunkType() reflect.Type { + return generic.TypeOf[T]() +} + +func (srp streamReaderPacker[T]) toStreamReaders(srs []streamReader) []*schema.StreamReader[T] { + ret := make([]*schema.StreamReader[T], len(srs)+1) + ret[0] = srp.sr + for i := 1; i < len(ret); i++ { + sr, ok := unpackStreamReader[T](srs[i-1]) + if !ok { + return nil + } + + ret[i] = sr + } + + return ret +} + +func (srp streamReaderPacker[T]) merge(isrs []streamReader) streamReader { + srs := srp.toStreamReaders(isrs) + + sr := schema.MergeStreamReaders(srs) + + return packStreamReader(sr) +} + +func (srp streamReaderPacker[T]) mergeWithNames(isrs []streamReader, names []string) streamReader { + srs := srp.toStreamReaders(isrs) + + sr := schema.InternalMergeNamedStreamReaders(srs, names) + + return packStreamReader(sr) +} + +func (srp streamReaderPacker[T]) withKey(key string) streamReader { + cvt := func(v T) (map[string]any, error) { + return map[string]any{key: v}, nil + } + + ret := schema.StreamReaderWithConvert[T, map[string]any](srp.sr, cvt) + + return packStreamReader(ret) +} + +func (srp streamReaderPacker[T]) toAnyStreamReader() *schema.StreamReader[any] { + return schema.StreamReaderWithConvert(srp.sr, func(t T) (any, error) { + return t, nil + }) +} + +func packStreamReader[T any](sr *schema.StreamReader[T]) streamReader { + return streamReaderPacker[T]{sr} +} + +func unpackStreamReader[T any](isr streamReader) (*schema.StreamReader[T], bool) { + c, ok := isr.(streamReaderPacker[T]) + if ok { + return c.sr, true + } + + typ := generic.TypeOf[T]() + if typ.Kind() == reflect.Interface { + return schema.StreamReaderWithConvert(isr.toAnyStreamReader(), func(t any) (T, error) { + return t.(T), nil + }), true + } + + return nil, false +} diff --git a/compose/stream_reader_test.go b/compose/stream_reader_test.go new file mode 100644 index 0000000..d104a50 --- /dev/null +++ b/compose/stream_reader_test.go @@ -0,0 +1,102 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "io" + "reflect" + "testing" + + "github.com/cloudwego/eino/schema" + "github.com/stretchr/testify/assert" +) + +func TestArrayStreamMerge(t *testing.T) { + + t.Run("unpack_to_equal_type", func(t *testing.T) { + a1 := []int{1, 2, 3} + a2 := []int{4, 5, 6} + a3 := []int{7, 8, 9} + s1 := schema.StreamReaderFromArray(a1) + s2 := schema.StreamReaderFromArray(a2) + s3 := schema.StreamReaderFromArray(a3) + + sp1 := streamReaderPacker[int]{sr: s1} + sp2 := streamReaderPacker[int]{sr: s2} + sp3 := streamReaderPacker[int]{sr: s3} + + sp := sp1.merge([]streamReader{sp2, sp3}) + + sr, ok := unpackStreamReader[int](sp) + if !ok { + t.Fatal("unexpected") + } + + defer sr.Close() + + var result []int + for { + chunk, err := sr.Recv() + if err == io.EOF { + break + } + assert.Nil(t, err) + result = append(result, chunk) + } + if !reflect.DeepEqual(result, append(append(a1, a2...), a3...)) { + t.Fatalf("result: %v error", result) + } + }) + + t.Run("unpack_to_father_type", func(t *testing.T) { + a1 := []*doctor{{say: "a"}, {say: "b"}, {say: "c"}} + a2 := []*doctor{{say: "d"}, {say: "e"}, {say: "f"}} + a3 := []*doctor{{say: "g"}, {say: "h"}, {say: "i"}} + s1 := schema.StreamReaderFromArray(a1) + s2 := schema.StreamReaderFromArray(a2) + s3 := schema.StreamReaderFromArray(a3) + + sp1 := streamReaderPacker[*doctor]{sr: s1} + sp2 := streamReaderPacker[*doctor]{sr: s2} + sp3 := streamReaderPacker[*doctor]{sr: s3} + + sp := sp1.merge([]streamReader{sp2, sp3}) + + sr, ok := unpackStreamReader[person](sp) + assert.True(t, ok) + + defer sr.Close() + + var result []person + for { + chunk, err := sr.Recv() + if err == io.EOF { + break + } + assert.Nil(t, err) + result = append(result, chunk) + } + + baseline := append(append(a1, a2...), a3...) + + assert.Len(t, result, len(baseline)) + + for idx := range result { + assert.Equal(t, baseline[idx].say, result[idx].Say()) + } + }) +} diff --git a/compose/tool_alias_test.go b/compose/tool_alias_test.go new file mode 100644 index 0000000..487132c --- /dev/null +++ b/compose/tool_alias_test.go @@ -0,0 +1,1178 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" +) + +type searchArgs struct { + Query string `json:"query"` +} + +func TestToolNameAliases(t *testing.T) { + ctx := context.Background() + + // Create test tool + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search for information", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string", Desc: "Search query"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "search result", nil + }) + + // Configure aliases + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"search_v1", "query", "find"}, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Test calling tool with alias + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "search_v1", // Using alias + Arguments: `{"query": "test"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Equal(t, "call_1", output[0].ToolCallID) + assert.Contains(t, output[0].Content, "search result") +} + +type searchArgsWithLimit struct { + Query string `json:"query"` + Limit int `json:"limit"` +} + +func TestArgumentsAliases(t *testing.T) { + ctx := context.Background() + + receivedArgs := "" + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search for information", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + "limit": {Type: "integer"}, + }), + }, func(ctx context.Context, args *searchArgsWithLimit) (string, error) { + b, _ := json.Marshal(args) + receivedArgs = string(b) + return "result", nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + ArgumentsAliases: map[string][]string{ + "query": {"q", "search_term"}, + "limit": {"max_results", "count"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Use alias parameters + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "search", + Arguments: `{"q": "test", "max_results": 10}`, // Using aliases + }, + }, + }) + + _, err = node.Invoke(ctx, input) + require.NoError(t, err) + + // Verify tool received canonical parameter names + var args map[string]any + err = json.Unmarshal([]byte(receivedArgs), &args) + require.NoError(t, err) + assert.Equal(t, "test", args["query"]) + assert.Equal(t, float64(10), args["limit"]) + assert.NotContains(t, args, "q") + assert.NotContains(t, args, "max_results") +} + +type emptyArgs struct{} + +func TestAliasConflict(t *testing.T) { + ctx := context.Background() + + tool1 := newTool(&schema.ToolInfo{Name: "search", Desc: "Search"}, func(ctx context.Context, args *emptyArgs) (string, error) { + return "result", nil + }) + tool2 := newTool(&schema.ToolInfo{Name: "query", Desc: "Query"}, func(ctx context.Context, args *emptyArgs) (string, error) { + return "result", nil + }) + + t.Run("tool name alias conflict", func(t *testing.T) { + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1, tool2}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + }, + "query": { + NameAliases: []string{"find"}, // Conflict: find already used by search + }, + }, + } + + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "conflicts with an alias already registered for") + }) + + t.Run("tool name alias conflicts with canonical name", func(t *testing.T) { + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1, tool2}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"query"}, // Conflict: "query" is tool2's canonical name + }, + }, + } + + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "conflicts with existing tool's canonical name") + }) + + t.Run("argument alias conflict", func(t *testing.T) { + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + "limit": {"q"}, // Conflict: q maps to multiple parameters + }, + }, + }, + } + + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "conflicting arg alias") + }) + + t.Run("arg alias conflicts with existing schema property", func(t *testing.T) { + searchWithParams := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + "limit": {Type: "integer"}, + }), + }, func(ctx context.Context, args *emptyArgs) (string, error) { + return "result", nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchWithParams}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + ArgumentsAliases: map[string][]string{ + "limit": {"query"}, // "query" is already a schema property + }, + }, + }, + } + + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "conflicts with existing schema property") + }) +} + +func TestArgumentsAliasesWithHandler(t *testing.T) { + ctx := context.Background() + + executionOrder := []string{} + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + executionOrder = append(executionOrder, "tool_invoke") + return "result", nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + }, + ToolArgumentsHandler: func(ctx context.Context, name, args string) (string, error) { + executionOrder = append(executionOrder, "args_handler") + // Handler receives the original model-returned name (alias) + assert.Equal(t, "search", name) + // Verify alias remapping has already been done + var m map[string]any + err := json.Unmarshal([]byte(args), &m) + require.NoError(t, err) + assert.Contains(t, m, "query") + assert.NotContains(t, m, "q") + return args, nil + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Call with alias name "find" and alias arg "q" + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "test"}`, + }, + }, + }) + + _, err = node.Invoke(ctx, input) + require.NoError(t, err) + + // Verify execution order: alias remapping → ToolArgumentsHandler → tool execution + assert.Equal(t, []string{"args_handler", "tool_invoke"}, executionOrder) +} + +func TestNonExistentToolInAliasConfig(t *testing.T) { + ctx := context.Background() + + tool1 := newTool(&schema.ToolInfo{Name: "search", Desc: "Search"}, func(ctx context.Context, args *emptyArgs) (string, error) { + return "result", nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{tool1}, + ToolAliases: map[string]ToolAliasConfig{ + "non_existent_tool": { // Non-existent tool + NameAliases: []string{"alias1"}, + }, + }, + } + + // Should not error — non-existent tool alias configs are silently skipped + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // The existing tool should still work normally + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "search", + Arguments: `{}`, + }, + }, + }) + output, err := node.Invoke(ctx, input) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "result") +} + +type weatherArgs struct { + Location string `json:"location"` +} + +func TestToolAliasesE2E(t *testing.T) { + ctx := context.Background() + + // Create multiple tools + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search for information", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + "limit": {Type: "integer"}, + }), + }, func(ctx context.Context, args *searchArgsWithLimit) (string, error) { + return "search result", nil + }) + + weatherTool := newTool(&schema.ToolInfo{ + Name: "weather", + Desc: "Get weather information", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "location": {Type: "string"}, + }), + }, func(ctx context.Context, args *weatherArgs) (string, error) { + return "weather result", nil + }) + + // Configure aliases for multiple tools + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool, weatherTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"search_v1", "query"}, + ArgumentsAliases: map[string][]string{ + "query": {"q", "search_term"}, + "limit": {"max_results"}, + }, + }, + "weather": { + NameAliases: []string{"get_weather"}, + ArgumentsAliases: map[string][]string{ + "location": {"loc", "city"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Construct message with multiple tool calls using different aliases + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "search_v1", // Tool name alias + Arguments: `{"q": "test", "max_results": 5}`, // Parameter aliases + }, + }, + { + ID: "call_2", + Function: schema.FunctionCall{ + Name: "get_weather", // Tool name alias + Arguments: `{"city": "Beijing"}`, // Parameter alias + }, + }, + }) + + output, err := node.Invoke(ctx, input) + require.NoError(t, err) + require.Len(t, output, 2) + + // Verify both tools executed successfully + assert.Equal(t, "call_1", output[0].ToolCallID) + assert.Equal(t, "call_2", output[1].ToolCallID) + assert.Contains(t, output[0].Content, "search result") + assert.Contains(t, output[1].Content, "weather result") +} + +func TestRemapArgsEdgeCases(t *testing.T) { + aliasMap := map[string]string{"q": "query"} + + t.Run("empty string", func(t *testing.T) { + result, err := remapArgs("", aliasMap) + assert.NoError(t, err) + assert.Equal(t, "", result) + }) + + t.Run("whitespace only", func(t *testing.T) { + result, err := remapArgs(" ", aliasMap) + assert.NoError(t, err) + assert.Equal(t, " ", result) + }) + + t.Run("non-object JSON", func(t *testing.T) { + result, err := remapArgs(`"hello"`, aliasMap) + assert.NoError(t, err) + assert.Equal(t, `"hello"`, result) + }) + + t.Run("JSON array", func(t *testing.T) { + result, err := remapArgs(`[1,2,3]`, aliasMap) + assert.NoError(t, err) + assert.Equal(t, `[1,2,3]`, result) + }) + + t.Run("invalid JSON", func(t *testing.T) { + result, err := remapArgs(`{invalid`, aliasMap) + assert.NoError(t, err) + assert.Equal(t, `{invalid`, result) + }) + + t.Run("alias and canonical both present", func(t *testing.T) { + // When both alias "q" and canonical "query" exist, alias is kept as-is (not deleted, not overwritten) + result, err := remapArgs(`{"q": "alias_val", "query": "canonical_val"}`, aliasMap) + assert.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal([]byte(result), &m)) + assert.Equal(t, "canonical_val", m["query"]) + assert.Equal(t, "alias_val", m["q"]) + }) + + t.Run("unknown fields preserved", func(t *testing.T) { + result, err := remapArgs(`{"q": "test", "unknown_field": 42}`, aliasMap) + assert.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal([]byte(result), &m)) + assert.Equal(t, "test", m["query"]) + assert.NotContains(t, m, "q") + assert.Equal(t, float64(42), m["unknown_field"]) + }) +} + +func TestCanonicalNameCallWithAliasConfigured(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "result: " + args.Query, nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Call with canonical name and canonical arg — should work normally + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "search", + Arguments: `{"query": "hello"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "result: hello") +} + +func TestEmptyAliasValidation(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{Name: "search", Desc: "Search"}, func(ctx context.Context, args *emptyArgs) (string, error) { + return "result", nil + }) + + t.Run("empty name alias", func(t *testing.T) { + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{""}, + }, + }, + } + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty name alias") + }) + + t.Run("empty arg alias", func(t *testing.T) { + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + ArgumentsAliases: map[string][]string{ + "query": {""}, + }, + }, + }, + } + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty argument alias") + }) + + t.Run("empty canonical arg key", func(t *testing.T) { + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + ArgumentsAliases: map[string][]string{ + "": {"q"}, + }, + }, + }, + } + _, err := NewToolNode(ctx, config) + require.Error(t, err) + assert.Contains(t, err.Error(), "empty canonical argument key") + }) +} + +func TestNameAliasSameAsCanonical(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{Name: "search", Desc: "Search"}, func(ctx context.Context, args *emptyArgs) (string, error) { + return "result", nil + }) + + // Alias same as canonical name — should be tolerated (skip, no error) + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"search", "find"}, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Both canonical and alias should work + for _, name := range []string{"search", "find"} { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: name, + Arguments: `{}`, + }, + }, + }) + output, err := node.Invoke(ctx, input) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "result") + } +} + +func TestToolAliasesWithDynamicToolList(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "search result: " + args.Query, nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Use dynamic ToolList via option — alias should still work + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "dynamic"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolList(searchTool)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "search result: dynamic") +} + +func TestToolNameAliasesStream(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search for information", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "stream result: " + args.Query, nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "hello"}`, + }, + }, + }) + + reader, err := node.Stream(ctx, input) + require.NoError(t, err) + + var chunks [][]*schema.Message + for { + chunk, err := reader.Recv() + if err != nil { + break + } + chunks = append(chunks, chunk) + } + + msgs, err := schema.ConcatMessageArray(chunks) + require.NoError(t, err) + require.Len(t, msgs, 1) + assert.Equal(t, "call_1", msgs[0].ToolCallID) + assert.Contains(t, msgs[0].Content, "stream result: hello") +} + +func TestEnhancedToolWithAliases(t *testing.T) { + ctx := context.Background() + + enhancedTool := &enhancedInvokableTool{ + info: &schema.ToolInfo{ + Name: "search", + Desc: "Enhanced search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, + fn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "enhanced: " + input.Text}, + }, + }, nil + }, + } + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{enhancedTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Call with alias name and alias arg + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "test"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Equal(t, "call_1", output[0].ToolCallID) + // Verify arg alias was remapped: "q" → "query" in the JSON passed to enhanced tool + assert.Contains(t, output[0].UserInputMultiContent[0].Text, "enhanced:") +} + +func TestDynamicToolListAliasRemoved(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "search result", nil + }) + + weatherTool := newTool(&schema.ToolInfo{ + Name: "weather", + Desc: "Weather", + }, func(ctx context.Context, args *emptyArgs) (string, error) { + return "weather result", nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool, weatherTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + // Dynamic tool list only contains weatherTool — "search" and its alias "find" should not be available + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{}`, + }, + }, + }) + + _, err = node.Invoke(ctx, input, WithToolList(weatherTool)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not found") +} + +func TestToolAliasesOptionOverridesGlobal(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "search result: " + args.Query, nil + }) + + weatherTool := newTool(&schema.ToolInfo{ + Name: "weather", + Desc: "Weather", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "location": {Type: "string"}, + }), + }, func(ctx context.Context, args *weatherArgs) (string, error) { + return "weather result: " + args.Location, nil + }) + + // Global aliases: search has alias "find" + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool, weatherTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + t.Run("opt ToolAliases overrides global in Invoke", func(t *testing.T) { + // opt.ToolAliases defines "lookup" as alias for search (not "find") + optAliases := map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"lookup"}, + ArgumentsAliases: map[string][]string{ + "query": {"keyword"}, + }, + }, + } + + // "lookup" should work with opt aliases + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "lookup", + Arguments: `{"keyword": "test"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolList(searchTool), WithToolAliases(optAliases)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "search result: test") + + // "find" (global alias) should NOT work when opt.ToolAliases is set + input2 := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_2", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "test"}`, + }, + }, + }) + + _, err = node.Invoke(ctx, input2, WithToolList(searchTool), WithToolAliases(optAliases)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not found") + }) + + t.Run("opt ToolAliases overrides global in Stream", func(t *testing.T) { + optAliases := map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"lookup"}, + ArgumentsAliases: map[string][]string{ + "query": {"keyword"}, + }, + }, + } + + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "lookup", + Arguments: `{"keyword": "stream_test"}`, + }, + }, + }) + + reader, err := node.Stream(ctx, input, WithToolList(searchTool), WithToolAliases(optAliases)) + require.NoError(t, err) + + var chunks [][]*schema.Message + for { + chunk, err := reader.Recv() + if err != nil { + break + } + chunks = append(chunks, chunk) + } + + msgs, err := schema.ConcatMessageArray(chunks) + require.NoError(t, err) + require.Len(t, msgs, 1) + assert.Contains(t, msgs[0].Content, "search result: stream_test") + }) + + t.Run("nil opt ToolAliases falls back to global filtered", func(t *testing.T) { + // No WithToolAliases — should use global "find" alias, filtered by ToolList + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "fallback"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolList(searchTool)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "search result: fallback") + }) + + t.Run("opt ToolAliases only without ToolList replaces global", func(t *testing.T) { + // Only WithToolAliases, no WithToolList — should use global tools with opt aliases + optAliases := map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"lookup"}, + ArgumentsAliases: map[string][]string{ + "query": {"keyword"}, + }, + }, + } + + // "lookup" (opt alias) should work + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "lookup", + Arguments: `{"keyword": "only_alias"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolAliases(optAliases)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "search result: only_alias") + + // "find" (global alias) should NOT work when opt.ToolAliases replaces global + input2 := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_2", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "test"}`, + }, + }, + }) + + _, err = node.Invoke(ctx, input2, WithToolAliases(optAliases)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not found") + }) + + t.Run("opt ToolAliases only without ToolList in Stream", func(t *testing.T) { + optAliases := map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"lookup"}, + }, + } + + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "lookup", + Arguments: `{"query": "stream_only_alias"}`, + }, + }, + }) + + reader, err := node.Stream(ctx, input, WithToolAliases(optAliases)) + require.NoError(t, err) + + var chunks [][]*schema.Message + for { + chunk, err := reader.Recv() + if err != nil { + break + } + chunks = append(chunks, chunk) + } + + msgs, err := schema.ConcatMessageArray(chunks) + require.NoError(t, err) + require.Len(t, msgs, 1) + assert.Contains(t, msgs[0].Content, "search result: stream_only_alias") + }) +} + +func TestAliasConfigForToolAddedViaOption(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "search result: " + args.Query, nil + }) + + weatherTool := newTool(&schema.ToolInfo{ + Name: "weather", + Desc: "Weather", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "location": {Type: "string"}, + }), + }, func(ctx context.Context, args *weatherArgs) (string, error) { + return "weather result: " + args.Location, nil + }) + + // New with only searchTool, but alias config includes weather tool + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + ArgumentsAliases: map[string][]string{ + "query": {"q"}, + }, + }, + "weather": { + NameAliases: []string{"forecast"}, + ArgumentsAliases: map[string][]string{ + "location": {"loc"}, + }, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + t.Run("weather alias works when tool passed via option", func(t *testing.T) { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "forecast", + Arguments: `{"loc": "Beijing"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolList(searchTool, weatherTool)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "weather result: Beijing") + }) + + t.Run("search alias still works with option tool list", func(t *testing.T) { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"q": "test"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolList(searchTool, weatherTool)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "search result: test") + }) +} + +func TestOptionWithToolListAndToolAliases(t *testing.T) { + ctx := context.Background() + + searchTool := newTool(&schema.ToolInfo{ + Name: "search", + Desc: "Search", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: "string"}, + }), + }, func(ctx context.Context, args *searchArgs) (string, error) { + return "search result: " + args.Query, nil + }) + + weatherTool := newTool(&schema.ToolInfo{ + Name: "weather", + Desc: "Weather", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "location": {Type: "string"}, + }), + }, func(ctx context.Context, args *weatherArgs) (string, error) { + return "weather result: " + args.Location, nil + }) + + config := &ToolsNodeConfig{ + Tools: []tool.BaseTool{searchTool}, + ToolAliases: map[string]ToolAliasConfig{ + "search": { + NameAliases: []string{"find"}, + }, + }, + } + + node, err := NewToolNode(ctx, config) + require.NoError(t, err) + + t.Run("opt aliases override global when both tool list and aliases provided", func(t *testing.T) { + optAliases := map[string]ToolAliasConfig{ + "weather": { + NameAliases: []string{"forecast"}, + ArgumentsAliases: map[string][]string{ + "location": {"loc"}, + }, + }, + } + + // "forecast" should work via opt aliases + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_1", + Function: schema.FunctionCall{ + Name: "forecast", + Arguments: `{"loc": "Shanghai"}`, + }, + }, + }) + + output, err := node.Invoke(ctx, input, WithToolList(searchTool, weatherTool), WithToolAliases(optAliases)) + require.NoError(t, err) + require.Len(t, output, 1) + assert.Contains(t, output[0].Content, "weather result: Shanghai") + + // "find" (global alias) should NOT work when opt aliases override + input2 := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call_2", + Function: schema.FunctionCall{ + Name: "find", + Arguments: `{"query": "test"}`, + }, + }, + }) + + _, err = node.Invoke(ctx, input2, WithToolList(searchTool, weatherTool), WithToolAliases(optAliases)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not found") + }) +} diff --git a/compose/tool_node.go b/compose/tool_node.go new file mode 100644 index 0000000..f65037e --- /dev/null +++ b/compose/tool_node.go @@ -0,0 +1,1310 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "runtime/debug" + "sort" + "strings" + "sync" + + "github.com/bytedance/sonic" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/internal/safe" + "github.com/cloudwego/eino/schema" +) + +type toolsNodeOptions struct { + ToolOptions []tool.Option + ToolList []tool.BaseTool + + ToolAliases map[string]ToolAliasConfig +} + +// ToolsNodeOption is the option func type for ToolsNode. +type ToolsNodeOption func(o *toolsNodeOptions) + +// WithToolOption adds tool options to the ToolsNode. +func WithToolOption(opts ...tool.Option) ToolsNodeOption { + return func(o *toolsNodeOptions) { + o.ToolOptions = append(o.ToolOptions, opts...) + } +} + +// WithToolList sets the tool list for the ToolsNode. +func WithToolList(tool ...tool.BaseTool) ToolsNodeOption { + return func(o *toolsNodeOptions) { + o.ToolList = tool + } +} + +// WithToolAliases sets the tool aliases for the ToolsNode call option. +// When used with WithToolList, it overrides the global alias configuration for the dynamic tool list. +// When used alone (without WithToolList), it replaces the global alias configuration while keeping the original tool list. +func WithToolAliases(toolAliases map[string]ToolAliasConfig) ToolsNodeOption { + return func(o *toolsNodeOptions) { + o.ToolAliases = toolAliases + } +} + +// ToolsNode represents a node capable of executing tools within a graph. +// The Graph Node interface is defined as follows: +// +// Invoke(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) ([]*schema.Message, error) +// Stream(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error) +// +// Input: An AssistantMessage containing ToolCalls +// Output: An array of ToolMessage where the order of elements corresponds to the order of ToolCalls in the input +type ToolsNode struct { + tuple *toolsTuple + tools []tool.BaseTool + unknownToolHandler func(ctx context.Context, name, input string) (string, error) + executeSequentially bool + toolArgumentsHandler func(ctx context.Context, name, input string) (string, error) + toolCallMiddlewares []InvokableToolMiddleware + streamToolCallMiddlewares []StreamableToolMiddleware + enhancedToolCallMiddlewares []EnhancedInvokableToolMiddleware + enhancedStreamToolCallMiddlewares []EnhancedStreamableToolMiddleware + toolAliasConfigs map[string]ToolAliasConfig +} + +// ToolInput represents the input parameters for a tool call execution. +type ToolInput struct { + // Name is the name of the tool to be executed. + Name string + // Arguments contains the arguments for the tool call. + Arguments string + // CallID is the unique identifier for this tool call. + CallID string + // CallOptions contains tool options for the execution. + CallOptions []tool.Option +} + +// ToolOutput represents the result of a non-streaming tool call execution. +type ToolOutput struct { + // Result contains the string output from the tool execution. + Result string +} + +// StreamToolOutput represents the result of a streaming tool call execution. +type StreamToolOutput struct { + // Result is a stream reader that provides access to the tool's streaming output. + Result *schema.StreamReader[string] +} + +// EnhancedInvokableToolOutput represents the result of a non-streaming enhanced tool call execution. +// It supports returning structured multimodal content (text, images, audio, video, files) from tools. +type EnhancedInvokableToolOutput struct { + // Result contains the structured multimodal output from the tool execution. + Result *schema.ToolResult +} + +// EnhancedStreamableToolOutput represents the result of a streaming enhanced tool call execution. +// It provides a stream reader for accessing multimodal content progressively. +type EnhancedStreamableToolOutput struct { + // Result is a stream reader that provides access to the tool's streaming multimodal output. + Result *schema.StreamReader[*schema.ToolResult] +} + +// InvokableToolEndpoint is the function signature for non-streaming tool calls. +type InvokableToolEndpoint func(ctx context.Context, input *ToolInput) (*ToolOutput, error) + +// StreamableToolEndpoint is the function signature for streaming tool calls. +type StreamableToolEndpoint func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) + +type EnhancedInvokableToolEndpoint func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) + +type EnhancedStreamableToolEndpoint func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) + +// InvokableToolMiddleware is a function that wraps InvokableToolEndpoint to add custom processing logic. +// It can be used to intercept, modify, or enhance tool call execution for non-streaming tools. +type InvokableToolMiddleware func(InvokableToolEndpoint) InvokableToolEndpoint + +// StreamableToolMiddleware is a function that wraps StreamableToolEndpoint to add custom processing logic. +// It can be used to intercept, modify, or enhance tool call execution for streaming tools. +type StreamableToolMiddleware func(StreamableToolEndpoint) StreamableToolEndpoint + +type EnhancedInvokableToolMiddleware func(EnhancedInvokableToolEndpoint) EnhancedInvokableToolEndpoint + +type EnhancedStreamableToolMiddleware func(EnhancedStreamableToolEndpoint) EnhancedStreamableToolEndpoint + +// ToolMiddleware groups middleware hooks for invokable and streamable tool calls. +type ToolMiddleware struct { + // Invokable contains middleware function for non-streaming tool calls. + // Note: This middleware only applies to tools that implement the InvokableTool interface. + Invokable InvokableToolMiddleware + + // Streamable contains middleware function for streaming tool calls. + // Note: This middleware only applies to tools that implement the StreamableTool interface. + Streamable StreamableToolMiddleware + + // EnhancedInvokable contains middleware function for non-streaming enhanced tool calls. + // Note: This middleware only applies to tools that implement the EnhancedInvokableTool interface. + EnhancedInvokable EnhancedInvokableToolMiddleware + + // EnhancedStreamable contains middleware function for streaming enhanced tool calls. + // Note: This middleware only applies to tools that implement the EnhancedStreamableTool interface. + EnhancedStreamable EnhancedStreamableToolMiddleware +} + +// ToolAliasConfig configures name and argument aliases for a single tool. +type ToolAliasConfig struct { + // NameAliases are alternative names for this tool. + // If the model returns any of these names, it will be resolved to the canonical tool name. + NameAliases []string + + // ArgumentsAliases maps canonical argument keys to their alias lists. + // key=canonical, value=[]alias. Applied to top-level JSON keys before tool execution. + // Example: {"query": ["q", "search_term"], "limit": ["max_results", "count"]} + ArgumentsAliases map[string][]string +} + +// ToolsNodeConfig is the config for ToolsNode. +type ToolsNodeConfig struct { + // Tools specify the list of tools can be called which are BaseTool but must implement InvokableTool or StreamableTool. + Tools []tool.BaseTool + + // ToolAliases configures name and argument aliases for tools. + // Key is the canonical tool name, value defines its aliases. + // This field is optional. When provided, tool name aliases will be resolved during tool dispatch, + // and argument aliases will be remapped before ToolArgumentsHandler (if configured) and tool execution. + // Execution order: ArgumentsAliases remapping → ToolArgumentsHandler → tool execution + ToolAliases map[string]ToolAliasConfig + + // UnknownToolsHandler handles tool calls for non-existent tools when LLM hallucinates. + // This field is optional. When not set, calling a non-existent tool will result in an error. + // When provided, if the LLM attempts to call a tool that doesn't exist in the Tools list, + // this handler will be invoked instead of returning an error, allowing graceful handling of hallucinated tools. + // Parameters: + // - ctx: The context for the tool call + // - name: The name of the non-existent tool + // - input: The tool call input generated by llm + // Returns: + // - string: The response to be returned as if the tool was executed + // - error: Any error that occurred during handling + UnknownToolsHandler func(ctx context.Context, name, input string) (string, error) + + // ExecuteSequentially determines whether tool calls should be executed sequentially (in order) or in parallel. + // When set to true, tool calls will be executed one after another in the order they appear in the input message. + // When set to false (default), tool calls will be executed in parallel. + ExecuteSequentially bool + + // ToolArgumentsHandler allows handling of tool arguments before execution. + // When provided, this function will be called for each tool call to process the arguments. + // Parameters: + // - ctx: The context for the tool call + // - name: The name of the tool being called + // - arguments: The original arguments string for the tool + // Returns: + // - string: The processed arguments string to be used for tool execution + // - error: Any error that occurred during preprocessing + ToolArgumentsHandler func(ctx context.Context, name, arguments string) (string, error) + + // ToolCallMiddlewares configures middleware for tool calls. + // Each element can contain Invokable and/or Streamable middleware. + // Invokable middleware only applies to tools implementing InvokableTool interface. + // Streamable middleware only applies to tools implementing StreamableTool interface. + ToolCallMiddlewares []ToolMiddleware +} + +// NewToolNode creates a new ToolsNode. +// e.g. +// +// conf := &ToolsNodeConfig{ +// Tools: []tool.BaseTool{invokableTool1, streamableTool2}, +// } +// toolsNode, err := NewToolNode(ctx, conf) +func NewToolNode(ctx context.Context, conf *ToolsNodeConfig) (*ToolsNode, error) { + var middlewares []InvokableToolMiddleware + var streamMiddlewares []StreamableToolMiddleware + var enhancedInvokableMiddlewares []EnhancedInvokableToolMiddleware + var enhancedStreamableMiddlewares []EnhancedStreamableToolMiddleware + + for _, m := range conf.ToolCallMiddlewares { + if m.Invokable != nil { + middlewares = append(middlewares, m.Invokable) + } + if m.Streamable != nil { + streamMiddlewares = append(streamMiddlewares, m.Streamable) + } + if m.EnhancedInvokable != nil { + enhancedInvokableMiddlewares = append(enhancedInvokableMiddlewares, m.EnhancedInvokable) + } + if m.EnhancedStreamable != nil { + enhancedStreamableMiddlewares = append(enhancedStreamableMiddlewares, m.EnhancedStreamable) + } + } + + params := convToolsParams{ + tools: conf.Tools, + aliasConfigs: conf.ToolAliases, + } + params.middlewares.invokable = middlewares + params.middlewares.streamable = streamMiddlewares + params.middlewares.enhancedInvokable = enhancedInvokableMiddlewares + params.middlewares.enhancedStreamable = enhancedStreamableMiddlewares + tuple, err := convTools(ctx, params) + if err != nil { + return nil, err + } + + return &ToolsNode{ + tuple: tuple, + tools: conf.Tools, + unknownToolHandler: conf.UnknownToolsHandler, + executeSequentially: conf.ExecuteSequentially, + toolArgumentsHandler: conf.ToolArgumentsHandler, + toolCallMiddlewares: middlewares, + streamToolCallMiddlewares: streamMiddlewares, + enhancedToolCallMiddlewares: enhancedInvokableMiddlewares, + enhancedStreamToolCallMiddlewares: enhancedStreamableMiddlewares, + toolAliasConfigs: conf.ToolAliases, + }, nil +} + +// ToolsInterruptAndRerunExtra carries interrupt metadata for ToolsNode reruns. +type ToolsInterruptAndRerunExtra struct { + // ToolCalls contains all tool calls from the original assistant message. + ToolCalls []schema.ToolCall + + // ExecutedTools maps tool call IDs to their string output for successfully executed standard tools. + ExecutedTools map[string]string + + // ExecutedEnhancedTools maps tool call IDs to their structured multimodal output for successfully executed enhanced tools. + ExecutedEnhancedTools map[string]*schema.ToolResult + + // RerunTools contains the IDs of tool calls that need to be re-executed. + RerunTools []string + + // RerunExtraMap stores additional metadata for each tool call that needs rerun, keyed by tool call ID. + RerunExtraMap map[string]any +} + +func init() { + schema.RegisterName[*ToolsInterruptAndRerunExtra]("_eino_compose_tools_interrupt_and_rerun_extra") + schema.RegisterName[*toolsInterruptAndRerunState]("_eino_compose_tools_interrupt_and_rerun_state") +} + +type toolsInterruptAndRerunState struct { + Input *schema.Message + ExecutedTools map[string]string + ExecutedEnhancedTools map[string]*schema.ToolResult + RerunTools []string +} + +type toolsTuple struct { + indexes map[string]int + meta []*executorMeta + endpoints []InvokableToolEndpoint + streamEndpoints []StreamableToolEndpoint + enhancedInvokableEndpoints []EnhancedInvokableToolEndpoint + enhancedStreamableEndpoints []EnhancedStreamableToolEndpoint + // argsAliasMap stores reverse argument alias mappings for each tool. + // key: canonical tool name, value: map[aliasKey]canonicalKey (alias → canonical direction) + argsAliasMap map[string]map[string]string + // canonicalNames stores the canonical name for each tool index + canonicalNames []string + // toolInfos stores the ToolInfo for each tool index, used for alias validation + toolInfos []*schema.ToolInfo +} + +// remapArgs replaces alias keys in the JSON arguments string with canonical keys. +// aliasMap: alias → canonical mapping +func remapArgs(args string, aliasMap map[string]string) (string, error) { + if len(aliasMap) == 0 { + return args, nil + } + + trimmed := strings.TrimSpace(args) + if trimmed == "" || trimmed[0] != '{' { + return args, nil + } + + var m map[string]json.RawMessage + if err := sonic.Unmarshal([]byte(args), &m); err != nil { + return args, nil + } + + changed := false + for alias, canonical := range aliasMap { + if v, ok := m[alias]; ok { + // Only replace if canonical key doesn't exist. + // If both alias and canonical are present (e.g. {"q":"a","query":"b"}), + // the alias key is kept as-is and passed through as an unknown field. + if _, exists := m[canonical]; !exists { + m[canonical] = v + delete(m, alias) + changed = true + } + } + } + + if !changed { + return args, nil + } + + b, err := sonic.Marshal(m) + return string(b), err +} + +type convToolsParams struct { + tools []tool.BaseTool + middlewares struct { + invokable []InvokableToolMiddleware + streamable []StreamableToolMiddleware + enhancedInvokable []EnhancedInvokableToolMiddleware + enhancedStreamable []EnhancedStreamableToolMiddleware + } + aliasConfigs map[string]ToolAliasConfig +} + +func (t *toolsTuple) applyAliasConfigs(aliasConfigs map[string]ToolAliasConfig) error { + t.argsAliasMap = make(map[string]map[string]string) + + sortedToolNames := make([]string, 0, len(aliasConfigs)) + for toolName := range aliasConfigs { + sortedToolNames = append(sortedToolNames, toolName) + } + sort.Strings(sortedToolNames) + + for _, toolName := range sortedToolNames { + aliasConfig := aliasConfigs[toolName] + var ( + toolIdx int + exists bool + ) + if toolIdx, exists = t.indexes[toolName]; !exists { + continue + } + + if err := t.applyNameAliases(toolName, toolIdx, aliasConfig.NameAliases); err != nil { + return err + } + + if err := t.applyArgsAliases(toolName, toolIdx, aliasConfig.ArgumentsAliases); err != nil { + return err + } + } + + return nil +} + +// applyNameAliases validates and registers name aliases for a single tool into the indexes map. +func (t *toolsTuple) applyNameAliases(toolName string, toolIdx int, nameAliases []string) error { + for _, alias := range nameAliases { + if strings.TrimSpace(alias) == "" { + return fmt.Errorf("tool '%s' has empty name alias", toolName) + } + if existingIdx, conflict := t.indexes[alias]; conflict { + if existingIdx != toolIdx { + conflictToolName := t.canonicalNames[existingIdx] + if alias == conflictToolName { + return fmt.Errorf("tool '%s': name alias '%s' conflicts with existing tool's canonical name", toolName, alias) + } + return fmt.Errorf("tool '%s': name alias '%s' conflicts with an alias already registered for tool '%s'", toolName, alias, conflictToolName) + } + continue + } + t.indexes[alias] = toolIdx + } + return nil +} + +// applyArgsAliases validates argument aliases against the tool schema and builds a reverse alias map for a single tool. +func (t *toolsTuple) applyArgsAliases(toolName string, toolIdx int, argumentsAliases map[string][]string) error { + if len(argumentsAliases) == 0 { + return nil + } + + schemaKeys := make(map[string]bool) + if info := t.toolInfos[toolIdx]; info != nil && info.ParamsOneOf != nil { + js, err := info.ParamsOneOf.ToJSONSchema() + if err != nil { + return fmt.Errorf("tool '%s': failed to parse JSON schema for alias validation: %w", toolName, err) + } + if js != nil && js.Properties != nil { + for pair := js.Properties.Oldest(); pair != nil; pair = pair.Next() { + schemaKeys[pair.Key] = true + } + } + } + + reverseMap := make(map[string]string) + sortedCanonicals := make([]string, 0, len(argumentsAliases)) + for canonical := range argumentsAliases { + sortedCanonicals = append(sortedCanonicals, canonical) + } + sort.Strings(sortedCanonicals) + + for _, canonical := range sortedCanonicals { + aliases := argumentsAliases[canonical] + if strings.TrimSpace(canonical) == "" { + return fmt.Errorf("tool '%s' has empty canonical argument key", toolName) + } + if strings.Contains(canonical, ".") { + return fmt.Errorf("tool '%s' has unsupported '.' in canonical argument key '%s': nested field matching is not yet supported", + toolName, canonical) + } + for _, alias := range aliases { + if strings.TrimSpace(alias) == "" { + return fmt.Errorf("tool '%s' has empty argument alias for canonical key '%s'", toolName, canonical) + } + if schemaKeys[alias] { + return fmt.Errorf("tool '%s' has arg alias '%s' that conflicts with existing schema property '%s'", + toolName, alias, alias) + } + if existingCanonical, conflict := reverseMap[alias]; conflict { + return fmt.Errorf("tool '%s' has conflicting arg alias '%s' mapped to both '%s' and '%s'", + toolName, alias, existingCanonical, canonical) + } + reverseMap[alias] = canonical + } + } + t.argsAliasMap[toolName] = reverseMap + + return nil +} + +func convTools(ctx context.Context, params convToolsParams) (*toolsTuple, error) { + ret := &toolsTuple{ + indexes: make(map[string]int), + meta: make([]*executorMeta, len(params.tools)), + endpoints: make([]InvokableToolEndpoint, len(params.tools)), + streamEndpoints: make([]StreamableToolEndpoint, len(params.tools)), + enhancedInvokableEndpoints: make([]EnhancedInvokableToolEndpoint, len(params.tools)), + enhancedStreamableEndpoints: make([]EnhancedStreamableToolEndpoint, len(params.tools)), + canonicalNames: make([]string, len(params.tools)), + toolInfos: make([]*schema.ToolInfo, len(params.tools)), + } + for idx, bt := range params.tools { + tl, err := bt.Info(ctx) + if err != nil { + return nil, fmt.Errorf("(NewToolNode) failed to get tool info at idx= %d: %w", idx, err) + } + + toolName := tl.Name + var ( + st tool.StreamableTool + it tool.InvokableTool + eiTool tool.EnhancedInvokableTool + esTool tool.EnhancedStreamableTool + + invokable InvokableToolEndpoint + streamable StreamableToolEndpoint + enhancedInvokable EnhancedInvokableToolEndpoint + enhancedStreamable EnhancedStreamableToolEndpoint + + ok bool + meta *executorMeta + ) + + meta = parseExecutorInfoFromComponent(components.ComponentOfTool, bt) + + if st, ok = bt.(tool.StreamableTool); ok { + streamable = wrapStreamToolCall(st, params.middlewares.streamable, !meta.isComponentCallbackEnabled) + } + + if it, ok = bt.(tool.InvokableTool); ok { + invokable = wrapToolCall(it, params.middlewares.invokable, !meta.isComponentCallbackEnabled) + } + + if eiTool, ok = bt.(tool.EnhancedInvokableTool); ok { + enhancedInvokable = wrapEnhancedInvokableToolCall(eiTool, params.middlewares.enhancedInvokable, !meta.isComponentCallbackEnabled) + } + + if esTool, ok = bt.(tool.EnhancedStreamableTool); ok { + enhancedStreamable = wrapEnhancedStreamableToolCall(esTool, params.middlewares.enhancedStreamable, !meta.isComponentCallbackEnabled) + } + + if st == nil && it == nil && eiTool == nil && esTool == nil { + return nil, fmt.Errorf("tool %s is not invokable, streamable, enhanced invokable or enhanced streamable", toolName) + } + if streamable == nil && invokable != nil { + streamable = invokableToStreamable(invokable) + } + if invokable == nil && streamable != nil { + invokable = streamableToInvokable(streamable) + } + + if enhancedStreamable == nil && enhancedInvokable != nil { + enhancedStreamable = enhancedInvokableToEnhancedStreamable(enhancedInvokable) + } + if enhancedInvokable == nil && enhancedStreamable != nil { + enhancedInvokable = enhancedStreamableToEnhancedInvokable(enhancedStreamable) + } + + ret.indexes[toolName] = idx + ret.meta[idx] = meta + ret.endpoints[idx] = invokable + ret.streamEndpoints[idx] = streamable + ret.enhancedInvokableEndpoints[idx] = enhancedInvokable + ret.enhancedStreamableEndpoints[idx] = enhancedStreamable + ret.canonicalNames[idx] = toolName + ret.toolInfos[idx] = tl + } + + if len(params.aliasConfigs) > 0 { + if err := ret.applyAliasConfigs(params.aliasConfigs); err != nil { + return nil, err + } + } + + return ret, nil +} + +func wrapToolCall(it tool.InvokableTool, middlewares []InvokableToolMiddleware, needCallback bool) InvokableToolEndpoint { + middleware := func(next InvokableToolEndpoint) InvokableToolEndpoint { + for i := len(middlewares) - 1; i >= 0; i-- { + next = middlewares[i](next) + } + return next + } + if needCallback { + it = &invokableToolWithCallback{it: it} + } + return middleware(func(ctx context.Context, input *ToolInput) (*ToolOutput, error) { + result, err := it.InvokableRun(ctx, input.Arguments, input.CallOptions...) + if err != nil { + return nil, err + } + return &ToolOutput{Result: result}, nil + }) +} + +func wrapStreamToolCall(st tool.StreamableTool, middlewares []StreamableToolMiddleware, needCallback bool) StreamableToolEndpoint { + middleware := func(next StreamableToolEndpoint) StreamableToolEndpoint { + for i := len(middlewares) - 1; i >= 0; i-- { + next = middlewares[i](next) + } + return next + } + if needCallback { + st = &streamableToolWithCallback{st: st} + } + return middleware(func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) { + result, err := st.StreamableRun(ctx, input.Arguments, input.CallOptions...) + if err != nil { + return nil, err + } + return &StreamToolOutput{Result: result}, nil + }) +} + +func wrapEnhancedInvokableToolCall(eiTool tool.EnhancedInvokableTool, middlewares []EnhancedInvokableToolMiddleware, needCallback bool) EnhancedInvokableToolEndpoint { + middleware := func(next EnhancedInvokableToolEndpoint) EnhancedInvokableToolEndpoint { + for i := len(middlewares) - 1; i >= 0; i-- { + next = middlewares[i](next) + } + return next + } + if needCallback { + eiTool = &enhancedInvokableToolWithCallback{eiTool: eiTool} + } + return middleware(func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) { + result, err := eiTool.InvokableRun(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...) + if err != nil { + return nil, err + } + return &EnhancedInvokableToolOutput{Result: result}, nil + }) +} + +func wrapEnhancedStreamableToolCall(est tool.EnhancedStreamableTool, middlewares []EnhancedStreamableToolMiddleware, needCallback bool) EnhancedStreamableToolEndpoint { + middleware := func(next EnhancedStreamableToolEndpoint) EnhancedStreamableToolEndpoint { + for i := len(middlewares) - 1; i >= 0; i-- { + next = middlewares[i](next) + } + return next + } + if needCallback { + est = &enhancedStreamableToolWithCallback{est: est} + } + return middleware(func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) { + result, err := est.StreamableRun(ctx, &schema.ToolArgument{Text: input.Arguments}, input.CallOptions...) + if err != nil { + return nil, err + } + return &EnhancedStreamableToolOutput{Result: result}, nil + }) +} + +type invokableToolWithCallback struct { + it tool.InvokableTool +} + +func (i *invokableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) { + return i.it.Info(ctx) +} + +func (i *invokableToolWithCallback) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + return invokeWithCallbacks(i.it.InvokableRun)(ctx, argumentsInJSON, opts...) +} + +type streamableToolWithCallback struct { + st tool.StreamableTool +} + +func (s *streamableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) { + return s.st.Info(ctx) +} + +func (s *streamableToolWithCallback) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + return streamWithCallbacks(s.st.StreamableRun)(ctx, argumentsInJSON, opts...) +} + +type enhancedInvokableToolWithCallback struct { + eiTool tool.EnhancedInvokableTool +} + +func (e *enhancedInvokableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) { + return e.eiTool.Info(ctx) +} + +func (e *enhancedInvokableToolWithCallback) InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + return invokeEnhancedWithCallbacks(e.eiTool.InvokableRun)(ctx, toolArgument, opts...) +} + +type enhancedStreamableToolWithCallback struct { + est tool.EnhancedStreamableTool +} + +func (e *enhancedStreamableToolWithCallback) Info(ctx context.Context) (*schema.ToolInfo, error) { + return e.est.Info(ctx) +} + +func (e *enhancedStreamableToolWithCallback) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return streamEnhancedWithCallbacks(e.est.StreamableRun)(ctx, toolArgument, opts...) +} + +func streamableToInvokable(e StreamableToolEndpoint) InvokableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*ToolOutput, error) { + so, err := e(ctx, input) + if err != nil { + return nil, err + } + o, err := concatStreamReader(so.Result) + if err != nil { + return nil, fmt.Errorf("failed to concat StreamableTool output message stream: %w", err) + } + return &ToolOutput{Result: o}, nil + } +} + +func invokableToStreamable(e InvokableToolEndpoint) StreamableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) { + o, err := e(ctx, input) + if err != nil { + return nil, err + } + return &StreamToolOutput{Result: schema.StreamReaderFromArray([]string{o.Result})}, nil + } +} + +func enhancedStreamableToEnhancedInvokable(e EnhancedStreamableToolEndpoint) EnhancedInvokableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) { + so, err := e(ctx, input) + if err != nil { + return nil, err + } + o, err := concatStreamReader(so.Result) + if err != nil { + return nil, fmt.Errorf("failed to concat EnhancedStreamableTool output message stream: %w", err) + } + return &EnhancedInvokableToolOutput{Result: o}, nil + } +} + +func enhancedInvokableToEnhancedStreamable(e EnhancedInvokableToolEndpoint) EnhancedStreamableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) { + o, err := e(ctx, input) + if err != nil { + return nil, err + } + return &EnhancedStreamableToolOutput{Result: schema.StreamReaderFromArray([]*schema.ToolResult{o.Result})}, nil + } +} + +func invokeEnhancedWithCallbacks(i func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error)) func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.ToolResult, error) { + return runWithCallbacks(i, onStart[*schema.ToolArgument], onEnd[*schema.ToolResult], onError) +} + +func streamEnhancedWithCallbacks(s func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error)) func(ctx context.Context, toolArgument *schema.ToolArgument, opts ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return runWithCallbacks(s, onStart[*schema.ToolArgument], onEndWithStreamOutput[*schema.ToolResult], onError) +} + +type toolCallTask struct { + // in + endpoint InvokableToolEndpoint + streamEndpoint StreamableToolEndpoint + enhancedInvokableEndpoint EnhancedInvokableToolEndpoint + enhancedStreamableEndpoint EnhancedStreamableToolEndpoint + meta *executorMeta + name string + arg string + callID string + useEnhanced bool + + // out + executed bool + output string + sOutput *schema.StreamReader[string] + enhancedOutput *schema.ToolResult + enhancedSOutput *schema.StreamReader[*schema.ToolResult] + err error +} + +func (tn *ToolsNode) genToolCallTasks(ctx context.Context, tuple *toolsTuple, + input *schema.Message, executedTools map[string]string, executedEnhancedTools map[string]*schema.ToolResult, isStream bool) ([]toolCallTask, error) { + + if input.Role != schema.Assistant { + return nil, fmt.Errorf("expected message role is Assistant, got %s", input.Role) + } + + n := len(input.ToolCalls) + if n == 0 { + return nil, errors.New("no tool call found in input message") + } + + toolCallTasks := make([]toolCallTask, n) + + for i := 0; i < n; i++ { + toolCall := input.ToolCalls[i] + if enhancedResult, executed := executedEnhancedTools[toolCall.ID]; executed { + toolCallTasks[i].name = toolCall.Function.Name + toolCallTasks[i].arg = toolCall.Function.Arguments + toolCallTasks[i].callID = toolCall.ID + toolCallTasks[i].executed = true + toolCallTasks[i].useEnhanced = true + if isStream { + toolCallTasks[i].enhancedSOutput = schema.StreamReaderFromArray([]*schema.ToolResult{enhancedResult}) + } else { + toolCallTasks[i].enhancedOutput = enhancedResult + } + continue + } + if result, executed := executedTools[toolCall.ID]; executed { + toolCallTasks[i].name = toolCall.Function.Name + toolCallTasks[i].arg = toolCall.Function.Arguments + toolCallTasks[i].callID = toolCall.ID + toolCallTasks[i].executed = true + toolCallTasks[i].useEnhanced = false + if isStream { + toolCallTasks[i].sOutput = schema.StreamReaderFromArray([]string{result}) + } else { + toolCallTasks[i].output = result + } + continue + } + index, ok := tuple.indexes[toolCall.Function.Name] + if !ok { + if tn.unknownToolHandler == nil { + return nil, fmt.Errorf("tool %s not found in toolsNode indexes", toolCall.Function.Name) + } + toolCallTasks[i] = newUnknownToolTask(toolCall.Function.Name, toolCall.Function.Arguments, toolCall.ID, tn.unknownToolHandler) + } else { + toolCallTasks[i].meta = tuple.meta[index] + toolCallTasks[i].name = toolCall.Function.Name + toolCallTasks[i].callID = toolCall.ID + + if tuple.enhancedInvokableEndpoints[index] != nil && tuple.enhancedStreamableEndpoints[index] != nil { + toolCallTasks[i].enhancedInvokableEndpoint = tuple.enhancedInvokableEndpoints[index] + toolCallTasks[i].enhancedStreamableEndpoint = tuple.enhancedStreamableEndpoints[index] + toolCallTasks[i].useEnhanced = true + } else { + toolCallTasks[i].endpoint = tuple.endpoints[index] + toolCallTasks[i].streamEndpoint = tuple.streamEndpoints[index] + toolCallTasks[i].useEnhanced = false + } + + // Get canonical tool name for looking up argument aliases + canonicalToolName := tuple.canonicalNames[index] + + // Process argument aliases remapping + args := toolCall.Function.Arguments + if aliasMap, hasAliases := tuple.argsAliasMap[canonicalToolName]; hasAliases { + remappedArgs, err := remapArgs(args, aliasMap) + if err != nil { + return nil, fmt.Errorf("failed to remap args for tool[name:%s]: %w", canonicalToolName, err) + } + args = remappedArgs + } + + if tn.toolArgumentsHandler != nil { + arg, err := tn.toolArgumentsHandler(ctx, canonicalToolName, args) + if err != nil { + return nil, fmt.Errorf("failed to executed tool[name:%s arguments:%s] arguments handler: %w", toolCall.Function.Name, args, err) + } + toolCallTasks[i].arg = arg + } else { + toolCallTasks[i].arg = args + } + } + } + + return toolCallTasks, nil +} + +func newUnknownToolTask(name, arg, callID string, unknownToolHandler func(ctx context.Context, name, input string) (string, error)) toolCallTask { + endpoint := func(ctx context.Context, input *ToolInput) (*ToolOutput, error) { + result, err := unknownToolHandler(ctx, input.Name, input.Arguments) + if err != nil { + return nil, err + } + return &ToolOutput{ + Result: result, + }, nil + } + return toolCallTask{ + endpoint: endpoint, + streamEndpoint: invokableToStreamable(endpoint), + meta: &executorMeta{ + component: components.ComponentOfTool, + isComponentCallbackEnabled: false, + componentImplType: "UnknownTool", + }, + name: name, + arg: arg, + callID: callID, + } +} + +func runToolCallTaskByInvoke(ctx context.Context, task *toolCallTask, opts ...tool.Option) { + if task.executed { + return + } + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{ + Name: task.name, + Type: task.meta.componentImplType, + Component: task.meta.component, + }) + + ctx = setToolCallInfo(ctx, &toolCallInfo{toolCallID: task.callID}) + ctx = appendToolAddressSegment(ctx, task.name, task.callID) + + if task.useEnhanced { + enhancedOutput, err := task.enhancedInvokableEndpoint(ctx, &ToolInput{ + Name: task.name, + Arguments: task.arg, + CallID: task.callID, + CallOptions: opts, + }) + if err != nil { + task.err = err + } else { + task.enhancedOutput = enhancedOutput.Result + task.executed = true + } + } else { + output, err := task.endpoint(ctx, &ToolInput{ + Name: task.name, + Arguments: task.arg, + CallID: task.callID, + CallOptions: opts, + }) + if err != nil { + task.err = err + } else { + task.output = output.Result + task.executed = true + } + } +} + +func runToolCallTaskByStream(ctx context.Context, task *toolCallTask, opts ...tool.Option) { + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{ + Name: task.name, + Type: task.meta.componentImplType, + Component: task.meta.component, + }) + + ctx = setToolCallInfo(ctx, &toolCallInfo{toolCallID: task.callID}) + ctx = appendToolAddressSegment(ctx, task.name, task.callID) + + if task.useEnhanced { + enhancedOutput, err := task.enhancedStreamableEndpoint(ctx, &ToolInput{ + Name: task.name, + Arguments: task.arg, + CallID: task.callID, + CallOptions: opts, + }) + if err != nil { + task.err = err + } else { + task.enhancedSOutput = enhancedOutput.Result + task.executed = true + } + } else { + output, err := task.streamEndpoint(ctx, &ToolInput{ + Name: task.name, + Arguments: task.arg, + CallID: task.callID, + CallOptions: opts, + }) + if err != nil { + task.err = err + } else { + task.sOutput = output.Result + task.executed = true + } + } +} + +func sequentialRunToolCall(ctx context.Context, + run func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option), + tasks []toolCallTask, opts ...tool.Option) { + + for i := range tasks { + if tasks[i].executed { + continue + } + run(ctx, &tasks[i], opts...) + } +} + +func parallelRunToolCall(ctx context.Context, + run func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option), + tasks []toolCallTask, opts ...tool.Option) { + + if len(tasks) == 1 { + run(ctx, &tasks[0], opts...) + return + } + + var wg sync.WaitGroup + for i := 1; i < len(tasks); i++ { + if tasks[i].executed { + continue + } + wg.Add(1) + go func(ctx_ context.Context, t *toolCallTask, opts ...tool.Option) { + defer wg.Done() + defer func() { + panicErr := recover() + if panicErr != nil { + t.err = safe.NewPanicErr(panicErr, debug.Stack()) + } + }() + run(ctx_, t, opts...) + }(ctx, &tasks[i], opts...) + } + + if !tasks[0].executed { + run(ctx, &tasks[0], opts...) + } + + wg.Wait() +} + +// buildTupleFromOpts rebuilds a toolsTuple when call options override tools or aliases. +func (tn *ToolsNode) buildTupleFromOpts(ctx context.Context, opt *toolsNodeOptions) (*toolsTuple, error) { + tools := opt.ToolList + if tools == nil { + tools = tn.tools + } + aliasConfigs := opt.ToolAliases + if aliasConfigs == nil { + aliasConfigs = tn.toolAliasConfigs + } + p := convToolsParams{ + tools: tools, + aliasConfigs: aliasConfigs, + } + p.middlewares.invokable = tn.toolCallMiddlewares + p.middlewares.streamable = tn.streamToolCallMiddlewares + p.middlewares.enhancedInvokable = tn.enhancedToolCallMiddlewares + p.middlewares.enhancedStreamable = tn.enhancedStreamToolCallMiddlewares + tuple, err := convTools(ctx, p) + if err != nil { + return nil, fmt.Errorf("failed to convert tool list from call option: %w", err) + } + return tuple, nil +} + +// Invoke calls the tools and collects the results of invokable tools. +// it's parallel if there are multiple tool calls in the input message. +func (tn *ToolsNode) Invoke(ctx context.Context, input *schema.Message, + opts ...ToolsNodeOption) ([]*schema.Message, error) { + + opt := getToolsNodeOptions(opts...) + tuple := tn.tuple + if opt.ToolList != nil || opt.ToolAliases != nil { + var err error + tuple, err = tn.buildTupleFromOpts(ctx, opt) + if err != nil { + return nil, err + } + } + + var executedTools map[string]string + var executedEnhancedTools map[string]*schema.ToolResult + if wasInterrupted, hasState, tnState := GetInterruptState[*toolsInterruptAndRerunState](ctx); wasInterrupted && hasState { + input = tnState.Input + if tnState.ExecutedTools != nil { + executedTools = tnState.ExecutedTools + } + if tnState.ExecutedEnhancedTools != nil { + executedEnhancedTools = tnState.ExecutedEnhancedTools + } + } + + tasks, err := tn.genToolCallTasks(ctx, tuple, input, executedTools, executedEnhancedTools, false) + if err != nil { + return nil, err + } + + if tn.executeSequentially { + sequentialRunToolCall(ctx, runToolCallTaskByInvoke, tasks, opt.ToolOptions...) + } else { + parallelRunToolCall(ctx, runToolCallTaskByInvoke, tasks, opt.ToolOptions...) + } + + n := len(tasks) + output := make([]*schema.Message, n) + + rerunExtra := &ToolsInterruptAndRerunExtra{ + ToolCalls: input.ToolCalls, + ExecutedTools: make(map[string]string), + ExecutedEnhancedTools: make(map[string]*schema.ToolResult), + RerunExtraMap: make(map[string]any), + } + rerunState := &toolsInterruptAndRerunState{ + Input: input, + ExecutedTools: make(map[string]string), + ExecutedEnhancedTools: make(map[string]*schema.ToolResult), + } + + var errs []error + for i := 0; i < n; i++ { + if tasks[i].err != nil { + info, ok := IsInterruptRerunError(tasks[i].err) + if !ok { + return nil, fmt.Errorf("failed to invoke tool[name:%s id:%s]: %w", tasks[i].name, tasks[i].callID, tasks[i].err) + } + + rerunExtra.RerunTools = append(rerunExtra.RerunTools, tasks[i].callID) + rerunState.RerunTools = append(rerunState.RerunTools, tasks[i].callID) + if info != nil { + rerunExtra.RerunExtraMap[tasks[i].callID] = info + } + + iErr := WrapInterruptAndRerunIfNeeded(ctx, + AddressSegment{ID: tasks[i].callID, Type: AddressSegmentTool}, tasks[i].err) + errs = append(errs, iErr) + continue + } + if tasks[i].executed { + if tasks[i].useEnhanced { + rerunExtra.ExecutedEnhancedTools[tasks[i].callID] = tasks[i].enhancedOutput + rerunState.ExecutedEnhancedTools[tasks[i].callID] = tasks[i].enhancedOutput + } else { + rerunExtra.ExecutedTools[tasks[i].callID] = tasks[i].output + rerunState.ExecutedTools[tasks[i].callID] = tasks[i].output + } + } + + if len(errs) == 0 { + if tasks[i].useEnhanced { + output[i] = schema.ToolMessage("", tasks[i].callID, schema.WithToolName(tasks[i].name)) + output[i].UserInputMultiContent, err = tasks[i].enhancedOutput.ToMessageInputParts() + if err != nil { + return nil, err + } + } else { + output[i] = schema.ToolMessage(tasks[i].output, tasks[i].callID, schema.WithToolName(tasks[i].name)) + } + + } + } + if len(errs) > 0 { + return nil, CompositeInterrupt(ctx, rerunExtra, rerunState, errs...) + } + + return output, nil +} + +// Stream calls the tools and collects the results of stream readers. +// it's parallel if there are multiple tool calls in the input message. +func (tn *ToolsNode) Stream(ctx context.Context, input *schema.Message, + opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error) { + + opt := getToolsNodeOptions(opts...) + tuple := tn.tuple + if opt.ToolList != nil || opt.ToolAliases != nil { + var err error + tuple, err = tn.buildTupleFromOpts(ctx, opt) + if err != nil { + return nil, err + } + } + + var executedTools map[string]string + var executedEnhancedTools map[string]*schema.ToolResult + if wasInterrupted, hasState, tnState := GetInterruptState[*toolsInterruptAndRerunState](ctx); wasInterrupted && hasState { + input = tnState.Input + if tnState.ExecutedTools != nil { + executedTools = tnState.ExecutedTools + } + if tnState.ExecutedEnhancedTools != nil { + executedEnhancedTools = tnState.ExecutedEnhancedTools + } + } + + tasks, err := tn.genToolCallTasks(ctx, tuple, input, executedTools, executedEnhancedTools, true) + if err != nil { + return nil, err + } + + if tn.executeSequentially { + sequentialRunToolCall(ctx, runToolCallTaskByStream, tasks, opt.ToolOptions...) + } else { + parallelRunToolCall(ctx, runToolCallTaskByStream, tasks, opt.ToolOptions...) + } + + n := len(tasks) + + rerunExtra := &ToolsInterruptAndRerunExtra{ + ToolCalls: input.ToolCalls, + ExecutedTools: make(map[string]string), + ExecutedEnhancedTools: make(map[string]*schema.ToolResult), + RerunExtraMap: make(map[string]any), + } + rerunState := &toolsInterruptAndRerunState{ + Input: input, + ExecutedTools: make(map[string]string), + ExecutedEnhancedTools: make(map[string]*schema.ToolResult), + } + var errs []error + // check rerun + for i := 0; i < n; i++ { + if tasks[i].err != nil { + info, ok := IsInterruptRerunError(tasks[i].err) + if !ok { + return nil, fmt.Errorf("failed to stream tool call %s: %w", tasks[i].callID, tasks[i].err) + } + + rerunExtra.RerunTools = append(rerunExtra.RerunTools, tasks[i].callID) + rerunState.RerunTools = append(rerunState.RerunTools, tasks[i].callID) + if info != nil { + rerunExtra.RerunExtraMap[tasks[i].callID] = info + } + iErr := WrapInterruptAndRerunIfNeeded(ctx, + AddressSegment{ID: tasks[i].callID, Type: AddressSegmentTool}, tasks[i].err) + errs = append(errs, iErr) + continue + } + } + + if len(errs) > 0 { + // concat and save tool output + for _, t := range tasks { + if t.executed { + if t.useEnhanced { + eo, err_ := concatStreamReader(t.enhancedSOutput) + if err_ != nil { + return nil, fmt.Errorf("failed to concat enhanced tool[name:%s id:%s]'s stream output: %w", t.name, t.callID, err_) + } + rerunExtra.ExecutedEnhancedTools[t.callID] = eo + rerunState.ExecutedEnhancedTools[t.callID] = eo + + } else { + o, err_ := concatStreamReader(t.sOutput) + if err_ != nil { + return nil, fmt.Errorf("failed to concat tool[name:%s id:%s]'s stream output: %w", t.name, t.callID, err_) + } + rerunExtra.ExecutedTools[t.callID] = o + rerunState.ExecutedTools[t.callID] = o + } + } + } + return nil, CompositeInterrupt(ctx, rerunExtra, rerunState, errs...) + } + + // common return + sOutput := make([]*schema.StreamReader[[]*schema.Message], n) + for i := 0; i < n; i++ { + index := i + callID := tasks[i].callID + callName := tasks[i].name + if tasks[i].useEnhanced { + cvt := func(tr *schema.ToolResult) ([]*schema.Message, error) { + ret := make([]*schema.Message, n) + ret[index] = schema.ToolMessage("", callID, schema.WithToolName(callName)) + ret[index].UserInputMultiContent, err = tr.ToMessageInputParts() + if err != nil { + return nil, err + } + return ret, nil + } + sOutput[i] = schema.StreamReaderWithConvert(tasks[i].enhancedSOutput, cvt) + } else { + cvt := func(s string) ([]*schema.Message, error) { + ret := make([]*schema.Message, n) + ret[index] = schema.ToolMessage(s, callID, schema.WithToolName(callName)) + return ret, nil + } + sOutput[i] = schema.StreamReaderWithConvert(tasks[i].sOutput, cvt) + } + + } + return schema.MergeStreamReaders(sOutput), nil +} + +// GetType returns the component type string for the Tools node. +func (tn *ToolsNode) GetType() string { + return "" +} + +func getToolsNodeOptions(opts ...ToolsNodeOption) *toolsNodeOptions { + o := &toolsNodeOptions{ + ToolOptions: make([]tool.Option, 0), + } + for _, opt := range opts { + opt(o) + } + return o +} + +type toolCallInfoKey struct{} +type toolCallInfo struct { + toolCallID string +} + +func setToolCallInfo(ctx context.Context, toolCallInfo *toolCallInfo) context.Context { + return context.WithValue(ctx, toolCallInfoKey{}, toolCallInfo) +} + +// GetToolCallID gets the current tool call id from the context. +func GetToolCallID(ctx context.Context) string { + v := ctx.Value(toolCallInfoKey{}) + if v == nil { + return "" + } + + info, ok := v.(*toolCallInfo) + if !ok { + return "" + } + + return info.toolCallID +} diff --git a/compose/tool_node_test.go b/compose/tool_node_test.go new file mode 100644 index 0000000..e24b42b --- /dev/null +++ b/compose/tool_node_test.go @@ -0,0 +1,1310 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "io" + "strconv" + "strings" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +const ( + toolNameOfUserCompany = "user_company" + toolIDOfUserCompany = "call_TRZhlagwBS0LpWbWPeZOvIXc" + + toolNameOfUserSalary = "user_salary" + toolIDOfUserSalary = "call_AqfoRW6fuF98k0o7696k2nzm" +) + +func TestToolsNode(t *testing.T) { + var err error + ctx := context.Background() + + userCompanyToolInfo := &schema.ToolInfo{ + Name: toolNameOfUserCompany, + Desc: "Query user's company and position information based on user's name and email", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Type: "string", + Desc: "User's name", + }, + "email": { + Type: "string", + Desc: "User's email", + }, + }), + } + + userSalaryToolInfo := &schema.ToolInfo{ + Name: toolNameOfUserSalary, + Desc: "Query user's salary information based on user's name and email", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Type: "string", + Desc: "User's name", + }, + "email": { + Type: "string", + Desc: "User's email", + }, + }), + } + + t.Run("success", func(t *testing.T) { + const ( + nodeOfTools = "tools" + nodeOfModel = "model" + ) + g := NewGraph[[]*schema.Message, []*schema.Message]() + + err = g.AddChatModelNode(nodeOfModel, &mockIntentChatModel{}) + assert.NoError(t, err) + + ui := newTool(userCompanyToolInfo, queryUserCompany) + us := newStreamableTool(userSalaryToolInfo, queryUserSalary) + + toolsNode, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{ui, us}, + }) + assert.NoError(t, err) + + err = g.AddToolsNode(nodeOfTools, toolsNode) + assert.NoError(t, err) + + err = g.AddEdge(START, nodeOfModel) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfModel, nodeOfTools) + assert.NoError(t, err) + + err = g.AddEdge(nodeOfTools, END) + assert.NoError(t, err) + + r, err := g.Compile(ctx) + assert.NoError(t, err) + + out, err := r.Invoke(ctx, []*schema.Message{}) + assert.NoError(t, err) + + msg := findMsgByToolCallID(out, toolIDOfUserCompany) + assert.Equal(t, toolIDOfUserCompany, msg.ToolCallID) + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","gender":"male","company":"bytedance","position":"CEO"}`, + msg.Content) + + msg = findMsgByToolCallID(out, toolIDOfUserSalary) + assert.Equal(t, toolIDOfUserSalary, msg.ToolCallID) + assert.Contains(t, msg.Content, + `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":5000}{"user_id":"zhangsan-zhangsan@bytedance.com","salary":3000}{"user_id":"zhangsan-zhangsan@bytedance.com","salary":2000}`) + + // 测试流式调用 + reader, err := r.Stream(ctx, []*schema.Message{}) + assert.NoError(t, err) + loops := 0 + userSalaryTimes := 0 + + defer reader.Close() + + var arrMsgs [][]*schema.Message + for ; loops < 10; loops++ { + msgs, err := reader.Recv() + if err == io.EOF { + break + } + + arrMsgs = append(arrMsgs, msgs) + + assert.NoError(t, err) + + assert.Len(t, msgs, 2) + if msg := findMsgByToolCallID(out, toolIDOfUserCompany); msg != nil { + assert.Equal(t, schema.Tool, msg.Role) + assert.Equal(t, toolIDOfUserCompany, msg.ToolCallID) + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","gender":"male","company":"bytedance","position":"CEO"}`, + msg.Content) + } else if msg := findMsgByToolCallID(out, toolIDOfUserSalary); msg != nil { + assert.Equal(t, schema.Tool, msg.Role) + assert.Equal(t, toolIDOfUserSalary, msg.ToolCallID) + + switch userSalaryTimes { + case 0: + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":5000}`, + msg.Content) + case 1: + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":3000}`, + msg.Content) + case 2: + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":2000}`, + msg.Content) + } + + userSalaryTimes++ + } else { + assert.Fail(t, "unexpected tool name") + } + } + + assert.Equal(t, 4, loops) + + msgs, err_ := schema.ConcatMessageArray(arrMsgs) + assert.NoError(t, err_) + msg = findMsgByToolCallID(msgs, toolIDOfUserCompany) + msg = findMsgByToolCallID(msgs, toolIDOfUserSalary) + + sr, sw := schema.Pipe[[]*schema.Message](2) + sw.Send([]*schema.Message{ + { + Role: schema.User, + Content: `hi, how are you`, + }, + }, nil) + sw.Send([]*schema.Message{ + { + Role: schema.User, + Content: `i'm fine'`, + }, + }, nil) + sw.Close() + + reader, err = r.Transform(ctx, sr) + assert.NoError(t, err) + + defer reader.Close() + + loops = 0 + userSalaryTimes = 0 + + for ; loops < 10; loops++ { + msgs, err := reader.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + + assert.Len(t, msgs, 2) + if msg := findMsgByToolCallID(out, toolIDOfUserCompany); msg != nil { + assert.Equal(t, schema.Tool, msg.Role) + assert.Equal(t, toolIDOfUserCompany, msg.ToolCallID) + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","gender":"male","company":"bytedance","position":"CEO"}`, + msg.Content) + } else if msg := findMsgByToolCallID(out, toolIDOfUserSalary); msg != nil { + assert.Equal(t, schema.Tool, msg.Role) + assert.Equal(t, toolIDOfUserSalary, msg.ToolCallID) + + switch userSalaryTimes { + case 0: + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":5000}`, + msg.Content) + case 1: + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":3000}`, + msg.Content) + case 2: + assert.JSONEq(t, `{"user_id":"zhangsan-zhangsan@bytedance.com","salary":2000}`, + msg.Content) + } + + userSalaryTimes++ + } else { + assert.Fail(t, "unexpected tool name") + } + } + + assert.Equal(t, 4, loops) + }) + + t.Run("order_consistency", func(t *testing.T) { + // Create a ToolsNode with multiple tools + ui := newTool(userCompanyToolInfo, queryUserCompany) + us := newTool(userSalaryToolInfo, queryUserSalary) + + toolsNode, err_ := NewToolNode(context.Background(), &ToolsNodeConfig{ + Tools: []tool.BaseTool{ui, us}, + }) + assert.NoError(t, err_) + + // Create an input message with multiple tool calls in a specific order + input := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserSalary, + Function: schema.FunctionCall{ + Name: toolNameOfUserSalary, + Arguments: `{"name": "zhangsan", "email": "zhangsan@bytedance.com"}`, + }, + }, + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: toolNameOfUserCompany, + Arguments: `{"name": "zhangsan", "email": "zhangsan@bytedance.com"}`, + }, + }, + }, + } + + // Invoke the ToolsNode + output, err_ := toolsNode.Invoke(context.Background(), input) + assert.NoError(t, err_) + + // Verify the order of output messages matches the order of input tool calls + assert.Equal(t, 2, len(output)) + assert.Equal(t, toolIDOfUserSalary, output[0].ToolCallID) + assert.Equal(t, toolIDOfUserCompany, output[1].ToolCallID) + + // Test with Stream method as well + streamer, err_ := toolsNode.Stream(context.Background(), input) + assert.NoError(t, err_) + defer streamer.Close() + + // Collect all stream outputs + var streamOutputs [][]*schema.Message + for { + chunk, err__ := streamer.Recv() + if err__ == io.EOF { + break + } + assert.NoError(t, err__) + streamOutputs = append(streamOutputs, chunk) + } + + // Verify each chunk maintains the correct order + for _, chunk := range streamOutputs { + if chunk[0] != nil { + assert.Equal(t, toolIDOfUserSalary, chunk[0].ToolCallID) + } + if chunk[1] != nil { + assert.Equal(t, toolIDOfUserCompany, chunk[1].ToolCallID) + } + } + + // Concatenate all stream outputs and verify final result + concatenated, err_ := schema.ConcatMessageArray(streamOutputs) + assert.NoError(t, err_) + assert.Equal(t, 2, len(concatenated)) + assert.Equal(t, toolIDOfUserSalary, concatenated[0].ToolCallID) + assert.Equal(t, toolIDOfUserCompany, concatenated[1].ToolCallID) + }) +} + +type userCompanyRequest struct { + Name string `json:"name"` + Email string `json:"email"` +} + +type userCompanyResponse struct { + UserID string `json:"user_id"` + Gender string `json:"gender"` + Company string `json:"company"` + Position string `json:"position"` +} + +func queryUserCompany(ctx context.Context, req *userCompanyRequest) (resp *userCompanyResponse, err error) { + callID := GetToolCallID(ctx) + if callID != toolIDOfUserCompany { + return nil, fmt.Errorf("invalid tool call id= %s", callID) + } + + return &userCompanyResponse{ + UserID: fmt.Sprintf("%v-%v", req.Name, req.Email), + Gender: "male", + Company: "bytedance", + Position: "CEO", + }, nil +} + +type userSalaryRequest struct { + Name string `json:"name"` + Email string `json:"email"` +} + +type userSalaryResponse struct { + UserID string `json:"user_id"` + Salary int `json:"salary"` +} + +func queryUserSalary(ctx context.Context, req *userSalaryRequest) (resp *schema.StreamReader[*userSalaryResponse], err error) { + callID := GetToolCallID(ctx) + if callID != toolIDOfUserSalary { + return nil, fmt.Errorf("invalid tool call id= %s", callID) + } + + sr, sw := schema.Pipe[*userSalaryResponse](10) + sw.Send(&userSalaryResponse{ + UserID: fmt.Sprintf("%v-%v", req.Name, req.Email), + Salary: 5000, + }, nil) + + sw.Send(&userSalaryResponse{ + UserID: fmt.Sprintf("%v-%v", req.Name, req.Email), + Salary: 3000, + }, nil) + + sw.Send(&userSalaryResponse{ + UserID: fmt.Sprintf("%v-%v", req.Name, req.Email), + Salary: 2000, + }, nil) + sw.Close() + return sr, nil +} + +type mockIntentChatModel struct{} + +func (m *mockIntentChatModel) BindTools(tools []*schema.ToolInfo) error { + return nil +} + +func (m *mockIntentChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return &schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: toolNameOfUserCompany, + Arguments: `{"name": "zhangsan", "email": "zhangsan@bytedance.com"}`, + }, + }, + { + ID: toolIDOfUserSalary, + Function: schema.FunctionCall{ + Name: toolNameOfUserSalary, + Arguments: `{"name": "zhangsan", "email": "zhangsan@bytedance.com"}`, + }, + }, + }, + }, nil +} + +func (m *mockIntentChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](2) + sw.Send(&schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: toolNameOfUserCompany, + Arguments: `{"name": "zhangsan", "email": "zhangsan@bytedance.com"}`, + }, + }, + }, + }, nil) + + sw.Send(&schema.Message{ + Role: schema.Assistant, + Content: "", + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserSalary, + Function: schema.FunctionCall{ + Name: toolNameOfUserSalary, + Arguments: `{"name": "zhangsan", "email": "zhangsan@bytedance.com"}`, + }, + }, + }, + }, nil) + + sw.Close() + + return sr, nil +} + +func TestToolsNodeOptions(t *testing.T) { + ctx := context.Background() + + t.Run("tool_option", func(t *testing.T) { + + g := NewGraph[*schema.Message, []*schema.Message]() + + mt := &mockTool{} + + tn, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{mt}, + }) + assert.NoError(t, err) + + err = g.AddToolsNode("tools", tn) + assert.NoError(t, err) + + err = g.AddEdge(START, "tools") + assert.NoError(t, err) + err = g.AddEdge("tools", END) + assert.NoError(t, err) + + r, err := g.Compile(ctx) + assert.NoError(t, err) + + out, err := r.Invoke(ctx, &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: "mock_tool", + Arguments: `{"name": "jack"}`, + }, + }, + }, + }, WithToolsNodeOption(WithToolOption(WithAge(10)))) + assert.NoError(t, err) + assert.Len(t, out, 1) + assert.JSONEq(t, `{"echo": "jack: 10"}`, out[0].Content) + + outMessages := make([][]*schema.Message, 0) + outStream, err := r.Stream(ctx, &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: "mock_tool", + Arguments: `{"name": "jack"}`, + }, + }, + }, + }, WithToolsNodeOption(WithToolOption(WithAge(10)))) + + assert.NoError(t, err) + + for { + msgs, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + outMessages = append(outMessages, msgs) + } + outStream.Close() + + msgs, err := internal.ConcatItems(outMessages) + assert.NoError(t, err) + + assert.Len(t, msgs, 1) + assert.JSONEq(t, `{"echo":"jack: 10"}`, msgs[0].Content) + }) + t.Run("tool_list", func(t *testing.T) { + + g := NewGraph[*schema.Message, []*schema.Message]() + + mt := &mockTool{} + + tn, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{}, + }) + assert.NoError(t, err) + + err = g.AddToolsNode("tools", tn) + assert.NoError(t, err) + + err = g.AddEdge(START, "tools") + assert.NoError(t, err) + err = g.AddEdge("tools", END) + assert.NoError(t, err) + + r, err := g.Compile(ctx) + assert.NoError(t, err) + + out, err := r.Invoke(ctx, &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: "mock_tool", + Arguments: `{"name": "jack"}`, + }, + }, + }, + }, WithToolsNodeOption(WithToolList(mt), WithToolOption(WithAge(10)))) + assert.NoError(t, err) + assert.Len(t, out, 1) + assert.JSONEq(t, `{"echo": "jack: 10"}`, out[0].Content) + + outMessages := make([][]*schema.Message, 0) + outStream, err := r.Stream(ctx, &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: toolIDOfUserCompany, + Function: schema.FunctionCall{ + Name: "mock_tool", + Arguments: `{"name": "jack"}`, + }, + }, + }, + }, WithToolsNodeOption(WithToolList(mt), WithToolOption(WithAge(10)))) + + assert.NoError(t, err) + + for { + msgs, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + outMessages = append(outMessages, msgs) + } + outStream.Close() + + msgs, err := internal.ConcatItems(outMessages) + assert.NoError(t, err) + + assert.Len(t, msgs, 1) + assert.JSONEq(t, `{"echo":"jack: 10"}`, msgs[0].Content) + }) + +} + +func findMsgByToolCallID(msgs []*schema.Message, toolCallID string) *schema.Message { + for _, msg := range msgs { + if msg.ToolCallID == toolCallID { + return msg + } + } + + return nil +} + +type mockToolOptions struct { + Age int +} + +func WithAge(age int) tool.Option { + return tool.WrapImplSpecificOptFn(func(o *mockToolOptions) { + o.Age = age + }) +} + +type mockToolRequest struct { + Name string `json:"name"` +} + +type mockToolResponse struct { + Echo string `json:"echo"` +} + +type mockTool struct{} + +func (m *mockTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "mock_tool", + Desc: "mock tool", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Type: "string", + Desc: "name", + Required: true, + }, + }), + }, nil +} + +func (m *mockTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + opt := tool.GetImplSpecificOptions(&mockToolOptions{}, opts...) + + req := &mockToolRequest{} + + if e := sonic.UnmarshalString(argumentsInJSON, req); e != nil { + return "", e + } + + resp := &mockToolResponse{ + Echo: fmt.Sprintf("%v: %v", req.Name, opt.Age), + } + + return sonic.MarshalString(resp) +} + +func (m *mockTool) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + sr, sw := schema.Pipe[string](1) + go func() { + defer sw.Close() + + opt := tool.GetImplSpecificOptions(&mockToolOptions{}, opts...) + + req := &mockToolRequest{} + + if e := sonic.UnmarshalString(argumentsInJSON, req); e != nil { + sw.Send("", e) + return + } + + resp := mockToolResponse{ + Echo: fmt.Sprintf("%v: %v", req.Name, opt.Age), + } + + output, err := sonic.MarshalString(resp) + if err != nil { + sw.Send("", err) + return + } + + for i := 0; i < len(output); i++ { + sw.Send(string(output[i]), nil) + } + }() + + return sr, nil +} + +func TestUnknownTool(t *testing.T) { + ctx := context.Background() + tn, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: nil, + UnknownToolsHandler: func(ctx context.Context, name, input string) (string, error) { + return "unknown", nil + }, + }) + assert.NoError(t, err) + + input := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + ID: "1", + Function: schema.FunctionCall{ + Name: "unknown1", + Arguments: `arg1`, + }, + }, + { + ID: "2", + Function: schema.FunctionCall{ + Name: "unknown2", + Arguments: `arg2`, + }, + }, + }, + } + + expected := []*schema.Message{ + { + Role: schema.Tool, + Content: "unknown", + ToolCallID: "1", + ToolName: "unknown1", + }, + { + Role: schema.Tool, + Content: "unknown", + ToolCallID: "2", + ToolName: "unknown2", + }, + } + + result, err := tn.Invoke(ctx, input) + assert.NoError(t, err) + assert.Equal(t, expected, result) + + streamResult, err := tn.Stream(ctx, input) + assert.NoError(t, err) + result = make([]*schema.Message, 2) + for { + chunk, err := streamResult.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + for i := range chunk { + if chunk[i] != nil { + result[i] = chunk[i] + } + } + } + assert.Equal(t, expected, result) +} + +func TestToolRerun(t *testing.T) { + type myToolRerunState struct { + In *schema.Message + } + + schema.Register[myToolRerunState]() + + tc := []schema.ToolCall{ + { + ID: "3", + Function: schema.FunctionCall{ + Name: "tool3", + Arguments: "input", + }, + }, + { + ID: "4", + Function: schema.FunctionCall{ + Name: "tool4", + Arguments: "input", + }, + }, + { + ID: "1", + Function: schema.FunctionCall{ + Name: "tool1", + Arguments: "input", + }, + }, + { + ID: "2", + Function: schema.FunctionCall{ + Name: "tool2", + Arguments: "input", + }, + }, + } + g := NewGraph[*schema.Message, string](WithGenLocalState(func(ctx context.Context) (state *myToolRerunState) { + return &myToolRerunState{In: &schema.Message{Role: schema.Assistant, ToolCalls: tc}} + })) + ctx := context.Background() + tn, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{&myTool1{}, &myTool2{}, &myTool3{t: t}, &myTool4{t: t}}, + }) + assert.NoError(t, err) + assert.NoError(t, g.AddToolsNode("tool node", tn)) + assert.NoError(t, g.AddLambdaNode("lambda", InvokableLambda(func(ctx context.Context, input []*schema.Message) (output string, err error) { + contents := make([]string, len(input)) + for _, m := range input { + callID := m.ToolCallID + callIDInt, err := strconv.Atoi(callID) + if err != nil { + return "", err + } + contents[callIDInt-1] = m.Content + } + sb := strings.Builder{} + for _, m := range contents { + sb.WriteString(m) + } + return sb.String(), nil + }))) + assert.NoError(t, g.AddEdge(START, "tool node")) + assert.NoError(t, g.AddEdge("tool node", "lambda")) + assert.NoError(t, g.AddEdge("lambda", END)) + + r, err := g.Compile(ctx, WithCheckPointStore(&inMemoryStore{m: map[string][]byte{}})) + assert.NoError(t, err) + + _, err = r.Stream(ctx, &schema.Message{Role: schema.Assistant, ToolCalls: tc}, WithCheckPointID("1")) + info, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + assert.Equal(t, []string{"tool node"}, info.RerunNodes) + assert.Equal(t, &ToolsInterruptAndRerunExtra{ + ToolCalls: tc, + RerunTools: []string{"1", "2"}, + RerunExtraMap: map[string]any{"1": "tool1 rerun extra", "2": "tool2 rerun extra"}, + ExecutedTools: map[string]string{ + "3": "tool3 input: input", + "4": "tool4 input: input", + }, + ExecutedEnhancedTools: make(map[string]*schema.ToolResult), + }, info.RerunNodesExtra["tool node"]) + + sr, err := r.Stream(ctx, nil, WithCheckPointID("1")) + assert.NoError(t, err) + result, err := concatStreamReader(sr) + assert.NoError(t, err) + assert.Equal(t, "tool1 input: inputtool2 input: inputtool3 input: inputtool4 input: input", result) +} + +func TestToolMiddleware(t *testing.T) { + ctx := context.Background() + t3 := &myTool3{t: t} + t4 := &myTool4{t: t} + tn, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{t3, t4}, + ToolCallMiddlewares: []ToolMiddleware{ + { + Invokable: func(endpoint InvokableToolEndpoint) InvokableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*ToolOutput, error) { + _, err := endpoint(ctx, input) + if err != nil { + return nil, err + } + return &ToolOutput{Result: "middleware1"}, nil + } + }, + }, + { + Streamable: func(endpoint StreamableToolEndpoint) StreamableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*StreamToolOutput, error) { + _, err := endpoint(ctx, input) + if err != nil { + return nil, err + } + return &StreamToolOutput{Result: schema.StreamReaderFromArray([]string{"middleware2"})}, nil + } + }, + }, + }, + }) + assert.NoError(t, err) + + messages, err := tn.Invoke(ctx, schema.AssistantMessage("", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "tool3", Arguments: ""}}, + {ID: "2", Function: schema.FunctionCall{Name: "tool4", Arguments: ""}}, + })) + assert.NoError(t, err) + assert.Len(t, messages, 2) + assert.Equal(t, "middleware1", messages[0].Content) + assert.Equal(t, "middleware2", messages[1].Content) + + t3.times, t4.times = 0, 0 // reset t3 t4 + messageStreams, err := tn.Stream(ctx, schema.AssistantMessage("", []schema.ToolCall{ + {ID: "1", Function: schema.FunctionCall{Name: "tool3", Arguments: ""}}, + {ID: "2", Function: schema.FunctionCall{Name: "tool4", Arguments: ""}}, + })) + assert.NoError(t, err) + var messageArray [][]*schema.Message + for { + chunk, err := messageStreams.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + messageArray = append(messageArray, chunk) + } + messages, err = schema.ConcatMessageArray(messageArray) + assert.Len(t, messages, 2) + assert.Equal(t, "middleware1", messages[0].Content) + assert.Equal(t, "middleware2", messages[1].Content) +} + +type myTool1 struct { + times uint +} + +func (m *myTool1) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "tool1"}, nil +} + +func (m *myTool1) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + if m.times == 0 { + m.times++ + return "", tool.Interrupt(ctx, "tool1 rerun extra") + } + return "tool1 input: " + argumentsInJSON, nil +} + +type myTool2 struct { + times uint +} + +func (m *myTool2) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "tool2"}, nil +} + +func (m *myTool2) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + if m.times == 0 { + m.times++ + return nil, tool.Interrupt(ctx, "tool2 rerun extra") + } + return schema.StreamReaderFromArray([]string{"tool2 input: ", argumentsInJSON}), nil +} + +type myTool3 struct { + t *testing.T + times int +} + +func (m *myTool3) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "tool3"}, nil +} + +func (m *myTool3) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + assert.Equal(m.t, 0, m.times) + m.times++ + return "tool3 input: " + argumentsInJSON, nil +} + +type myTool4 struct { + t *testing.T + times int +} + +func (m *myTool4) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "tool4"}, nil +} + +func (m *myTool4) StreamableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (*schema.StreamReader[string], error) { + assert.Equal(m.t, 0, m.times) + m.times++ + return schema.StreamReaderFromArray([]string{"tool4 input: ", argumentsInJSON}), nil +} + +func newTool[I, O any](info *schema.ToolInfo, f func(ctx context.Context, in I) (O, error)) tool.InvokableTool { + return &invokableTool[I, O]{ + info: info, + fn: f, + } +} + +type invokableTool[I, O any] struct { + info *schema.ToolInfo + fn func(ctx context.Context, in I) (O, error) +} + +func (f *invokableTool[I, O]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return f.info, nil +} + +func (f *invokableTool[I, O]) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + t := generic.NewInstance[I]() + err := sonic.UnmarshalString(argumentsInJSON, t) + if err != nil { + return "", err + } + o, err := f.fn(ctx, t) + if err != nil { + return "", err + } + return sonic.MarshalString(o) +} + +func newStreamableTool[I, O any](info *schema.ToolInfo, f func(ctx context.Context, in I) (*schema.StreamReader[O], error)) tool.StreamableTool { + return &streamableTool[I, O]{ + info: info, + fn: f, + } +} + +type streamableTool[I, O any] struct { + info *schema.ToolInfo + fn func(ctx context.Context, in I) (*schema.StreamReader[O], error) +} + +func (f *streamableTool[I, O]) Info(ctx context.Context) (*schema.ToolInfo, error) { + return f.info, nil +} +func (f *streamableTool[I, O]) StreamableRun(ctx context.Context, argumentsInJSON string, _ ...tool.Option) (*schema.StreamReader[string], error) { + t := generic.NewInstance[I]() + err := sonic.UnmarshalString(argumentsInJSON, t) + if err != nil { + return nil, err + } + sr, err := f.fn(ctx, t) + if err != nil { + return nil, err + } + return schema.StreamReaderWithConvert(sr, func(o O) (string, error) { + return sonic.MarshalString(o) + }), nil +} + +type enhancedInvokableTool struct { + info *schema.ToolInfo + fn func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) +} + +func (e *enhancedInvokableTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return e.info, nil +} + +func (e *enhancedInvokableTool) InvokableRun(ctx context.Context, toolArgument *schema.ToolArgument, _ ...tool.Option) (*schema.ToolResult, error) { + return e.fn(ctx, toolArgument) +} + +type enhancedStreamableTool struct { + info *schema.ToolInfo + fn func(ctx context.Context, input *schema.ToolArgument) (*schema.StreamReader[*schema.ToolResult], error) +} + +func (e *enhancedStreamableTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return e.info, nil +} + +func (e *enhancedStreamableTool) StreamableRun(ctx context.Context, toolArgument *schema.ToolArgument, _ ...tool.Option) (*schema.StreamReader[*schema.ToolResult], error) { + return e.fn(ctx, toolArgument) +} + +func TestEnhancedToolNode(t *testing.T) { + ctx := context.Background() + + enhancedInvokable := &enhancedInvokableTool{ + info: &schema.ToolInfo{ + Name: "enhanced_invokable_tool", + Desc: "test enhanced invokable tool", + }, + fn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "invokable result: " + input.Text}, + }, + }, nil + }, + } + + enhancedStreamable := &enhancedStreamableTool{ + info: &schema.ToolInfo{ + Name: "enhanced_streamable_tool", + Desc: "test enhanced streamable tool", + }, + fn: func(ctx context.Context, input *schema.ToolArgument) (*schema.StreamReader[*schema.ToolResult], error) { + results := []*schema.ToolResult{ + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: "stream part 1: " + input.Text}}}, + {Parts: []schema.ToolOutputPart{{Type: schema.ToolPartTypeText, Text: " stream part 2"}}}, + } + return schema.StreamReaderFromArray(results), nil + }, + } + + toolNode, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{enhancedInvokable, enhancedStreamable}, + }) + assert.NoError(t, err) + assert.NotNil(t, toolNode) + + t.Run("enhanced invokable tool", func(t *testing.T) { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call1", + Function: schema.FunctionCall{ + Name: "enhanced_invokable_tool", + Arguments: "test input", + }, + }, + }) + + output, err := toolNode.Invoke(ctx, input) + assert.NoError(t, err) + assert.Len(t, output, 1) + assert.Equal(t, schema.Tool, output[0].Role) + assert.Equal(t, "call1", output[0].ToolCallID) + }) + + t.Run("enhanced streamable tool", func(t *testing.T) { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call2", + Function: schema.FunctionCall{ + Name: "enhanced_streamable_tool", + Arguments: "test stream", + }, + }, + }) + + streamReader, err := toolNode.Stream(ctx, input) + assert.NoError(t, err) + assert.NotNil(t, streamReader) + + var messages []*schema.Message + for { + chunk, err := streamReader.Recv() + if err != nil { + break + } + if chunk != nil { + messages = append(messages, chunk...) + } + } + message, err := schema.ConcatMessages(messages) + assert.NoError(t, err) + assert.Len(t, messages, 2) + assert.Equal(t, schema.Tool, messages[0].Role) + assert.Equal(t, "call2", messages[0].ToolCallID) + assert.Contains(t, message.UserInputMultiContent[0].Text, "stream part") + }) +} + +func TestEnhancedToolConversion(t *testing.T) { + ctx := context.Background() + + enhancedInvokable := &enhancedInvokableTool{ + info: &schema.ToolInfo{ + Name: "enhanced_only_invokable", + Desc: "test enhanced invokable only", + }, + fn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Type: schema.ToolPartTypeText, Text: "enhanced: " + input.Text}, + }, + }, nil + }, + } + + toolNode, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{enhancedInvokable}, + }) + assert.NoError(t, err) + + t.Run("enhanced invokable auto-converts to streamable", func(t *testing.T) { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call1", + Function: schema.FunctionCall{ + Name: "enhanced_only_invokable", + Arguments: "test", + }, + }, + }) + + streamReader, err := toolNode.Stream(ctx, input) + assert.NoError(t, err) + assert.NotNil(t, streamReader) + + var messages []*schema.Message + for { + chunk, err := streamReader.Recv() + if err != nil { + break + } + if chunk != nil { + messages = append(messages, chunk...) + } + } + assert.Len(t, messages, 1) + }) +} + +func TestEnhancedToolMiddleware(t *testing.T) { + ctx := context.Background() + + var invokableMiddlewareCalled bool + var streamableMiddlewareCalled bool + + enhancedInvokable := &enhancedInvokableTool{ + info: &schema.ToolInfo{ + Name: "enhanced_tool_with_middleware", + Desc: "test enhanced tool with middleware", + }, + fn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Text: "result", Type: schema.ToolPartTypeText}, + }, + }, nil + }, + } + + toolNode, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{enhancedInvokable}, + ToolCallMiddlewares: []ToolMiddleware{ + { + EnhancedInvokable: func(next EnhancedInvokableToolEndpoint) EnhancedInvokableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*EnhancedInvokableToolOutput, error) { + invokableMiddlewareCalled = true + return next(ctx, input) + } + }, + EnhancedStreamable: func(next EnhancedStreamableToolEndpoint) EnhancedStreamableToolEndpoint { + return func(ctx context.Context, input *ToolInput) (*EnhancedStreamableToolOutput, error) { + streamableMiddlewareCalled = true + return next(ctx, input) + } + }, + }, + }, + }) + assert.NoError(t, err) + + t.Run("enhanced invokable middleware", func(t *testing.T) { + invokableMiddlewareCalled = false + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call1", + Function: schema.FunctionCall{ + Name: "enhanced_tool_with_middleware", + Arguments: "test", + }, + }, + }) + + _, err := toolNode.Invoke(ctx, input) + assert.NoError(t, err) + assert.True(t, invokableMiddlewareCalled) + }) + + t.Run("enhanced streamable middleware", func(t *testing.T) { + streamableMiddlewareCalled = false + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call2", + Function: schema.FunctionCall{ + Name: "enhanced_tool_with_middleware", + Arguments: "test", + }, + }, + }) + + streamReader, err := toolNode.Stream(ctx, input) + assert.NoError(t, err) + for { + _, err := streamReader.Recv() + if err != nil { + break + } + } + assert.False(t, streamableMiddlewareCalled) + }) +} + +func TestEnhancedToolPriority(t *testing.T) { + ctx := context.Background() + + enhancedInvokable := &enhancedInvokableTool{ + info: &schema.ToolInfo{ + Name: "test_tool", + Desc: "test tool with both enhanced and regular", + }, + fn: func(ctx context.Context, input *schema.ToolArgument) (*schema.ToolResult, error) { + return &schema.ToolResult{ + Parts: []schema.ToolOutputPart{ + {Text: "enhanced result", Type: schema.ToolPartTypeText}, + }, + }, nil + }, + } + + toolNode, err := NewToolNode(ctx, &ToolsNodeConfig{ + Tools: []tool.BaseTool{enhancedInvokable}, + }) + assert.NoError(t, err) + + t.Run("enhanced tool is used when available", func(t *testing.T) { + input := schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "call1", + Function: schema.FunctionCall{ + Name: "test_tool", + Arguments: "test", + }, + }, + }) + + output, err := toolNode.Invoke(ctx, input) + assert.NoError(t, err) + assert.Len(t, output, 1) + assert.Contains(t, output[0].UserInputMultiContent[0].Text, "enhanced result") + }) +} diff --git a/compose/types.go b/compose/types.go new file mode 100644 index 0000000..54f8e2b --- /dev/null +++ b/compose/types.go @@ -0,0 +1,47 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "github.com/cloudwego/eino/components" +) + +type component = components.Component + +// built-in component types in graph node. +// it represents the type of the most primitive executable object provided by the user. +const ( + ComponentOfUnknown component = "Unknown" + ComponentOfGraph component = "Graph" + ComponentOfWorkflow component = "Workflow" + ComponentOfChain component = "Chain" + ComponentOfPassthrough component = "Passthrough" + ComponentOfToolsNode component = "ToolsNode" + ComponentOfAgenticToolsNode component = "AgenticToolsNode" + ComponentOfLambda component = "Lambda" +) + +// NodeTriggerMode controls the triggering mode of graph nodes. +type NodeTriggerMode string + +const ( + // AnyPredecessor means that the node will be triggered when any of its predecessors is included in the previous completed super step. + // Ref:https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/orchestration_design_principles/#runtime-engine + AnyPredecessor NodeTriggerMode = "any_predecessor" + // AllPredecessor means that the current node will only be triggered when all of its predecessor nodes have finished running. + AllPredecessor NodeTriggerMode = "all_predecessor" +) diff --git a/compose/types_composable.go b/compose/types_composable.go new file mode 100644 index 0000000..bfd3618 --- /dev/null +++ b/compose/types_composable.go @@ -0,0 +1,31 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "reflect" +) + +// AnyGraph the identifiers for composable and compilable Graph[I, O]、Chain[I, O] in Eino. +type AnyGraph interface { + getGenericHelper() *genericHelper + compile(ctx context.Context, options *graphCompileOptions) (*composableRunnable, error) + inputType() reflect.Type + outputType() reflect.Type + component() component +} diff --git a/compose/types_lambda.go b/compose/types_lambda.go new file mode 100644 index 0000000..05bbab3 --- /dev/null +++ b/compose/types_lambda.go @@ -0,0 +1,265 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/schema" +) + +// Invoke is the type of the invokable lambda function. +type Invoke[I, O, TOption any] func(ctx context.Context, input I, opts ...TOption) (output O, err error) + +// Stream is the type of the streamable lambda function. +type Stream[I, O, TOption any] func(ctx context.Context, + input I, opts ...TOption) (output *schema.StreamReader[O], err error) + +// Collect is the type of the collectable lambda function. +type Collect[I, O, TOption any] func(ctx context.Context, + input *schema.StreamReader[I], opts ...TOption) (output O, err error) + +// Transform is the type of the transformable lambda function. +type Transform[I, O, TOption any] func(ctx context.Context, + input *schema.StreamReader[I], opts ...TOption) (output *schema.StreamReader[O], err error) + +// InvokeWOOpt is the type of the invokable lambda function without options. +type InvokeWOOpt[I, O any] func(ctx context.Context, input I) (output O, err error) + +// StreamWOOpt is the type of the streamable lambda function without options. +type StreamWOOpt[I, O any] func(ctx context.Context, + input I) (output *schema.StreamReader[O], err error) + +// CollectWOOpt is the type of the collectable lambda function without options. +type CollectWOOpt[I, O any] func(ctx context.Context, + input *schema.StreamReader[I]) (output O, err error) + +// TransformWOOpts is the type of the transformable lambda function without options. +type TransformWOOpts[I, O any] func(ctx context.Context, + input *schema.StreamReader[I]) (output *schema.StreamReader[O], err error) + +// Lambda is the node that wraps the user provided lambda function. +// It can be used as a node in Graph or Chain (include Parallel and Branch). +// Create a Lambda by using AnyLambda/InvokableLambda/StreamableLambda/CollectableLambda/TransformableLambda. +// eg. +// +// lambda := compose.InvokableLambda(func(ctx context.Context, input string) (output string, err error) { +// return input, nil +// }) +type Lambda struct { + executor *composableRunnable +} + +type lambdaOpts struct { + // same as executorMeta.isComponentCallbackEnabled + // indicates whether the executable lambda user provided could execute the callback aspect itself. + // if it could, the callback in the corresponding graph node won't be executed anymore + enableComponentCallback bool + + // same as executorMeta.componentImplType + // for AnyLambda, the value comes from the user's explicit config + // if componentImplType is empty, then the class name or func name in the instance will be inferred, but no guarantee. + componentImplType string +} + +// LambdaOpt is the option for creating a Lambda. +type LambdaOpt func(o *lambdaOpts) + +// WithLambdaCallbackEnable enables the callback aspect of the lambda function. +func WithLambdaCallbackEnable(y bool) LambdaOpt { + return func(o *lambdaOpts) { + o.enableComponentCallback = y + } +} + +// WithLambdaType sets the type of the lambda function. +func WithLambdaType(t string) LambdaOpt { + return func(o *lambdaOpts) { + o.componentImplType = t + } +} + +type unreachableOption struct{} + +// InvokableLambdaWithOption creates a Lambda with invokable lambda function and options. +func InvokableLambdaWithOption[I, O, TOption any](i Invoke[I, O, TOption], opts ...LambdaOpt) *Lambda { + return anyLambda(i, nil, nil, nil, opts...) +} + +// InvokableLambda creates a Lambda with invokable lambda function without options. +func InvokableLambda[I, O any](i InvokeWOOpt[I, O], opts ...LambdaOpt) *Lambda { + f := func(ctx context.Context, input I, opts_ ...unreachableOption) (output O, err error) { + return i(ctx, input) + } + + return anyLambda(f, nil, nil, nil, opts...) +} + +// StreamableLambdaWithOption creates a Lambda with streamable lambda function and options. +func StreamableLambdaWithOption[I, O, TOption any](s Stream[I, O, TOption], opts ...LambdaOpt) *Lambda { + return anyLambda(nil, s, nil, nil, opts...) +} + +// StreamableLambda creates a Lambda with streamable lambda function without options. +func StreamableLambda[I, O any](s StreamWOOpt[I, O], opts ...LambdaOpt) *Lambda { + f := func(ctx context.Context, input I, opts_ ...unreachableOption) ( + output *schema.StreamReader[O], err error) { + + return s(ctx, input) + } + + return anyLambda(nil, f, nil, nil, opts...) +} + +// CollectableLambdaWithOption creates a Lambda with collectable lambda function and options. +func CollectableLambdaWithOption[I, O, TOption any](c Collect[I, O, TOption], opts ...LambdaOpt) *Lambda { + return anyLambda(nil, nil, c, nil, opts...) +} + +// CollectableLambda creates a Lambda with collectable lambda function without options. +func CollectableLambda[I, O any](c CollectWOOpt[I, O], opts ...LambdaOpt) *Lambda { + f := func(ctx context.Context, input *schema.StreamReader[I], + opts_ ...unreachableOption) (output O, err error) { + + return c(ctx, input) + } + + return anyLambda(nil, nil, f, nil, opts...) +} + +// TransformableLambdaWithOption creates a Lambda with transformable lambda function and options. +func TransformableLambdaWithOption[I, O, TOption any](t Transform[I, O, TOption], opts ...LambdaOpt) *Lambda { + return anyLambda(nil, nil, nil, t, opts...) +} + +// TransformableLambda creates a Lambda with transformable lambda function without options. +func TransformableLambda[I, O any](t TransformWOOpts[I, O], opts ...LambdaOpt) *Lambda { + + f := func(ctx context.Context, input *schema.StreamReader[I], + opts_ ...unreachableOption) (output *schema.StreamReader[O], err error) { + + return t(ctx, input) + } + + return anyLambda(nil, nil, nil, f, opts...) +} + +// AnyLambda creates a Lambda with any lambda function. +// you can only implement one or more of the four lambda functions, and the rest use nil. +// eg. +// +// invokeFunc := func(ctx context.Context, input string, opts ...myOption) (output string, err error) { +// // ... +// } +// streamFunc := func(ctx context.Context, input string, opts ...myOption) (output *schema.StreamReader[string], err error) { +// // ... +// } +// +// lambda := compose.AnyLambda(invokeFunc, streamFunc, nil, nil) +func AnyLambda[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption], + c Collect[I, O, TOption], t Transform[I, O, TOption], opts ...LambdaOpt) (*Lambda, error) { + + if i == nil && s == nil && c == nil && t == nil { + return nil, fmt.Errorf("needs to have at least one of four lambda types: invoke/stream/collect/transform, got none") + } + + return anyLambda(i, s, c, t, opts...), nil +} + +func anyLambda[I, O, TOption any](i Invoke[I, O, TOption], s Stream[I, O, TOption], + c Collect[I, O, TOption], t Transform[I, O, TOption], opts ...LambdaOpt) *Lambda { + + opt := getLambdaOpt(opts...) + + executor := runnableLambda(i, s, c, t, + !opt.enableComponentCallback, + ) + executor.meta = &executorMeta{ + component: ComponentOfLambda, + isComponentCallbackEnabled: opt.enableComponentCallback, + componentImplType: opt.componentImplType, + } + + return &Lambda{ + executor: executor, + } +} + +func getLambdaOpt(opts ...LambdaOpt) *lambdaOpts { + opt := &lambdaOpts{ + enableComponentCallback: false, + componentImplType: "", + } + + for _, optFn := range opts { + optFn(opt) + } + return opt +} + +// ToList creates a Lambda that converts input I to a []I. +// It's useful when you want to convert a single input to a list of inputs. +// eg. +// +// lambda := compose.ToList[*schema.Message]() +// chain := compose.NewChain[[]*schema.Message, []*schema.Message]() +// +// chain.AddChatModel(chatModel) // chatModel returns *schema.Message, but we need []*schema.Message +// chain.AddLambda(lambda) // convert *schema.Message to []*schema.Message +func ToList[I any](opts ...LambdaOpt) *Lambda { + i := func(ctx context.Context, input I, opts_ ...unreachableOption) (output []I, err error) { + return []I{input}, nil + } + + f := func(ctx context.Context, inputS *schema.StreamReader[I], opts_ ...unreachableOption) (outputS *schema.StreamReader[[]I], err error) { + return schema.StreamReaderWithConvert(inputS, func(i I) ([]I, error) { + return []I{i}, nil + }), nil + } + + return anyLambda(i, nil, nil, f, opts...) +} + +// MessageParser creates a lambda that parses a message into an object T, usually used after a chatmodel. +// usage: +// +// parser := schema.NewMessageJSONParser[MyStruct](&schema.MessageJSONParseConfig{ +// ParseFrom: schema.MessageParseFromContent, +// }) +// parserLambda := MessageParser(parser) +// +// chain := NewChain[*schema.Message, MyStruct]() +// chain.AppendChatModel(chatModel) +// chain.AppendLambda(parserLambda) +// +// r, err := chain.Compile(context.Background()) +// +// // parsed is a MyStruct object +// parsed, err := r.Invoke(context.Background(), &schema.Message{ +// Role: schema.MessageRoleUser, +// Content: "return a json string for my struct", +// }) +func MessageParser[T any](p schema.MessageParser[T], opts ...LambdaOpt) *Lambda { + i := func(ctx context.Context, input *schema.Message, opts_ ...unreachableOption) (output T, err error) { + return p.Parse(ctx, input) + } + + opts = append([]LambdaOpt{WithLambdaType("MessageParse")}, opts...) + + return anyLambda(i, nil, nil, nil, opts...) +} diff --git a/compose/types_lambda_test.go b/compose/types_lambda_test.go new file mode 100644 index 0000000..96afc0b --- /dev/null +++ b/compose/types_lambda_test.go @@ -0,0 +1,212 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/schema" +) + +func TestLambda(t *testing.T) { + t.Run("InvokableLambda", func(t *testing.T) { + ld := InvokableLambdaWithOption( + func(ctx context.Context, input string, opts ...any) (output string, err error) { + return "good", nil + }, + WithLambdaCallbackEnable(false), + WithLambdaType("ForTest"), + ) + + assert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "ForTest", ld.executor.meta.componentImplType) + + ld = InvokableLambda( + func(ctx context.Context, input string) (output string, err error) { + return "good", nil + }, + WithLambdaCallbackEnable(false), + WithLambdaType("ForTest"), + ) + + assert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "ForTest", ld.executor.meta.componentImplType) + }) + + t.Run("StreamableLambda", func(t *testing.T) { + ld := StreamableLambdaWithOption( + func(ctx context.Context, input string, opts ...any) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](1) + sw.Close() + return sr, nil + }, + WithLambdaCallbackEnable(false), + WithLambdaType("ForTest"), + ) + + assert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "ForTest", ld.executor.meta.componentImplType) + + ld = StreamableLambda( + func(ctx context.Context, input string) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](1) + sw.Close() + return sr, nil + }, + WithLambdaCallbackEnable(false), + WithLambdaType("ForTest"), + ) + + assert.Equal(t, false, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "ForTest", ld.executor.meta.componentImplType) + }) + + t.Run("CollectableLambda", func(t *testing.T) { + ld := CollectableLambdaWithOption( + func(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output string, err error) { + return "good", nil + }, + WithLambdaCallbackEnable(true), + ) + + assert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "", ld.executor.meta.componentImplType) + + ld = CollectableLambda( + func(ctx context.Context, input *schema.StreamReader[string]) (output string, err error) { + return "good", nil + }, + WithLambdaCallbackEnable(true), + ) + + assert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "", ld.executor.meta.componentImplType) + }) + + t.Run("TransformableLambda", func(t *testing.T) { + ld := TransformableLambdaWithOption( + func(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](1) + sw.Close() + return sr, nil + }, + WithLambdaCallbackEnable(true), + ) + + assert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "", ld.executor.meta.componentImplType) + + ld = TransformableLambda( + func(ctx context.Context, input *schema.StreamReader[string]) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](1) + sw.Close() + return sr, nil + }, + WithLambdaCallbackEnable(true), + ) + + assert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "", ld.executor.meta.componentImplType) + }) + + t.Run("AnyLambda", func(t *testing.T) { + ld, err := AnyLambda[string, string]( + func(ctx context.Context, input string, opts ...any) (output string, err error) { + return "good", nil + }, + func(ctx context.Context, input string, opts ...any) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](1) + sw.Close() + return sr, nil + }, + func(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output string, err error) { + return "good", nil + }, + func(ctx context.Context, input *schema.StreamReader[string], opts ...any) (output *schema.StreamReader[string], err error) { + sr, sw := schema.Pipe[string](1) + sw.Close() + return sr, nil + }, + WithLambdaCallbackEnable(true), + WithLambdaType("ForTest"), + ) + assert.NoError(t, err) + + assert.Equal(t, true, ld.executor.meta.isComponentCallbackEnabled) + assert.Equal(t, ComponentOfLambda, ld.executor.meta.component) + assert.Equal(t, "ForTest", ld.executor.meta.componentImplType) + }) +} + +type TestStructForParse struct { + ID int `json:"id"` +} + +func TestMessageParser(t *testing.T) { + t.Run("parse from content", func(t *testing.T) { + parser := schema.NewMessageJSONParser[TestStructForParse](&schema.MessageJSONParseConfig{ + ParseFrom: schema.MessageParseFromContent, + }) + + parserLambda := MessageParser(parser) + + chain := NewChain[*schema.Message, TestStructForParse]() + chain.AppendLambda(parserLambda) + + r, err := chain.Compile(context.Background()) + assert.Nil(t, err) + + parsed, err := r.Invoke(context.Background(), &schema.Message{ + Content: `{"id": 1}`, + }) + assert.Nil(t, err) + assert.Equal(t, 1, parsed.ID) + }) + + t.Run("parse from tool call", func(t *testing.T) { + parser := schema.NewMessageJSONParser[*TestStructForParse](&schema.MessageJSONParseConfig{ + ParseFrom: schema.MessageParseFromToolCall, + }) + + parserLambda := MessageParser(parser) + + chain := NewChain[*schema.Message, *TestStructForParse]() + chain.AppendLambda(parserLambda) + + r, err := chain.Compile(context.Background()) + assert.Nil(t, err) + + parsed, err := r.Invoke(context.Background(), &schema.Message{ + ToolCalls: []schema.ToolCall{ + {Function: schema.FunctionCall{Arguments: `{"id": 1}`}}, + }, + }) + assert.Nil(t, err) + assert.Equal(t, 1, parsed.ID) + }) +} diff --git a/compose/utils.go b/compose/utils.go new file mode 100644 index 0000000..c348c95 --- /dev/null +++ b/compose/utils.go @@ -0,0 +1,324 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + + "github.com/cloudwego/eino/callbacks" + icb "github.com/cloudwego/eino/internal/callbacks" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +type on[T any] func(context.Context, T) (context.Context, T) + +func onStart[T any](ctx context.Context, input T) (context.Context, T) { + return icb.On(ctx, input, icb.OnStartHandle[T], callbacks.TimingOnStart, true) +} + +func onEnd[T any](ctx context.Context, output T) (context.Context, T) { + return icb.On(ctx, output, icb.OnEndHandle[T], callbacks.TimingOnEnd, false) +} + +func onStartWithStreamInput[T any](ctx context.Context, input *schema.StreamReader[T]) ( + context.Context, *schema.StreamReader[T]) { + + return icb.On(ctx, input, icb.OnStartWithStreamInputHandle[T], callbacks.TimingOnStartWithStreamInput, true) +} + +func genericOnStartWithStreamInputHandle(ctx context.Context, input streamReader, + runInfo *icb.RunInfo, handlers []icb.Handler) (context.Context, streamReader) { + + handlers = generic.Reverse(handlers) + + cpy := input.copy + + handle := func(ctx context.Context, handler icb.Handler, in streamReader) context.Context { + in_, ok := unpackStreamReader[icb.CallbackInput](in) + if !ok { + panic("impossible") + } + + return handler.OnStartWithStreamInput(ctx, runInfo, in_) + } + + return icb.OnWithStreamHandle(ctx, input, handlers, cpy, handle) +} + +func genericOnStartWithStreamInput(ctx context.Context, input streamReader) (context.Context, streamReader) { + return icb.On(ctx, input, genericOnStartWithStreamInputHandle, callbacks.TimingOnStartWithStreamInput, true) +} + +func onEndWithStreamOutput[T any](ctx context.Context, output *schema.StreamReader[T]) ( + context.Context, *schema.StreamReader[T]) { + + return icb.On(ctx, output, icb.OnEndWithStreamOutputHandle[T], callbacks.TimingOnEndWithStreamOutput, false) +} + +func genericOnEndWithStreamOutputHandle(ctx context.Context, output streamReader, + runInfo *icb.RunInfo, handlers []icb.Handler) (context.Context, streamReader) { + + cpy := output.copy + + handle := func(ctx context.Context, handler icb.Handler, out streamReader) context.Context { + out_, ok := unpackStreamReader[icb.CallbackOutput](out) + if !ok { + panic("impossible") + } + + return handler.OnEndWithStreamOutput(ctx, runInfo, out_) + } + + return icb.OnWithStreamHandle(ctx, output, handlers, cpy, handle) +} + +func genericOnEndWithStreamOutput(ctx context.Context, output streamReader) (context.Context, streamReader) { + return icb.On(ctx, output, genericOnEndWithStreamOutputHandle, callbacks.TimingOnEndWithStreamOutput, false) +} + +func onError(ctx context.Context, err error) (context.Context, error) { + return icb.On(ctx, err, icb.OnErrorHandle, callbacks.TimingOnError, false) +} + +func runWithCallbacks[I, O, TOption any](r func(context.Context, I, ...TOption) (O, error), + onStart on[I], onEnd on[O], onError on[error]) func(context.Context, I, ...TOption) (O, error) { + + return func(ctx context.Context, input I, opts ...TOption) (output O, err error) { + ctx, input = onStart(ctx, input) + + output, err = r(ctx, input, opts...) + if err != nil { + ctx, err = onError(ctx, err) + return output, err + } + + ctx, output = onEnd(ctx, output) + + return output, nil + } +} + +func invokeWithCallbacks[I, O, TOption any](i Invoke[I, O, TOption]) Invoke[I, O, TOption] { + return runWithCallbacks(i, onStart[I], onEnd[O], onError) +} + +func onGraphStart(ctx context.Context, input any, isStream bool) (context.Context, any) { + if isStream { + return genericOnStartWithStreamInput(ctx, input.(streamReader)) + } + return onStart(ctx, input) +} + +func onGraphEnd(ctx context.Context, output any, isStream bool) (context.Context, any) { + if isStream { + return genericOnEndWithStreamOutput(ctx, output.(streamReader)) + } + return onEnd(ctx, output) +} + +func onGraphError(ctx context.Context, err error) (context.Context, error) { + return onError(ctx, err) +} + +func streamWithCallbacks[I, O, TOption any](s Stream[I, O, TOption]) Stream[I, O, TOption] { + return runWithCallbacks(s, onStart[I], onEndWithStreamOutput[O], onError) +} + +func collectWithCallbacks[I, O, TOption any](c Collect[I, O, TOption]) Collect[I, O, TOption] { + return runWithCallbacks(c, onStartWithStreamInput[I], onEnd[O], onError) +} + +func transformWithCallbacks[I, O, TOption any](t Transform[I, O, TOption]) Transform[I, O, TOption] { + return runWithCallbacks(t, onStartWithStreamInput[I], onEndWithStreamOutput[O], onError) +} + +func initGraphCallbacks(ctx context.Context, info *nodeInfo, meta *executorMeta, opts ...Option) context.Context { + ri := &callbacks.RunInfo{} + if meta != nil { + ri.Component = meta.component + ri.Type = meta.componentImplType + } + + if info != nil { + ri.Name = info.name + } + + var cbs []callbacks.Handler + for i := range opts { + if len(opts[i].handler) != 0 && len(opts[i].paths) == 0 { + cbs = append(cbs, opts[i].handler...) + } + } + + if len(cbs) == 0 { + return icb.ReuseHandlers(ctx, ri) + } + + return icb.AppendHandlers(ctx, ri, cbs...) +} + +func initNodeCallbacks(ctx context.Context, key string, info *nodeInfo, meta *executorMeta, opts ...Option) context.Context { + ri := &callbacks.RunInfo{} + if meta != nil { + ri.Component = meta.component + ri.Type = meta.componentImplType + } + + if info != nil { + ri.Name = info.name + } + + var cbs []callbacks.Handler + for i := range opts { + if len(opts[i].handler) != 0 { + if len(opts[i].paths) != 0 { + for _, k := range opts[i].paths { + if len(k.path) == 1 && k.path[0] == key { + cbs = append(cbs, opts[i].handler...) + break + } + } + } + } + } + + if len(cbs) == 0 { + return icb.ReuseHandlers(ctx, ri) + } + + return icb.AppendHandlers(ctx, ri, cbs...) +} + +func streamChunkConvertForCBOutput[O any](o O) (callbacks.CallbackOutput, error) { + return o, nil +} + +func streamChunkConvertForCBInput[I any](i I) (callbacks.CallbackInput, error) { + return i, nil +} + +func toAnyList[T any](in []T) []any { + ret := make([]any, len(in)) + for i := range in { + ret[i] = in[i] + } + return ret +} + +type assignableType uint8 + +const ( + assignableTypeMustNot assignableType = iota + assignableTypeMust + assignableTypeMay +) + +func checkAssignable(input, arg reflect.Type) assignableType { + if arg == nil || input == nil { + return assignableTypeMustNot + } + + if arg == input { + return assignableTypeMust + } + + if arg.Kind() == reflect.Interface && input.Implements(arg) { + return assignableTypeMust + } + if input.Kind() == reflect.Interface { + if arg.Implements(input) { + return assignableTypeMay + } + return assignableTypeMustNot + } + + return assignableTypeMustNot +} + +func extractOption(nodes map[string]*chanCall, opts ...Option) (map[string][]any, error) { + optMap := map[string][]any{} + for _, opt := range opts { + if len(opt.paths) == 0 { + // common, discard callback, filter option by type + if len(opt.options) == 0 { + continue + } + for name, c := range nodes { + if c.action.optionType == nil { + // subgraph + optMap[name] = append(optMap[name], opt) + } else if reflect.TypeOf(opt.options[0]) == c.action.optionType { // assume that types of options are the same + optMap[name] = append(optMap[name], opt.options...) + } + } + } + for _, path := range opt.paths { + if len(path.path) == 0 { + return nil, fmt.Errorf("call option has designated an empty path") + } + + var curNode *chanCall + var ok bool + if curNode, ok = nodes[path.path[0]]; !ok { + return nil, fmt.Errorf("option has designated an unknown node: %s", path) + } + curNodeKey := path.path[0] + + if len(path.path) == 1 { + if len(opt.options) == 0 { + // sub graph common callbacks has been added to ctx in initNodeCallback and won't be passed to subgraph only pass options + // node callback also won't be passed + continue + } + if curNode.action.optionType == nil { + nOpt := opt.deepCopy() + nOpt.paths = []*NodePath{} + optMap[curNodeKey] = append(optMap[curNodeKey], nOpt) + } else { + // designate to component + if curNode.action.optionType != reflect.TypeOf(opt.options[0]) { // assume that types of options are the same + return nil, fmt.Errorf("option type[%s] is different from which the designated node[%s] expects[%s]", + reflect.TypeOf(opt.options[0]).String(), path, curNode.action.optionType.String()) + } + optMap[curNodeKey] = append(optMap[curNodeKey], opt.options...) + } + } else { + if curNode.action.optionType != nil { + // component + return nil, fmt.Errorf("cannot designate sub path of a component, path:%s", path) + } + // designate to sub graph's nodes + nOpt := opt.deepCopy() + nOpt.paths = []*NodePath{NewNodePath(path.path[1:]...)} + optMap[curNodeKey] = append(optMap[curNodeKey], nOpt) + } + } + } + + return optMap, nil +} + +func mapToList(m map[string]any) []any { + ret := make([]any, 0, len(m)) + for _, v := range m { + ret = append(ret, v) + } + return ret +} diff --git a/compose/utils_test.go b/compose/utils_test.go new file mode 100644 index 0000000..07a781c --- /dev/null +++ b/compose/utils_test.go @@ -0,0 +1,114 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/generic" +) + +type good interface { + ThisIsGood() bool +} + +type good2 interface { + ThisIsGood2() bool +} + +type good3 interface { + ThisIsGood() bool +} + +type goodImpl struct{} + +func (g *goodImpl) ThisIsGood() bool { + return true +} + +type goodNotImpl struct{} + +func TestValidateType(t *testing.T) { + + t.Run("equal_type", func(t *testing.T) { + arg := generic.TypeOf[int]() + input := generic.TypeOf[int]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMust, result) + }) + + t.Run("unequal_type", func(t *testing.T) { + arg := generic.TypeOf[int]() + input := generic.TypeOf[string]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMustNot, result) + }) + + t.Run("implement_interface", func(t *testing.T) { + arg := generic.TypeOf[good]() + input := generic.TypeOf[*goodImpl]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMust, result) + }) + + t.Run("may_implement_interface", func(t *testing.T) { + arg := generic.TypeOf[*goodImpl]() + input := generic.TypeOf[good]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMay, result) + }) + + t.Run("not_implement_interface", func(t *testing.T) { + arg := generic.TypeOf[good]() + input := generic.TypeOf[*goodNotImpl]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMustNot, result) + }) + + t.Run("interface_unequal_interface", func(t *testing.T) { + arg := generic.TypeOf[good]() + input := generic.TypeOf[good2]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMustNot, result) + }) + + t.Run("interface_equal_interface", func(t *testing.T) { + arg := generic.TypeOf[good]() + input := generic.TypeOf[good3]() + + result := checkAssignable(input, arg) + assert.Equal(t, assignableTypeMust, result) + }) +} + +func TestStreamChunkConvert(t *testing.T) { + o, err := streamChunkConvertForCBOutput(1) + assert.Nil(t, err) + assert.Equal(t, o, 1) + + i, err := streamChunkConvertForCBInput(1) + assert.Nil(t, err) + assert.Equal(t, i, 1) +} diff --git a/compose/values_merge.go b/compose/values_merge.go new file mode 100644 index 0000000..c08a12e --- /dev/null +++ b/compose/values_merge.go @@ -0,0 +1,82 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "fmt" + "reflect" + + "github.com/cloudwego/eino/internal" +) + +// RegisterValuesMergeFunc registers a function to merge outputs from multiple nodes when fan-in. +// It's used to define how to merge for a specific type. +// For maps that already have a default merge function, you don't need to register a new one unless you want to customize the merge logic. +func RegisterValuesMergeFunc[T any](fn func([]T) (T, error)) { + internal.RegisterValuesMergeFunc(fn) +} + +type mergeOptions struct { + streamMergeWithSourceEOF bool + names []string +} + +// the caller should ensure len(vs) > 1 +func mergeValues(vs []any, opts *mergeOptions) (any, error) { + v0 := reflect.ValueOf(vs[0]) + t0 := v0.Type() + + if fn := internal.GetMergeFunc(t0); fn != nil { + return fn(vs) + } + + // merge StreamReaders + if s, ok := vs[0].(streamReader); ok { + t := s.getChunkType() + if internal.GetMergeFunc(t) == nil { + return nil, fmt.Errorf("(mergeValues | stream type)"+ + " unsupported chunk type: %v", t) + } + + ss := make([]streamReader, len(vs)-1) + for i := 0; i < len(ss); i++ { + sri, ok_ := vs[i+1].(streamReader) + if !ok_ { + return nil, fmt.Errorf("(mergeStream) unexpected type. "+ + "expect: %v, got: %v", t0, reflect.TypeOf(vs[i])) + } + + if st := sri.getChunkType(); st != t { + return nil, fmt.Errorf("(mergeStream) chunk type mismatch. "+ + "expect: %v, got: %v", t, st) + } + + ss[i] = sri + } + + if opts != nil && opts.streamMergeWithSourceEOF { + ms := s.mergeWithNames(ss, opts.names) + return ms, nil + } + + ms := s.merge(ss) + + return ms, nil + } + + return nil, fmt.Errorf("(mergeValues) unsupported type: %v", t0) +} diff --git a/compose/values_merge_test.go b/compose/values_merge_test.go new file mode 100644 index 0000000..a848397 --- /dev/null +++ b/compose/values_merge_test.go @@ -0,0 +1,256 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "fmt" + "io" + "sort" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +func Test_mergeValues(t *testing.T) { + t.Run("merge maps", func(t *testing.T) { + m1 := map[int]int{1: 1, 2: 2, 3: 3, 4: 4} + m2 := map[int]int{5: 5, 6: 6, 7: 7, 8: 8} + m3 := map[int]int{9: 9, 10: 10, 11: 11} + + t.Run("regular", func(t *testing.T) { + mergedM, err := mergeValues([]any{m1, m2, m3}, nil) + assert.NoError(t, err) + + m := mergedM.(map[int]int) + + // len(m) == len(m1) + len(m2) + len(m3) + assert.Equal(t, len(m), len(m1)+len(m2)+len(m3)) + }) + + t.Run("duplicated key", func(t *testing.T) { + _, err := mergeValues([]any{m1, m2, m3, map[int]int{1: 1}}, nil) + assert.ErrorContains(t, err, "duplicated key") + }) + + t.Run("type mismatch", func(t *testing.T) { + _, err := mergeValues([]any{m1, m2, m3, map[int]string{1: "1"}}, nil) + assert.ErrorContains(t, err, "type mismatch") + }) + }) + + t.Run("merge stream", func(t *testing.T) { + ass := []any{ + packStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{1: "1"}})), + packStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{2: "2"}})), + packStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{3: "3", 4: "4"}})), + } + isr, err := mergeValues(ass, nil) + require.NoError(t, err) + ret, ok := unpackStreamReader[map[int]string](isr.(streamReader)) + require.True(t, ok) + defer ret.Close() + + got := make(map[int]string) + for i := 0; i < 3; i++ { + m, err := ret.Recv() + require.NoError(t, err) + for k, v := range m { + got[k] = v + } + } + _, err = ret.Recv() + require.ErrorIs(t, err, io.EOF) + + assert.Equal(t, map[int]string{ + 1: "1", + 2: "2", + 3: "3", + 4: "4", + }, got) + }) + + t.Run("merge stream with source EOF", func(t *testing.T) { + ass := []any{ + packStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{1: "1"}})), + packStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{2: "2"}})), + packStreamReader(schema.StreamReaderFromArray[map[int]string]([]map[int]string{{3: "3", 4: "4"}})), + } + opts := &mergeOptions{ + streamMergeWithSourceEOF: true, + names: []string{ + "source0", + "source1", + "source2", + }, + } + isr, err := mergeValues(ass, opts) + require.NoError(t, err) + ret, ok := unpackStreamReader[map[int]string](isr.(streamReader)) + require.True(t, ok) + defer ret.Close() + + got := make(map[int]string) + endedSources := make(map[string]bool) + + for { + m, e := ret.Recv() + if e != nil { + if sourceName, ok_ := schema.GetSourceName(e); ok_ { + t.Logf("Source '%s' ended", sourceName) + endedSources[sourceName] = true + continue + } + if e == io.EOF { + // This EOF means all chunks from all sources that were not SourceEOF have been merged and sent. + // Or, if all sources send SourceEOF first, this io.EOF means the merged stream itself is now empty. + break + } + + require.NoError(t, e) // Fail on any other error + } + // If streamMergeWithSourceEOF is true, the final merged result comes as a single map chunk + // after all SourceEOFs (if any non-empty streams existed) or directly if all streams were empty. + for k, v := range m { + got[k] = v + } + } + + // Check that all expected sources have ended if they were part of opts.names + for i := 0; i < len(ass); i++ { + expectedSourceName := opts.names[i] + assert.True(t, endedSources[expectedSourceName], "Expected source %s to have sent SourceEOF", expectedSourceName) + } + + // The final 'got' map should contain all items because streamMergeWithSourceEOF merges them at the end. + assert.Equal(t, map[int]string{ + 1: "1", + 2: "2", + 3: "3", + 4: "4", + }, got) + }) + + type TestType struct { + A int + B []string + } + + RegisterValuesMergeFunc(func(vs []*TestType) (*TestType, error) { + ret := &TestType{} + for _, v := range vs { + if v == nil { + continue + } + if ret.A < 0 { + return nil, fmt.Errorf("test error: %v", ret.A) + } + ret.A += v.A + ret.B = append(ret.B, v.B...) + } + sort.Strings(ret.B) + return ret, nil + }) + + t.Run("custom merge", func(t *testing.T) { + t.Run("regular", func(t *testing.T) { + vs := []any{ + &TestType{A: 0, B: []string{}}, + &TestType{A: 1, B: []string{"1"}}, + &TestType{A: 2, B: []string{"2", "22"}}, + &TestType{A: 3, B: []string{"3", "33", "333"}}, + } + ret, err := mergeValues(vs, nil) + require.NoError(t, err) + assert.Equal(t, &TestType{ + A: 6, + B: []string{"1", "2", "22", "3", "33", "333"}, + }, ret) + }) + + t.Run("custom error", func(t *testing.T) { + vs := []any{ + &TestType{A: 0, B: []string{}}, + &TestType{A: 1, B: []string{"1"}}, + &TestType{A: -2, B: []string{"2", "22"}}, + &TestType{A: 3, B: []string{"3", "33", "333"}}, + } + _, err := mergeValues(vs, nil) + require.ErrorContains(t, err, "test error") + }) + + t.Run("type mismatch", func(t *testing.T) { + vs := []any{ + &TestType{A: 0, B: []string{}}, + &TestType{A: 1, B: []string{"1"}}, + &TestType{A: 2, B: []string{"2", "22"}}, + "test3", + } + _, err := mergeValues(vs, nil) + require.ErrorContains(t, err, "type mismatch") + }) + + t.Run("stream", func(t *testing.T) { + ass := []any{ + packStreamReader(schema.StreamReaderFromArray([]*TestType{ + {A: 0, B: []string{}}, + })), + packStreamReader(schema.StreamReaderFromArray([]*TestType{ + {A: 1, B: []string{"1"}}, + })), + packStreamReader(schema.StreamReaderFromArray([]*TestType{ + {A: 2, B: []string{"2", "22"}}, + })), + packStreamReader(schema.StreamReaderFromArray([]*TestType{ + {A: 3, B: []string{"3", "33", "333"}}, + })), + } + isr, err := mergeValues(ass, nil) + require.NoError(t, err) + ret, ok := unpackStreamReader[*TestType](isr.(streamReader)) + require.True(t, ok) + defer ret.Close() + + var vs []any + for i := 0; i < 4; i++ { + v, err := ret.Recv() + require.NoError(t, err) + vs = append(vs, v) + } + + _, err = ret.Recv() + require.ErrorIs(t, err, io.EOF) + + merged, err := mergeValues(vs, nil) + require.NoError(t, err) + + assert.Equal(t, &TestType{ + A: 6, + B: []string{"1", "2", "22", "3", "33", "333"}, + }, merged) + }) + }) + + t.Run("unregistered type", func(t *testing.T) { + type Unregistered TestType + _, err := mergeValues([]any{&Unregistered{}}, nil) + assert.ErrorContains(t, err, "unsupported type") + }) + +} diff --git a/compose/workflow.go b/compose/workflow.go new file mode 100644 index 0000000..6b50962 --- /dev/null +++ b/compose/workflow.go @@ -0,0 +1,545 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "fmt" + "reflect" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +// WorkflowNode is the node of the Workflow. +type WorkflowNode struct { + g *graph + key string + addInputs []func() error + staticValues map[string]any + dependencySetter func(fromNodeKey string, typ dependencyType) + mappedFieldPath map[string]any +} + +// Workflow is wrapper of graph, replacing AddEdge with declaring dependencies and field mappings between nodes. +// Under the hood it uses NodeTriggerMode(AllPredecessor), so does not support cycles. +type Workflow[I, O any] struct { + g *graph + workflowNodes map[string]*WorkflowNode + workflowBranches []*WorkflowBranch + dependencies map[string]map[string]dependencyType +} + +type dependencyType int + +const ( + normalDependency dependencyType = iota + noDirectDependency + branchDependency +) + +// NewWorkflow creates a new Workflow. +func NewWorkflow[I, O any](opts ...NewGraphOption) *Workflow[I, O] { + options := &newGraphOptions{} + for _, opt := range opts { + opt(options) + } + + wf := &Workflow[I, O]{ + g: newGraphFromGeneric[I, O]( + ComponentOfWorkflow, + options.withState, + options.stateType, + opts, + ), + workflowNodes: make(map[string]*WorkflowNode), + dependencies: make(map[string]map[string]dependencyType), + } + + return wf +} + +// Compile builds the workflow into a runnable graph. +func (wf *Workflow[I, O]) Compile(ctx context.Context, opts ...GraphCompileOption) (Runnable[I, O], error) { + return compileAnyGraph[I, O](ctx, wf, opts...) +} + +// AddChatModelNode adds a chat model node and returns it. +func (wf *Workflow[I, O]) AddChatModelNode(key string, chatModel model.BaseChatModel, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddChatModelNode(key, chatModel, opts...) + return wf.initNode(key) +} + +// AddAgenticModelNode adds an agentic model node and returns it. +func (wf *Workflow[I, O]) AddAgenticModelNode(key string, agenticModel model.AgenticModel, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddAgenticModelNode(key, agenticModel, opts...) + return wf.initNode(key) +} + +// AddChatTemplateNode adds a chat template node and returns it. +func (wf *Workflow[I, O]) AddChatTemplateNode(key string, chatTemplate prompt.ChatTemplate, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddChatTemplateNode(key, chatTemplate, opts...) + return wf.initNode(key) +} + +// AddAgenticChatTemplateNode adds an agentic chat template node and returns it. +func (wf *Workflow[I, O]) AddAgenticChatTemplateNode(key string, chatTemplate prompt.AgenticChatTemplate, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddAgenticChatTemplateNode(key, chatTemplate, opts...) + return wf.initNode(key) +} + +// AddToolsNode adds a tools node and returns it. +func (wf *Workflow[I, O]) AddToolsNode(key string, tools *ToolsNode, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddToolsNode(key, tools, opts...) + return wf.initNode(key) +} + +// AddAgenticToolsNode adds an agentic tools node and returns it. +func (wf *Workflow[I, O]) AddAgenticToolsNode(key string, tools *AgenticToolsNode, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddAgenticToolsNode(key, tools, opts...) + return wf.initNode(key) +} + +// AddRetrieverNode adds a retriever node and returns it. +func (wf *Workflow[I, O]) AddRetrieverNode(key string, retriever retriever.Retriever, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddRetrieverNode(key, retriever, opts...) + return wf.initNode(key) +} + +// AddEmbeddingNode adds an embedding node and returns it. +func (wf *Workflow[I, O]) AddEmbeddingNode(key string, embedding embedding.Embedder, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddEmbeddingNode(key, embedding, opts...) + return wf.initNode(key) +} + +// AddIndexerNode adds an indexer node to the workflow and returns it. +func (wf *Workflow[I, O]) AddIndexerNode(key string, indexer indexer.Indexer, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddIndexerNode(key, indexer, opts...) + return wf.initNode(key) +} + +// AddLoaderNode adds a document loader node to the workflow and returns it. +func (wf *Workflow[I, O]) AddLoaderNode(key string, loader document.Loader, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddLoaderNode(key, loader, opts...) + return wf.initNode(key) +} + +// AddDocumentTransformerNode adds a document transformer node and returns it. +func (wf *Workflow[I, O]) AddDocumentTransformerNode(key string, transformer document.Transformer, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddDocumentTransformerNode(key, transformer, opts...) + return wf.initNode(key) +} + +// AddGraphNode adds a nested graph node to the workflow and returns it. +func (wf *Workflow[I, O]) AddGraphNode(key string, graph AnyGraph, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddGraphNode(key, graph, opts...) + return wf.initNode(key) +} + +// AddLambdaNode adds a lambda node to the workflow and returns it. +func (wf *Workflow[I, O]) AddLambdaNode(key string, lambda *Lambda, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddLambdaNode(key, lambda, opts...) + return wf.initNode(key) +} + +// End returns the WorkflowNode representing END node. +func (wf *Workflow[I, O]) End() *WorkflowNode { + if node, ok := wf.workflowNodes[END]; ok { + return node + } + return wf.initNode(END) +} + +// AddPassthroughNode adds a passthrough node to the workflow and returns it. +func (wf *Workflow[I, O]) AddPassthroughNode(key string, opts ...GraphAddNodeOpt) *WorkflowNode { + _ = wf.g.AddPassthroughNode(key, opts...) + return wf.initNode(key) +} + +// AddInput creates both data and execution dependencies between nodes. +// It configures how data flows from the predecessor node (fromNodeKey) to the current node, +// and ensures the current node only executes after the predecessor completes. +// +// Parameters: +// - fromNodeKey: the key of the predecessor node +// - inputs: field mappings that specify how data should flow from the predecessor +// to the current node. If no mappings are provided, the entire output of the +// predecessor will be used as input. +// +// Example: +// +// // Map between specific field +// node.AddInput("userNode", MapFields("user.name", "displayName")) +// +// // Use entire output +// node.AddInput("dataNode") +// +// Returns the current node for method chaining. +func (n *WorkflowNode) AddInput(fromNodeKey string, inputs ...*FieldMapping) *WorkflowNode { + return n.addDependencyRelation(fromNodeKey, inputs, &workflowAddInputOpts{}) +} + +type workflowAddInputOpts struct { + // noDirectDependency indicates whether to create a data mapping without establishing + // a direct execution dependency. When true, the current node can access data from + // the predecessor node but its execution is not directly blocked by it. + noDirectDependency bool + // dependencyWithoutInput indicates whether to create an execution dependency + // without any data mapping. When true, the current node will wait for the + // predecessor node to complete but won't receive any data from it. + dependencyWithoutInput bool +} + +// WorkflowAddInputOpt configures behavior of AddInputWithOptions. +type WorkflowAddInputOpt func(*workflowAddInputOpts) + +func getAddInputOpts(opts []WorkflowAddInputOpt) *workflowAddInputOpts { + opt := &workflowAddInputOpts{} + for _, o := range opts { + o(opt) + } + return opt +} + +// WithNoDirectDependency creates a data mapping without establishing a direct execution dependency. +// The predecessor node will still complete before the current node executes, but through indirect +// execution paths rather than a direct dependency. +// +// In a workflow graph, node dependencies typically serve two purposes: +// 1. Execution order: determining when nodes should run +// 2. Data flow: specifying how data passes between nodes +// +// This option separates these concerns by: +// - Creating data mapping from the predecessor to the current node +// - Relying on the predecessor's path to reach the current node through other nodes +// that have direct execution dependencies +// +// Example: +// +// node.AddInputWithOptions("dataNode", mappings, WithNoDirectDependency()) +// +// Important: +// +// 1. Branch scenarios: When connecting nodes on different sides of a branch, +// WithNoDirectDependency MUST be used to let the branch itself handle the +// execution order, preventing incorrect dependencies that could bypass the branch. +// +// 2. Execution guarantee: The predecessor will still complete before the current +// node executes because the predecessor must have a path (through other nodes) +// that eventually reaches the current node. +// +// 3. Graph validity: There MUST be a path from the predecessor that eventually +// reaches the current node through other nodes with direct dependencies. +// This ensures the execution order while avoiding redundant direct dependencies. +// +// Common use cases: +// - Cross-branch data access where the branch handles execution order +// - Avoiding redundant dependencies when a path already exists +func WithNoDirectDependency() WorkflowAddInputOpt { + return func(opt *workflowAddInputOpts) { + opt.noDirectDependency = true + } +} + +// AddInputWithOptions creates a dependency between nodes with custom configuration options. +// It allows fine-grained control over both data flow and execution dependencies. +// +// Parameters: +// - fromNodeKey: the key of the predecessor node +// - inputs: field mappings that specify how data flows from the predecessor to the current node. +// If no mappings are provided, the entire output of the predecessor will be used as input. +// - opts: configuration options that control how the dependency is established +// +// Example: +// +// // Create data mapping without direct execution dependency +// node.AddInputWithOptions("dataNode", mappings, WithNoDirectDependency()) +// +// Returns the current node for method chaining. +func (n *WorkflowNode) AddInputWithOptions(fromNodeKey string, inputs []*FieldMapping, opts ...WorkflowAddInputOpt) *WorkflowNode { + return n.addDependencyRelation(fromNodeKey, inputs, getAddInputOpts(opts)) +} + +// AddDependency creates an execution-only dependency between nodes. +// The current node will wait for the predecessor node to complete before executing, +// but no data will be passed between them. +// +// Parameters: +// - fromNodeKey: the key of the predecessor node that must complete before this node starts +// +// Example: +// +// // Wait for "setupNode" to complete before executing +// node.AddDependency("setupNode") +// +// This is useful when: +// - You need to ensure execution order without data transfer +// - The predecessor performs setup or initialization that must complete first +// - You want to explicitly separate execution dependencies from data flow +// +// Returns the current node for method chaining. +func (n *WorkflowNode) AddDependency(fromNodeKey string) *WorkflowNode { + return n.addDependencyRelation(fromNodeKey, nil, &workflowAddInputOpts{dependencyWithoutInput: true}) +} + +// SetStaticValue sets a static value for a field path that will be available +// during workflow execution. These values are determined at compile time and +// remain constant throughout the workflow's lifecycle. +// +// Example: +// +// node.SetStaticValue(FieldPath{"query"}, "static query") +func (n *WorkflowNode) SetStaticValue(path FieldPath, value any) *WorkflowNode { + n.staticValues[path.join()] = value + return n +} + +func (n *WorkflowNode) addDependencyRelation(fromNodeKey string, inputs []*FieldMapping, options *workflowAddInputOpts) *WorkflowNode { + for _, input := range inputs { + input.fromNodeKey = fromNodeKey + } + + if options.noDirectDependency { + n.addInputs = append(n.addInputs, func() error { + var paths []FieldPath + for _, input := range inputs { + paths = append(paths, input.targetPath()) + } + if err := n.checkAndAddMappedPath(paths); err != nil { + return err + } + + if err := n.g.addEdgeWithMappings(fromNodeKey, n.key, true, false, inputs...); err != nil { + return err + } + n.dependencySetter(fromNodeKey, noDirectDependency) + return nil + }) + } else if options.dependencyWithoutInput { + n.addInputs = append(n.addInputs, func() error { + if len(inputs) > 0 { + return fmt.Errorf("dependency without input should not have inputs. node: %s, fromNode: %s, inputs: %v", n.key, fromNodeKey, inputs) + } + if err := n.g.addEdgeWithMappings(fromNodeKey, n.key, false, true); err != nil { + return err + } + n.dependencySetter(fromNodeKey, normalDependency) + return nil + }) + } else { + n.addInputs = append(n.addInputs, func() error { + var paths []FieldPath + for _, input := range inputs { + paths = append(paths, input.targetPath()) + } + if err := n.checkAndAddMappedPath(paths); err != nil { + return err + } + + if err := n.g.addEdgeWithMappings(fromNodeKey, n.key, false, false, inputs...); err != nil { + return err + } + n.dependencySetter(fromNodeKey, normalDependency) + return nil + }) + } + + return n +} + +func (n *WorkflowNode) checkAndAddMappedPath(paths []FieldPath) error { + if v, ok := n.mappedFieldPath[""]; ok { + if _, ok = v.(struct{}); ok { + return fmt.Errorf("entire output has already been mapped for node: %s", n.key) + } + } else { + if len(paths) == 0 { + n.mappedFieldPath[""] = struct{}{} + return nil + } else { + n.mappedFieldPath[""] = map[string]any{} + } + } + + for _, targetPath := range paths { + m := n.mappedFieldPath[""].(map[string]any) + var traversed FieldPath + for i, path := range targetPath { + traversed = append(traversed, path) + if v, ok := m[path]; ok { + if _, ok = v.(struct{}); ok { + return fmt.Errorf("two terminal field paths conflict for node %s: %v, %v", n.key, traversed, targetPath) + } + } + + if i < len(targetPath)-1 { + m[path] = make(map[string]any) + m = m[path].(map[string]any) + } else { + m[path] = struct{}{} + } + } + } + + return nil +} + +// WorkflowBranch represents a branch added to a workflow. +// Each branch may define its own end nodes and mappings. +type WorkflowBranch struct { + fromNodeKey string + *GraphBranch +} + +// AddBranch adds a branch to the workflow. +// +// End Nodes Field Mappings: +// End nodes of the branch are required to define their own field mappings. +// This is a key distinction between Graph's Branch and Workflow's Branch: +// - Graph's Branch: Automatically passes its input to the selected node. +// - Workflow's Branch: Does not pass its input to the selected node. +func (wf *Workflow[I, O]) AddBranch(fromNodeKey string, branch *GraphBranch) *WorkflowBranch { + wb := &WorkflowBranch{ + fromNodeKey: fromNodeKey, + GraphBranch: branch, + } + + wf.workflowBranches = append(wf.workflowBranches, wb) + return wb +} + +// AddEnd connects a node to END with optional field mappings. +// Deprecated: use *Workflow[I,O].End() to obtain a WorkflowNode instance for END, then work with it just like a normal WorkflowNode. +func (wf *Workflow[I, O]) AddEnd(fromNodeKey string, inputs ...*FieldMapping) *Workflow[I, O] { + for _, input := range inputs { + input.fromNodeKey = fromNodeKey + } + _ = wf.g.addEdgeWithMappings(fromNodeKey, END, false, false, inputs...) + return wf +} + +func (wf *Workflow[I, O]) compile(ctx context.Context, options *graphCompileOptions) (*composableRunnable, error) { + if wf.g.buildError != nil { + return nil, wf.g.buildError + } + + for _, wb := range wf.workflowBranches { + for endNode := range wb.endNodes { + if endNode == END { + if _, ok := wf.dependencies[END]; !ok { + wf.dependencies[END] = make(map[string]dependencyType) + } + wf.dependencies[END][wb.fromNodeKey] = branchDependency + } else { + n := wf.workflowNodes[endNode] + n.dependencySetter(wb.fromNodeKey, branchDependency) + } + } + _ = wf.g.addBranch(wb.fromNodeKey, wb.GraphBranch, true) + } + + for _, n := range wf.workflowNodes { + for _, addInput := range n.addInputs { + if err := addInput(); err != nil { + return nil, err + } + } + n.addInputs = nil + } + + for _, n := range wf.workflowNodes { + if len(n.staticValues) > 0 { + value := make(map[string]any, len(n.staticValues)) + var paths []FieldPath + for path, v := range n.staticValues { + value[path] = v + paths = append(paths, splitFieldPath(path)) + } + + if err := n.checkAndAddMappedPath(paths); err != nil { + return nil, err + } + + pair := handlerPair{ + invoke: func(in any) (any, error) { + values := []any{in, value} + return mergeValues(values, nil) + }, + transform: func(in streamReader) streamReader { + sr := schema.StreamReaderFromArray([]map[string]any{value}) + newS, err := mergeValues([]any{in, packStreamReader(sr)}, nil) + if err != nil { + errSR, errSW := schema.Pipe[map[string]any](1) + errSW.Send(nil, err) + errSW.Close() + return packStreamReader(errSR) + } + + return newS.(streamReader) + }, + } + + for i := range paths { + wf.g.fieldMappingRecords[n.key] = append(wf.g.fieldMappingRecords[n.key], ToFieldPath(paths[i])) + } + + wf.g.handlerPreNode[n.key] = []handlerPair{pair} + } + } + + // TODO: check indirect edges are legal + + return wf.g.compile(ctx, options) +} + +func (wf *Workflow[I, O]) initNode(key string) *WorkflowNode { + n := &WorkflowNode{ + g: wf.g, + key: key, + staticValues: make(map[string]any), + dependencySetter: func(fromNodeKey string, typ dependencyType) { + if _, ok := wf.dependencies[key]; !ok { + wf.dependencies[key] = make(map[string]dependencyType) + } + wf.dependencies[key][fromNodeKey] = typ + }, + mappedFieldPath: make(map[string]any), + } + wf.workflowNodes[key] = n + return n +} + +func (wf *Workflow[I, O]) getGenericHelper() *genericHelper { + return wf.g.getGenericHelper() +} + +func (wf *Workflow[I, O]) inputType() reflect.Type { + return wf.g.inputType() +} + +func (wf *Workflow[I, O]) outputType() reflect.Type { + return wf.g.outputType() +} + +func (wf *Workflow[I, O]) component() component { + return wf.g.component() +} diff --git a/compose/workflow_test.go b/compose/workflow_test.go new file mode 100644 index 0000000..d0d4293 --- /dev/null +++ b/compose/workflow_test.go @@ -0,0 +1,1603 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package compose + +import ( + "context" + "errors" + "fmt" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/internal/mock/components/embedding" + "github.com/cloudwego/eino/internal/mock/components/indexer" + "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +func TestWorkflow(t *testing.T) { + ctx := context.Background() + + type structA struct { + Field1 string + Field2 int + Field3 []any + } + + type structB struct { + Field1 string + Field2 int + } + + type structC struct { + Field1 string + } + + type structE struct { + Field1 string + Field2 string + Field3 []any + } + + type structF struct { + Field1 string + Field2 string + Field3 []any + B int + StateTemp string + } + RegisterStreamChunkConcatFunc(func(ts []*structF) (*structF, error) { + ret := &structF{} + for _, tt := range ts { + ret.Field1 += tt.Field1 + ret.Field2 += tt.Field2 + ret.Field3 = append(ret.Field3, tt.Field3...) + ret.B += tt.B + ret.StateTemp += tt.StateTemp + } + return ret, nil + }) + + type state struct { + temp string + } + + type structEnd struct { + Field1 string + } + + subGraph := NewGraph[string, *structB]() + _ = subGraph.AddLambdaNode( + "1", + InvokableLambda(func(ctx context.Context, input string) (*structB, error) { + return &structB{Field1: input, Field2: 33}, nil + }), + ) + _ = subGraph.AddEdge(START, "1") + _ = subGraph.AddEdge("1", END) + + subChain := NewChain[any, *structC](). + AppendLambda(InvokableLambda(func(_ context.Context, in any) (*structC, error) { + return &structC{Field1: fmt.Sprintf("%d", in)}, nil + })) + + type struct2 struct { + F map[string]any + } + subWorkflow := NewWorkflow[[]any, []any]() + subWorkflow.AddLambdaNode( + "1", + InvokableLambda(func(_ context.Context, in []any) ([]any, error) { + return in, nil + }), + WithOutputKey("key")). + AddInput(START) // []any -> map["key"][]any + subWorkflow.AddLambdaNode( + "2", + InvokableLambda(func(_ context.Context, in []any) ([]any, error) { + return in, nil + }), + WithInputKey("key"), + WithOutputKey("key1")). + AddInput("1") // map["key"][]any -> []any -> map["key1"][]any + subWorkflow.AddLambdaNode( + "3", + InvokableLambda(func(_ context.Context, in struct2) (map[string]any, error) { + return in.F, nil + }), + ). + AddInput("2", ToField("F")) // map["key1"][]any -> map["F"]map["key1"][]any -> struct2{F: map["key1"]any} -> map["key1"][]any + subWorkflow.AddLambdaNode( + "4", + InvokableLambda(func(_ context.Context, in []any) ([]any, error) { + return in, nil + }), + WithInputKey("key1"), + ). + AddInput("3") // map["key1"][]any -> []any + subWorkflow.End().AddInput("4") + + w := NewWorkflow[*structA, *structEnd](WithGenLocalState(func(context.Context) *state { return &state{} })) + + w. + AddGraphNode("B", subGraph, + WithStatePostHandler(func(ctx context.Context, out *structB, state *state) (*structB, error) { + state.temp = out.Field1 + return out, nil + })). + AddInput(START, FromField("Field1")) + + w. + AddGraphNode("C", subChain). + AddInput(START, FromField("Field2")) + + w. + AddGraphNode("D", subWorkflow). + AddInput(START, FromField("Field3")) + + w. + AddLambdaNode( + "E", + TransformableLambda(func(_ context.Context, in *schema.StreamReader[structE]) (*schema.StreamReader[structE], error) { + return schema.StreamReaderWithConvert(in, func(in structE) (structE, error) { + if len(in.Field1) > 0 { + in.Field1 = "E:" + in.Field1 + } + if len(in.Field2) > 0 { + in.Field2 = "E:" + in.Field2 + } + + return in, nil + }), nil + }), + WithStreamStatePreHandler(func(ctx context.Context, in *schema.StreamReader[structE], state *state) (*schema.StreamReader[structE], error) { + temp := state.temp + return schema.StreamReaderWithConvert(in, func(v structE) (structE, error) { + if len(v.Field3) > 0 { + v.Field3 = append(v.Field3, "Pre:"+temp) + } + + return v, nil + }), nil + }), + WithStreamStatePostHandler(func(ctx context.Context, out *schema.StreamReader[structE], state *state) (*schema.StreamReader[structE], error) { + return schema.StreamReaderWithConvert(out, func(v structE) (structE, error) { + if len(v.Field1) > 0 { + v.Field1 = v.Field1 + "+Post" + } + return v, nil + }), nil + })). + AddInput("B", MapFields("Field1", "Field1")). + AddInput("C", MapFields("Field1", "Field2")). + AddInput("D", ToField("Field3")) + + w. + AddLambdaNode( + "F", + InvokableLambda(func(ctx context.Context, in *structF) (string, error) { + return fmt.Sprintf("%v_%v_%v_%v_%v", in.Field1, in.Field2, in.Field3, in.B, in.StateTemp), nil + }), + WithStatePreHandler(func(ctx context.Context, in *structF, state *state) (*structF, error) { + in.StateTemp = state.temp + return in, nil + }), + ). + AddInput("B", MapFields("Field2", "B")). + AddInput("E", + MapFields("Field1", "Field1"), + MapFields("Field2", "Field2"), + MapFields("Field3", "Field3"), + ) + + w.End().AddInput("F", ToField("Field1")) + + compiled, err := w.Compile(ctx) + assert.NoError(t, err) + + input := &structA{ + Field1: "1", + Field2: 2, + Field3: []any{ + 1, "good", + }, + } + out, err := compiled.Invoke(ctx, input) + assert.NoError(t, err) + assert.Equal(t, &structEnd{"E:1+Post_E:2_[1 good Pre:1]_33_1"}, out) + + outStream, err := compiled.Stream(ctx, input) + assert.NoError(t, err) + defer outStream.Close() + for { + chunk, err := outStream.Recv() + if err != nil { + if err == io.EOF { + break + } + + t.Error(err) + return + } + + assert.Equal(t, &structEnd{"E:1+Post_E:2_[1 good Pre:1]_33_1"}, chunk) + } +} + +func TestWorkflowWithMap(t *testing.T) { + ctx := context.Background() + + type structA struct { + F1 any + } + + wf := NewWorkflow[map[string]any, map[string]any]() + wf.AddLambdaNode("lambda1", InvokableLambda(func(ctx context.Context, in map[string]any) (map[string]any, error) { + return in, nil + })).AddInput(START, MapFields("map_key", "lambda1_key")) + wf.AddLambdaNode("lambda2", InvokableLambda(func(ctx context.Context, in *structA) (*structA, error) { + return in, nil + })).AddInput(START, MapFields("map_key", "F1")) + wf.End().AddInput("lambda1", MapFields("lambda1_key", "end_lambda1")) + wf.End().AddInput("lambda2", MapFields("F1", "end_lambda2")) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[string]any{"map_key": "value"}) + assert.NoError(t, err) + assert.Equal(t, map[string]any{"end_lambda1": "value", "end_lambda2": "value"}, out) +} + +func TestWorkflowWithNestedFieldMappings(t *testing.T) { + ctx := context.Background() + + type structA struct { + F1 string + } + + type structB struct { + F1 *structA + F2 map[string]any + F3 int + F4 any + F5 map[string]structA + F6 structA + } + + t.Run("from struct.struct.field", func(t *testing.T) { + wf := NewWorkflow[*structB, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, &structB{ + F1: &structA{ + F1: "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + + wf = NewWorkflow[*structB, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F1", "F2"})) + _, err = wf.Compile(ctx) + assert.ErrorContains(t, err, "has no field[F2]") + }) + + t.Run("to struct.(non-ptr)struct.field", func(t *testing.T) { + wf := NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F6", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, &structB{ + F6: structA{ + F1: "hello", + }, + }, out) + }) + + t.Run("to map.(non-ptr)struct.field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]structA]() + wf.End().AddInput(START, ToFieldPath([]string{"key", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]structA{ + "key": { + F1: "hello", + }, + }, out) + }) + + t.Run("from map.map.field", func(t *testing.T) { + wf := NewWorkflow[map[string]map[string]string, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[string]map[string]string{ + "F1": { + "F1": "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + + _, err = r.Invoke(ctx, map[string]map[string]string{ + "F1": { + "F2": "hello", + }, + }) + + var ie *internalError + assert.True(t, errors.As(err, &ie)) + var myErr *errMapKeyNotFound + assert.True(t, errors.As(ie.origError, &myErr)) + }) + + t.Run("from struct.map.field", func(t *testing.T) { + wf := NewWorkflow[*structB, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F2", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, &structB{ + F2: map[string]any{ + "F1": "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + + _, err = r.Invoke(ctx, &structB{ + F2: map[string]any{ + "F2": "hello", + }, + }) + var ie *internalError + assert.True(t, errors.As(err, &ie)) + var myErr *errMapKeyNotFound + assert.True(t, errors.As(ie.origError, &myErr)) + }) + + t.Run("from map.struct.field", func(t *testing.T) { + wf := NewWorkflow[map[string]*structA, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[string]*structA{ + "F1": { + F1: "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + + wf = NewWorkflow[map[string]*structA, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F1", "F2"})) + _, err = wf.Compile(ctx) + assert.ErrorContains(t, err, "has no field[F2]") + }) + + t.Run("from map[string]any.field", func(t *testing.T) { + wf := NewWorkflow[map[string]any, string]() + wf.End().AddInput(START, FromFieldPath([]string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[string]any{ + "F1": &structA{ + F1: "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + + out, err = r.Invoke(ctx, map[string]any{ + "F1": map[string]any{ + "F1": "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + + _, err = r.Invoke(ctx, map[string]any{ + "F1": 1, + }) + + var ie *internalError + assert.True(t, errors.As(err, &ie)) + var myErr *errInterfaceNotValidForFieldMapping + assert.True(t, errors.As(ie.origError, &myErr)) + }) + + t.Run("to struct.struct.field", func(t *testing.T) { + wf := NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, &structB{ + F1: &structA{ + F1: "hello", + }, + }, out) + + wf = NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F1", "F2"})) + _, err = wf.Compile(ctx) + assert.ErrorContains(t, err, "has no field[F2]") + }) + + t.Run("to map.map.field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]map[string]string]() + wf.End().AddInput(START, ToFieldPath([]string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]map[string]string{ + "F1": { + "F1": "hello", + }, + }, out) + + wf1 := NewWorkflow[string, map[string]map[string]int]() + wf1.End().AddInput(START, ToFieldPath([]string{"F1", "F1"})) + _, err = wf1.Compile(ctx) + assert.ErrorContains(t, err, "field[string]-[int] is absolutely not assignable") + }) + + t.Run("to struct.map.field", func(t *testing.T) { + wf := NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F2", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, &structB{ + F2: map[string]any{ + "F1": "hello", + }, + }, out) + }) + + t.Run("to map.struct.struct.field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]*structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F1", "F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]*structB{ + "F1": { + F1: &structA{ + F1: "hello", + }, + }, + }, out) + }) + + t.Run("to struct.map.struct.field", func(t *testing.T) { + wf := NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F5", "key", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, &structB{ + F5: map[string]structA{ + "key": { + F1: "hello", + }, + }, + }, out) + }) + + t.Run("to map.map.struct(non-ptr).field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]map[string]structA]() + wf.End().AddInput(START, ToFieldPath([]string{"key1", "key2", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]map[string]structA{ + "key1": { + "key2": { + F1: "hello", + }, + }, + }, out) + }) + + t.Run("to struct.int.field", func(t *testing.T) { + wf := NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F3", "F1", "F1"})) + _, err := wf.Compile(ctx) + assert.ErrorContains(t, err, "type[int] is not valid") + }) + + t.Run("to struct.any.field", func(t *testing.T) { + wf := NewWorkflow[string, *structB]() + wf.End().AddInput(START, ToFieldPath([]string{"F4", "F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, &structB{ + F4: map[string]any{ + "F1": map[string]any{ + "F1": "hello", + }, + }, + }, out) + }) + + t.Run("to map.any.any.field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.End().AddInput(START, ToFieldPath([]string{"Key1", "Key2", "Key3"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "Key1": map[string]any{ + "Key2": map[string]any{ + "Key3": "hello", + }, + }, + }, out) + }) + + t.Run("to any", func(t *testing.T) { + wf := NewWorkflow[string, any]() + wf.End().AddInput(START) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, "hello", out) + }) + + t.Run("to any.field", func(t *testing.T) { + wf := NewWorkflow[string, any]() + wf.End().AddInput(START, ToFieldPath([]string{"Key1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "Key1": "hello", + }, out) + }) + + t.Run("to interface.field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]fmt.Stringer]() + wf.End().AddInput(START, ToFieldPath([]string{"Key1", "A"})) + _, err := wf.Compile(ctx) + assert.ErrorContains(t, err, "static check failed for mapping [from start to Key1\u001FA(field)], "+ + "the successor has intermediate interface type fmt.Stringer") + }) + + t.Run("both to map.any, and to map.any.field", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.End().AddInput(START, ToFieldPath([]string{"Key1"}), ToFieldPath([]string{"Key1", "Key2"})) + _, err := wf.Compile(ctx) + assert.ErrorContains(t, err, "two terminal field paths conflict") + }) + + t.Run("to map.any.any.field1, and to map.any.any.field2", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.End().AddInput(START, ToFieldPath([]string{"Key1", "Key2", "key3"}), ToFieldPath([]string{"Key1", "Key2", "key4"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "Key1": map[string]any{ + "Key2": map[string]any{ + "key3": "hello", + "key4": "hello", + }, + }, + }, out) + }) + + t.Run("from nested to nested", func(t *testing.T) { + wf := NewWorkflow[map[string]any, *structB]() + wf.End().AddInput(START, MapFieldPaths([]string{"key1", "key2"}, []string{"F1", "F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[string]any{ + "key1": map[string]any{ + "key2": "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, &structB{ + F1: &structA{ + F1: "hello", + }, + }, out) + }) + + t.Run("from nested to normal", func(t *testing.T) { + wf := NewWorkflow[map[string]any, *structA]() + wf.End().AddInput(START, MapFieldPaths(FieldPath{"key1", "key2"}, FieldPath{"F1"})) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[string]any{ + "key1": map[string]any{ + "key2": "hello", + }, + }) + assert.NoError(t, err) + assert.Equal(t, &structA{ + F1: "hello", + }, out) + }) +} + +func TestWorkflowCompile(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + + t.Run("compile without add end", func(t *testing.T) { + w := NewWorkflow[*schema.Message, []*schema.Message]() + w.AddToolsNode("1", &ToolsNode{}).AddInput(START) + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "end node not set") + }) + + t.Run("type mismatch", func(t *testing.T) { + w := NewWorkflow[string, string]() + w.AddToolsNode("1", &ToolsNode{}).AddInput(START) + w.End().AddInput("1") + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, " mismatch") + }) + + t.Run("predecessor's output not struct/struct ptr/map, mapping has FromField", func(t *testing.T) { + w := NewWorkflow[[]*schema.Document, []string]() + + w.AddIndexerNode("indexer", indexer.NewMockIndexer(ctrl)).AddInput(START, FromField("F1")) + w.End().AddInput("indexer") + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "predecessor output type should be struct") + }) + + t.Run("successor's input not struct/struct ptr/map, mapping has ToField", func(t *testing.T) { + w := NewWorkflow[[]string, [][]float64]() + w.AddEmbeddingNode("embedder", embedding.NewMockEmbedder(ctrl)).AddInput(START, ToField("F1")) + w.End().AddInput("embedder") + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "successor input type should be struct") + }) + + t.Run("map to non existing field in predecessor", func(t *testing.T) { + w := NewWorkflow[*schema.Message, []*schema.Message]() + w.AddToolsNode("tools_node", &ToolsNode{}).AddInput(START, FromField("non_exist")) + w.End().AddInput("tools_node") + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "type[schema.Message] has no field[non_exist]") + }) + + t.Run("map to not exported field in successor", func(t *testing.T) { + w := NewWorkflow[string, *FieldMapping]() + w.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })).AddInput(START) + w.End().AddInput("1", ToField("to")) + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "has an unexported field[to]") + }) + + t.Run("map from not exported field in predecessor", func(t *testing.T) { + w := NewWorkflow[*FieldMapping, string]() + w.End().AddInput(START, FromField("from")) + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "has an unexported field[from]") + }) + + t.Run("duplicate node key", func(t *testing.T) { + w := NewWorkflow[[]*schema.Message, []*schema.Message]() + w.AddChatModelNode("1", model.NewMockChatModel(ctrl)).AddInput(START) + w.AddToolsNode("1", &ToolsNode{}).AddInput("1") + w.End().AddInput("1") + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "node '1' already present") + }) + + t.Run("from non-existing node", func(t *testing.T) { + w := NewWorkflow[*schema.Message, []*schema.Message]() + w.AddToolsNode("1", &ToolsNode{}).AddInput(START) + w.End().AddInput("2") + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "edge start node '2' needs to be added to graph first") + }) + + t.Run("to map with non-string key type", func(t *testing.T) { + w := NewWorkflow[string, map[int]any]() + w.End().AddInput(START, ToField("1")) + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "type[map[int]interface {}] is not a map with string or string alias key") + + type stringAlias string + w1 := NewWorkflow[string, map[stringAlias]any]() + w1.End().AddInput(START, ToField("1")) + r, err := w1.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[stringAlias]any{ + "1": "hello", + }, out) + }) + + t.Run("from map with non-string key type", func(t *testing.T) { + w := NewWorkflow[map[int]any, string]() + w.End().AddInput(START, FromField("1")) + _, err := w.Compile(ctx) + assert.ErrorContains(t, err, "type[map[int]interface {}] is not a map with string or string alias key") + type stringAlias string + w1 := NewWorkflow[map[stringAlias]any, string]() + w1.End().AddInput(START, FromField("1")) + r, err := w1.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, map[stringAlias]any{ + "1": "hello", + }) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + }) +} + +func TestFanInToSameDest(t *testing.T) { + t.Run("traditional outputKey fan-in with map[string]any", func(t *testing.T) { + wf := NewWorkflow[string, []*schema.Message]() + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in, nil + }), WithOutputKey("q1")).AddInput(START) + wf.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in, nil + }), WithOutputKey("q2")).AddInput(START) + wf.AddChatTemplateNode("prompt", prompt.FromMessages(schema.Jinja2, schema.UserMessage("{{q1}}_{{q2}}"))). + AddInput("1", MapFields("q1", "q1")). + AddInput("2", MapFields("q2", "q2")) + wf.End().AddInput("prompt") + c, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := c.Invoke(context.Background(), "query") + assert.NoError(t, err) + assert.Equal(t, []*schema.Message{{Role: schema.User, Content: "query_query"}}, out) + }) + + t.Run("fan-in to a field of map", func(t *testing.T) { + type dest struct { + F map[string]any + } + + type in struct { + A string + B int + } + + wf := NewWorkflow[in, dest]() + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in, nil + }), WithOutputKey("A")).AddInput(START, FromField("A")) + wf.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, in int) (output int, err error) { + return in, nil + }), WithOutputKey("B")).AddInput(START, FromField("B")) + wf.End().AddInput("1", ToField("F")).AddInput("2", ToField("F")) + _, err := wf.Compile(context.Background()) + assert.ErrorContains(t, err, "two terminal field paths conflict for node end: [F], [F]") + }) +} + +func TestIndirectEdge(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_" + in, nil + })).AddInput(START) + + wf.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, in map[string]string) (output string, err error) { + return in["1"] + "_" + in[START], nil + })).AddInput("1", ToField("1")). + AddInputWithOptions(START, []*FieldMapping{ToField(START)}, WithNoDirectDependency()) + + wf.End().AddInput("2", ToField("2")). + AddInputWithOptions("1", []*FieldMapping{ToField("1")}, WithNoDirectDependency()) + + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := r.Invoke(context.Background(), "query") + assert.NoError(t, err) + assert.Equal(t, map[string]any{"1": "query_query", "2": "query_query_query"}, out) +} + +func TestDependencyWithNoInput(t *testing.T) { + t.Run("simple case", func(t *testing.T) { + wf := NewWorkflow[string, string]() + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return "useless", nil + })).AddInput(START) + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_done", nil + })).AddDependency("0").AddInputWithOptions(START, nil, WithNoDirectDependency()) + wf.End().AddInput("1") + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := r.Invoke(context.Background(), "hello") + assert.NoError(t, err) + assert.Equal(t, "hello_done", out) + }) + + t.Run("simple control flow: [Start] --> [Node '0'] --> [End]", func(t *testing.T) { + // [Start] --> [Node "0"] --> [End] + wf := NewWorkflow[map[string]any, map[string]any]() + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) { + return map[string]any{ + "result": "result from node 0", + }, nil + })).AddDependency(START) + wf.End().AddInput("0", ToField("final_result")). + AddInputWithOptions(START, []*FieldMapping{ToField("final_from_start")}, WithNoDirectDependency()) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + ret, err := r.Invoke(context.Background(), map[string]any{ + "input": "hello", + }) + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "final_result": map[string]any{ + "result": "result from node 0", + }, + "final_from_start": map[string]any{ + "input": "hello", + }, + }, ret) + + sRet, err := r.Stream(context.Background(), map[string]any{ + "input": "hello", + }) + assert.NoError(t, err) + ret, err = concatStreamReader(sRet) + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "final_result": map[string]any{ + "result": "result from node 0", + }, + "final_from_start": map[string]any{ + "input": "hello", + }, + }, ret) + }) +} + +func TestStaticValue(t *testing.T) { + t.Run("prefill map", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) { + return in, nil + })). + AddInput(START, ToField(START)). + SetStaticValue(FieldPath{"prefilled"}, "yo-ho") + wf.End().AddInput("0") + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := r.Invoke(context.Background(), "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{"prefilled": "yo-ho", START: "hello"}, out) + streamOut, err := r.Stream(context.Background(), "hello") + assert.NoError(t, err) + out = map[string]any{} + for { + chunk, err := streamOut.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + for k, v := range chunk { + out[k] = v + } + } + assert.Equal(t, map[string]any{"prefilled": "yo-ho", START: "hello"}, out) + }) + + t.Run("static value and to-all mapping conflict", func(t *testing.T) { + wf := NewWorkflow[map[string]any, map[string]any]() + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) { + return in, nil + })). + AddInput(START). + SetStaticValue( + FieldPath{"prefilled"}, + "yo-ho", + ) + wf.End().AddInput("0") + _, err := wf.Compile(context.Background()) + assert.ErrorContains(t, err, "entire output has already been mapped for node: 0") + }) + + t.Run("static value and dynamic mapping conflict", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) { + return in, nil + })). + AddInput(START, ToField("prefilled")). + SetStaticValue(FieldPath{"prefilled"}, "yo-ho") + wf.End().AddInput("0") + _, err := wf.Compile(context.Background()) + assert.ErrorContains(t, err, "two terminal field paths conflict for node 0: [prefilled], [prefilled]") + }) + + t.Run("all inputs are static values", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) { + return in, nil + })). + AddDependency(START). + SetStaticValue(FieldPath{"a", "b"}, "a_b"). + SetStaticValue(FieldPath{"c", "d"}, "c_d"). + SetStaticValue(FieldPath{"a", "d"}, "a_d") + wf.End().AddInput("0") + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := r.Invoke(context.Background(), "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "a": map[string]any{ + "b": "a_b", + "d": "a_d", + }, + "c": map[string]any{ + "d": "c_d", + }, + }, out) + + type a struct { + B string + D string + } + + type s struct { + A a + C map[string]any + } + + wf1 := NewWorkflow[string, *s]() + wf1.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in map[string]any) (output map[string]any, err error) { + return in, nil + })). + AddDependency(START). + SetStaticValue(FieldPath{"A", "B"}, "a_b"). + SetStaticValue(FieldPath{"C", "D"}, "c_d"). + SetStaticValue(FieldPath{"A", "D"}, "a_d") + wf1.End().AddInput("0", MapFieldPaths(FieldPath{"A", "B"}, FieldPath{"A", "B"}), + MapFieldPaths(FieldPath{"A", "D"}, FieldPath{"A", "D"}), + MapFields("C", "C")) + r1, err := wf1.Compile(context.Background()) + assert.NoError(t, err) + out1, err := r1.Stream(context.Background(), "hello") + assert.NoError(t, err) + outChunk, err := out1.Recv() + out1.Close() + assert.Equal(t, &s{ + A: a{ + B: "a_b", + D: "a_d", + }, + C: map[string]any{ + "D": "c_d", + }, + }, outChunk) + }) +} + +func TestBranch(t *testing.T) { + ctx := context.Background() + t.Run("simple branch: one predecessor, two successor, one of them is END", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_" + in, nil + })).AddInputWithOptions(START, nil, WithNoDirectDependency()) + + wf.AddPassthroughNode("branch_1").AddInput(START, ToField(START)) + + branch := NewGraphBranch(func(ctx context.Context, in map[string]any) (string, error) { + if in[START] == "hello" { + return "1", nil + } + return END, nil + }, map[string]bool{ + "1": true, + END: true, + }) + wf.AddBranch("branch_1", branch) + wf.End().AddInput("1", ToField("1")).AddInputWithOptions(START, []*FieldMapping{ToField(START)}, WithNoDirectDependency()) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": "hello_hello", + START: "hello", + }, out) + out, err = r.Invoke(ctx, "world") + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + START: "world", + }, out) + }) + + t.Run("multiple predecessors", func(t *testing.T) { + wf := NewWorkflow[string, map[string]any]() + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_" + in, nil + })).AddInput(START) + wf.AddLambdaNode("2", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_" + in, nil + })).AddInputWithOptions("1", nil, WithNoDirectDependency()) + wf.AddLambdaNode("0", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_" + in, nil + })).AddInput(START) + + wf.AddPassthroughNode("branch_1").AddInput(START, ToField(START)).AddInput("1", ToField("1")).AddDependency("0") + wf.AddBranch("branch_1", NewGraphBranch(func(ctx context.Context, in map[string]any) (string, error) { + if in[START].(string) == "hello" { + return "2", nil + } + return END, nil + }, map[string]bool{ + "2": true, + END: true, + })) + wf.End().AddInput("2", ToField("2")).AddInputWithOptions(START, []*FieldMapping{ToField(START)}, WithNoDirectDependency()) + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "hello") + assert.NoError(t, err) + assert.Equal(t, map[string]any{"2": "hello_hello_hello_hello", START: "hello"}, out) + out, err = r.Invoke(ctx, "world") + assert.NoError(t, err) + assert.Equal(t, map[string]any{START: "world"}, out) + }) + t.Run("empty input for node after branch", func(t *testing.T) { + wf := NewWorkflow[map[string]any, map[string]any]() + wf.AddLambdaNode("start_1", InvokableLambda(func(ctx context.Context, input map[string]any) (map[string]any, error) { + return map[string]any{}, nil + })).AddInput("start") + wf.AddLambdaNode("branch_1", InvokableLambda(func(ctx context.Context, input map[string]any) (map[string]any, error) { + return map[string]any{}, nil + })) + wf.AddPassthroughNode("my_branch").AddInput("start_1") + wf.AddBranch("my_branch", NewGraphBranch(func(ctx context.Context, input map[string]any) (string, error) { + return END, nil + }, map[string]bool{ + "branch_1": true, + END: true, + })) + wf.End().AddInput("branch_1") + runner, err := wf.Compile(context.Background()) + assert.NoError(t, err) + resp, err := runner.Invoke(context.Background(), map[string]any{}) + assert.NoError(t, err) + assert.Equal(t, resp, (map[string]any)(nil)) + }) +} + +type goodInterface interface { + GOOD() +} +type goodStruct struct{} + +func (g *goodStruct) GOOD() {} + +func TestMayAssignableFieldMapping(t *testing.T) { + type in struct { + A goodInterface + } + wf := NewWorkflow[in, *goodStruct]() + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input *goodStruct) (output goodInterface, err error) { return input, nil })). + AddInput(START, FromField("A")) + wf.End().AddInput("1") + ctx := context.Background() + r, err := wf.Compile(ctx) + assert.NoError(t, err) + result, err := r.Invoke(ctx, in{A: &goodStruct{}}) + assert.NoError(t, err) + result.GOOD() +} + +func TestNilValue(t *testing.T) { + t.Run("from map key with a nil value to map key", func(t *testing.T) { + wf := NewWorkflow[map[string]any, map[string]any]() + wf.End().AddInput(START, MapFields("a", "a")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + + result, err := r.Invoke(context.Background(), map[string]any{"a": nil}) + assert.NoError(t, err) + assert.Equal(t, map[string]any{"a": nil}, result) + }) + + t.Run("from nil struct field to map key", func(t *testing.T) { + type in struct { + A *string + } + wf := NewWorkflow[in, map[string]any]() + wf.End().AddInput(START, MapFields("A", "A")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + result, err := r.Invoke(context.Background(), in{A: nil}) + assert.NoError(t, err) + assert.Equal(t, map[string]any{"A": (*string)(nil)}, result) + }) + + t.Run("from map key with a nil value to struct field", func(t *testing.T) { + type out struct { + A *string + } + wf := NewWorkflow[map[string]any, out]() + wf.End().AddInput(START, MapFields("A", "A")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + result, err := r.Invoke(context.Background(), map[string]any{"A": nil}) + assert.NoError(t, err) + assert.Equal(t, out{A: (*string)(nil)}, result) + }) + + t.Run("from nil struct field to struct field", func(t *testing.T) { + type inOut struct { + A *string + } + wf := NewWorkflow[inOut, inOut]() + wf.End().AddInput(START, MapFields("A", "A")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + result, err := r.Invoke(context.Background(), inOut{A: nil}) + assert.NoError(t, err) + assert.Equal(t, inOut{A: (*string)(nil)}, result) + }) + + t.Run("from nil to a type that can't be nil", func(t *testing.T) { + wf := NewWorkflow[map[string]any, int]() + wf.End().AddInput(START, FromField("a")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]any{"a": nil}) + assert.ErrorContains(t, err, "runtime check failed for mapping [from a(field) of start], field[]-[int] is absolutely not assignable") + }) + + t.Run("from nil to a map other than map[string]any", func(t *testing.T) { + wf := NewWorkflow[map[string]any, map[string]fmt.Stringer]() + wf.End().AddInput(START, MapFields("a", "a")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := r.Invoke(context.Background(), map[string]any{"a": nil}) + assert.Equal(t, map[string]fmt.Stringer{ + "a": nil, + }, out) + }) +} + +func TestStreamFieldMap(t *testing.T) { + t.Run("multiple incomplete chunks in source stream", func(t *testing.T) { + wf := NewWorkflow[map[string]any, map[string]any]() + wf.End().AddInput(START, MapFields("a", "a"), MapFields("b", "b")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + + sr, sw := schema.Pipe[map[string]any](2) + sw.Send(map[string]any{"a": 1}, nil) + sw.Send(map[string]any{"b": 2}, nil) + sw.Close() + outputS, err := r.Transform(context.Background(), sr) + assert.NoError(t, err) + result, err := concatStreamReader(outputS) + assert.NoError(t, err) + assert.Equal(t, map[string]any{"a": 1, "b": 2}, result) + }) +} + +func TestRuntimeTypeCheck(t *testing.T) { + g := NewWorkflow[map[string]any, any]() + + _ = g. + AddLambdaNode("A", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })). + AddInput(START, FromField("A")) + + _ = g.AddLambdaNode("B", InvokableLambda(func(ctx context.Context, input string) (output string, err error) { + return input, nil + })). + AddInput(START, FromField("B")) + + _ = g.AddLambdaNode("MergeA", InvokableLambda(func(ctx context.Context, input map[string]any) (output map[string]any, err error) { + return input, nil + })). + AddInput("A", ToField("a")). + AddInput("B", ToField("b")) + + g.End().AddInput("MergeA") + + ctx := context.Background() + r, err := g.Compile(ctx) + assert.NoError(t, err) + result, err := r.Stream(ctx, map[string]any{"A": "1", "B": "2"}) + assert.NoError(t, err) + chunk, err := result.Recv() + assert.NoError(t, err) + assert.Equal(t, map[string]any{"a": "1", "b": "2"}, chunk) + chunk, err = result.Recv() + assert.True(t, errors.Is(err, io.EOF)) +} + +func TestIntermediateMappingSource(t *testing.T) { + t.Run("intermediate any source is nil", func(t *testing.T) { + wf := NewWorkflow[map[string]any, any]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a", "b"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]any{ + "a": nil, + }) + assert.ErrorContains(t, err, "intermediate source value on path=[a b] is nil for type [interface {}]") + outStream, err := r.Transform(context.Background(), schema.StreamReaderFromArray([]map[string]any{ + { + "a": nil, + }, + { + "b": "ok", + }, + })) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "intermediate source value on path=[a b] is nil for type [interface {}]") + outStream.Close() + }) + + t.Run("intermediate map source is nil", func(t *testing.T) { + wf := NewWorkflow[map[string]any, any]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), nil) + assert.ErrorContains(t, err, "intermediate source value on path=[a] is nil for map type [map[string]interface {}]") + outStream, err := r.Stream(context.Background(), nil) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "intermediate source value on path=[a] is nil for map type [map[string]interface {}]") + outStream.Close() + }) + + t.Run("intermediate map ptr source is nil", func(t *testing.T) { + wf := NewWorkflow[*map[string]any, any]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), nil) + assert.ErrorContains(t, err, "intermediate source value on path=[a] is nil for type [*map[string]interface {}]") + outStream, err := r.Stream(context.Background(), nil) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "intermediate source value on path=[a] is nil for type [*map[string]interface {}]") + outStream.Close() + }) + + t.Run("intermediate struct ptr source is nil", func(t *testing.T) { + type inner struct { + A string + } + + wf := NewWorkflow[map[string]*inner, string]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"I", "A"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]*inner{"I": nil}) + assert.ErrorContains(t, err, "intermediate source value on path=[I A] is nil") + outStream, err := r.Stream(context.Background(), map[string]*inner{"I": nil}) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "intermediate source value on path=[I A] is nil") + outStream.Close() + }) + + t.Run("intermediate interface source is nil", func(t *testing.T) { + wf := NewWorkflow[map[string]fmt.Stringer, string]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a", "b"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]fmt.Stringer{"a": nil}) + assert.ErrorContains(t, err, "intermediate source value on path=[a b] is nil for type [fmt.Stringer]") + outStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{"a": nil}) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "intermediate source value on path=[a b] is nil for type [fmt.Stringer]") + outStream.Close() + }) + + t.Run("intermediate interface source valid", func(t *testing.T) { + wf := NewWorkflow[map[string]fmt.Stringer, string]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a", "A"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + out, err := r.Invoke(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello"}}) + assert.NoError(t, err) + assert.Equal(t, "hello", out) + outStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello"}}) + assert.NoError(t, err) + out, err = outStream.Recv() + assert.NoError(t, err) + assert.Equal(t, "hello", out) + outStream.Close() + }) + + t.Run("intermediate interface source, source field not found at request time", func(t *testing.T) { + wf := NewWorkflow[map[string]fmt.Stringer, string]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a", "B"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello"}}) + assert.ErrorContains(t, err, "field mapping from a struct field, but field not found. field=B") + outStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello"}}) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "field mapping from a struct field, but field not found. field=B") + outStream.Close() + }) + + t.Run("intermediate interface source, source field not exported at request time", func(t *testing.T) { + wf := NewWorkflow[map[string]fmt.Stringer, string]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a", "c"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello", c: "c"}}) + assert.ErrorContains(t, err, "field mapping from a struct field, but field not exported.") + outStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello", c: "c"}}) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "field mapping from a struct field, but field not exported.") + outStream.Close() + }) + + t.Run("intermediate interface source, type mismatch at request time", func(t *testing.T) { + wf := NewWorkflow[map[string]fmt.Stringer, int]() + wf.End().AddInput(START, FromFieldPath(FieldPath{"a", "A"})) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + _, err = r.Invoke(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello"}}) + assert.ErrorContains(t, err, "runtime check failed for mapping [from a\x1fA(field) of start], field[string]-[int] is absolutely not assignable") + outStream, err := r.Stream(context.Background(), map[string]fmt.Stringer{"a": &goodStruct2{A: "hello"}}) + assert.NoError(t, err) + _, err = outStream.Recv() + assert.ErrorContains(t, err, "runtime check failed for mapping [from a\u001FA(field) of start], field[string]-[int] is absolutely not assignable") + outStream.Close() + }) +} + +type goodStruct2 struct { + A string + c string +} + +func (g *goodStruct2) String() string { + return g.A +} + +func TestSetFanInMergeConfig_RealStreamNode_Workflow(t *testing.T) { + wf := NewWorkflow[int, map[string]int]() + + wf.AddLambdaNode("s1", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) { + sr, sw := schema.Pipe[int](2) + sw.Send(input+1, nil) + sw.Send(input+2, nil) + sw.Close() + return sr, nil + })).AddInput(START) + + wf.AddLambdaNode("s2", StreamableLambda(func(ctx context.Context, input int) (*schema.StreamReader[int], error) { + sr, sw := schema.Pipe[int](2) + sw.Send(input+10, nil) + sw.Send(input+20, nil) + sw.Close() + return sr, nil + })).AddInput(START) + + wf.End().AddInput("s1", ToField("s1")).AddInput("s2", ToField("s2")) + + r, err := wf.Compile(context.Background(), + WithFanInMergeConfig(map[string]FanInMergeConfig{END: {StreamMergeWithSourceEOF: true}})) + assert.NoError(t, err) + + sr, err := r.Stream(context.Background(), 1) + assert.NoError(t, err) + + merged := make(map[string]map[int]bool) + var sourceEOFCount int + sourceNames := make(map[string]bool) + for { + m, e := sr.Recv() + if e != nil { + if name, ok := schema.GetSourceName(e); ok { + sourceEOFCount++ + sourceNames[name] = true + continue + } + if e == io.EOF { + break + } + assert.NoError(t, e) + } + + for k, v := range m { + if merged[k] == nil { + merged[k] = make(map[int]bool) + } + merged[k][v] = true + } + } + + assert.Equal(t, map[string]map[int]bool{"s1": {2: true, 3: true}, "s2": {11: true, 21: true}}, merged) + assert.Equal(t, 2, sourceEOFCount, "should receive SourceEOF for each input stream when StreamMergeWithSourceEOF is true") + assert.True(t, sourceNames["s1"], "should receive SourceEOF from s1") + assert.True(t, sourceNames["s2"], "should receive SourceEOF from s2") +} + +func TestCustomExtractor(t *testing.T) { + t.Run("custom extract from array element", func(t *testing.T) { + wf := NewWorkflow[[]int, map[string]int]() + wf.End().AddInput(START, ToField("a", WithCustomExtractor(func(input any) (any, error) { + return input.([]int)[0], nil + }))) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + result, err := r.Invoke(context.Background(), []int{1, 2}) + assert.NoError(t, err) + assert.Equal(t, map[string]int{"a": 1}, result) + }) + + t.Run("mix custom extract with normal mapping", func(t *testing.T) { + wf := NewWorkflow[map[string]any, map[string]int]() + wf.End().AddInput(START, ToField("a", WithCustomExtractor(func(input any) (any, error) { + return input.(map[string]any)["a"].([]any)[0].(map[string]any)["c"], nil + })), MapFields("b", "b")) + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + result, err := r.Invoke(context.Background(), map[string]any{ + "a": []any{ + map[string]any{ + "c": 1, + }, + }, + "b": 2, + }) + assert.NoError(t, err) + assert.Equal(t, map[string]int{"a": 1, "b": 2}, result) + }) +} + +func TestAddDependency(t *testing.T) { + ctx := context.Background() + + wf := NewWorkflow[string, any]() + + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, in string) (output string, err error) { + return in + "_" + in, nil + })).AddDependency(START) + + wf.End().AddDependency("1") + + r, err := wf.Compile(ctx) + assert.NoError(t, err) + out, err := r.Invoke(ctx, "input") + assert.NoError(t, err) + assert.Equal(t, nil, out) +} + +func TestIndirectDependencyWithBranch(t *testing.T) { + t.Run("data only mapping across branch", func(t *testing.T) { + wf := NewWorkflow[[]int, map[string]any]() + + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input int) (output int, err error) { + return input + 1, nil + })). + AddInputWithOptions(START, []*FieldMapping{ToField("", WithCustomExtractor(func(input any) (any, error) { + inputList := input.([]int) + if len(inputList) == 0 { + return nil, fmt.Errorf("input list is empty") + } + return input.([]int)[0], nil + }))}, WithNoDirectDependency()) + + wf.AddBranch(START, NewGraphBranch(func(ctx context.Context, in []int) (endNode string, err error) { + if len(in) > 0 { + return "1", nil + } + + return END, nil + }, map[string]bool{"1": true, END: true})) + + wf.End(). + AddInput("1", ToField("output")). + SetStaticValue(FieldPath{"static"}, 2) + + r, err := wf.Compile(context.Background()) + assert.NoError(t, err) + + // skip lambda node "1" + out, err := r.Invoke(context.Background(), nil) + assert.NoError(t, err) + assert.Equal(t, out, map[string]any{"static": 2}) + + // choose lambda node "1" + out, err = r.Invoke(context.Background(), []int{1}) + assert.NoError(t, err) + assert.Equal(t, out, map[string]any{"output": 2, "static": 2}) + }) + + t.Run("data only mapping across branch, with interrupt after branch", func(t *testing.T) { + wf := NewWorkflow[[]int, map[string]any]() + + wf.AddLambdaNode("1", InvokableLambda(func(ctx context.Context, input int) (output int, err error) { + return input + 1, nil + })). + AddInputWithOptions(START, []*FieldMapping{ToField("", WithCustomExtractor(func(input any) (any, error) { + inputList := input.([]int) + if len(inputList) == 0 { + return nil, fmt.Errorf("input list is empty") + } + return input.([]int)[0], nil + }))}, WithNoDirectDependency()) + + wf.AddBranch(START, NewGraphBranch(func(ctx context.Context, in []int) (endNode string, err error) { + if len(in) > 0 { + return "1", nil + } + + return END, nil + }, map[string]bool{"1": true, END: true})) + + wf.End(). + AddInput("1", ToField("output")). + SetStaticValue(FieldPath{"static"}, 2) + + r, err := wf.Compile(context.Background(), WithCheckPointStore(newInMemoryStore()), + WithInterruptBeforeNodes([]string{"1"})) + assert.NoError(t, err) + + // skip lambda node "1" + out, err := r.Invoke(context.Background(), nil) + assert.NoError(t, err) + assert.Equal(t, out, map[string]any{"static": 2}) + + // choose lambda node "1" + _, err = r.Invoke(context.Background(), []int{1}, WithCheckPointID("123")) + _, ok := ExtractInterruptInfo(err) + assert.True(t, ok) + + out, err = r.Invoke(context.Background(), nil, WithCheckPointID("123")) + assert.NoError(t, err) + assert.Equal(t, out, map[string]any{"output": 2, "static": 2}) + }) +} diff --git a/doc.go b/doc.go new file mode 100644 index 0000000..9c53751 --- /dev/null +++ b/doc.go @@ -0,0 +1,19 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package eino provides building blocks for agent workflows, +// tools, and composable graph utilities. +package eino diff --git a/flow/agent/agent_option.go b/flow/agent/agent_option.go new file mode 100644 index 0000000..9d16fbe --- /dev/null +++ b/flow/agent/agent_option.go @@ -0,0 +1,71 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package agent defines common option types used by agents and multi-agents. +package agent + +import "github.com/cloudwego/eino/compose" + +// AgentOption is the common option type for various agent and multi-agent implementations. +// For options intended to use with underlying graph or components, use WithComposeOptions to specify. +// For options intended to use with particular agent/multi-agent implementations, use WrapImplSpecificOptFn to specify. +type AgentOption struct { + implSpecificOptFn any + composeOptions []compose.Option +} + +// GetComposeOptions returns all compose options from the given agent options. +func GetComposeOptions(opts ...AgentOption) []compose.Option { + var result []compose.Option + for _, opt := range opts { + result = append(result, opt.composeOptions...) + } + + return result +} + +// WithComposeOptions returns an agent option that specifies compose options. +func WithComposeOptions(opts ...compose.Option) AgentOption { + return AgentOption{ + composeOptions: opts, + } +} + +// WrapImplSpecificOptFn returns an agent option that specifies a function to modify the implementation-specific options. +func WrapImplSpecificOptFn[T any](optFn func(*T)) AgentOption { + return AgentOption{ + implSpecificOptFn: optFn, + } +} + +// GetImplSpecificOptions returns the implementation-specific options from the given agent options. +func GetImplSpecificOptions[T any](base *T, opts ...AgentOption) *T { + if base == nil { + base = new(T) + } + + for i := range opts { + opt := opts[i] + if opt.implSpecificOptFn != nil { + optFn, ok := opt.implSpecificOptFn.(func(*T)) + if ok { + optFn(base) + } + } + } + + return base +} diff --git a/flow/agent/multiagent/host/callback.go b/flow/agent/multiagent/host/callback.go new file mode 100644 index 0000000..e5bcc63 --- /dev/null +++ b/flow/agent/multiagent/host/callback.go @@ -0,0 +1,100 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package host + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/flow/agent" + "github.com/cloudwego/eino/schema" + template "github.com/cloudwego/eino/utils/callbacks" +) + +// MultiAgentCallback is the callback interface for host multi-agent. +type MultiAgentCallback interface { + OnHandOff(ctx context.Context, info *HandOffInfo) context.Context +} + +// HandOffInfo is the info which will be passed to MultiAgentCallback.OnHandOff, representing a hand off event. +type HandOffInfo struct { + ToAgentName string + Argument string +} + +// ConvertCallbackHandlers converts []host.MultiAgentCallback to callbacks.Handler. +func ConvertCallbackHandlers(handlers ...MultiAgentCallback) callbacks.Handler { + onChatModelEnd := func(ctx context.Context, info *callbacks.RunInfo, output *model.CallbackOutput) context.Context { + msg := output.Message + if msg == nil || msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 { + return ctx + } + + for _, cb := range handlers { + for _, toolCall := range msg.ToolCalls { + ctx = cb.OnHandOff(ctx, &HandOffInfo{ + ToAgentName: toolCall.Function.Name, + Argument: toolCall.Function.Arguments, + }) + } + } + + return ctx + } + + onChatModelEndWithStreamOutput := func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context { + go func() { + msg, err := schema.ConcatMessageStream(schema.StreamReaderWithConvert(output, + func(m *model.CallbackOutput) (*schema.Message, error) { + return m.Message, nil + })) + if err != nil { + fmt.Printf("concat message stream for host multi-agent failed: %v", err) + return + } + + for _, cb := range handlers { + for _, tc := range msg.ToolCalls { + _ = cb.OnHandOff(ctx, &HandOffInfo{ + ToAgentName: tc.Function.Name, + Argument: tc.Function.Arguments, + }) + } + } + }() + + return ctx + } + + return template.NewHandlerHelper().ChatModel(&template.ModelCallbackHandler{ + OnEnd: onChatModelEnd, + OnEndWithStreamOutput: onChatModelEndWithStreamOutput, + }).Handler() +} + +// convertCallbacks reads graph call options, extract host.MultiAgentCallback and convert it to callbacks.Handler. +func convertCallbacks(opts ...agent.AgentOption) callbacks.Handler { + agentOptions := agent.GetImplSpecificOptions(&options{}, opts...) + if len(agentOptions.agentCallbacks) == 0 { + return nil + } + + handlers := agentOptions.agentCallbacks + return ConvertCallbackHandlers(handlers...) +} diff --git a/flow/agent/multiagent/host/compose.go b/flow/agent/multiagent/host/compose.go new file mode 100644 index 0000000..6a65cfb --- /dev/null +++ b/flow/agent/multiagent/host/compose.go @@ -0,0 +1,335 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package host + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + "github.com/cloudwego/eino/schema" +) + +const ( + defaultHostNodeKey = "host" // the key of the host node in the graph + defaultHostPrompt = "decide which tool is best for the task and call only the best tool." + specialistsAnswersCollectorNodeKey = "specialist_answers_collect" + singleIntentAnswerNodeKey = "single_intent_answer" + multiIntentSummarizeNodeKey = "multi_intents_summarize" + defaultSummarizerPrompt = "summarize the answers from the specialists into a single answer." + map2ListConverterNodeKey = "map_to_list" +) + +type state struct { + msgs []*schema.Message + isMultipleIntents bool +} + +// NewMultiAgent creates a new host multi-agent system. +// +// IMPORTANT!! For models that don't output tool calls in the first streaming chunk (e.g. Claude) +// the default StreamToolCallChecker may not work properly since it only checks the first chunk for tool calls. +// In such cases, you need to implement a custom StreamToolCallChecker that can properly detect tool calls. +func NewMultiAgent(ctx context.Context, config *MultiAgentConfig) (*MultiAgent, error) { + if err := config.validate(); err != nil { + return nil, err + } + + hostKeyName := defaultHostNodeKey + if config.HostNodeName != "" { + hostKeyName = config.HostNodeName + } + + var ( + hostPrompt = config.Host.SystemPrompt + name = config.Name + toolCallChecker = config.StreamToolCallChecker + ) + + if len(hostPrompt) == 0 { + hostPrompt = defaultHostPrompt + } + + if len(name) == 0 { + name = "host multi agent" + } + + if toolCallChecker == nil { + toolCallChecker = firstChunkStreamToolCallChecker + } + + g := compose.NewGraph[[]*schema.Message, *schema.Message]( + compose.WithGenLocalState(func(context.Context) *state { return &state{} })) + + if err := g.AddPassthroughNode(specialistsAnswersCollectorNodeKey); err != nil { + return nil, err + } + + agentTools := make([]*schema.ToolInfo, 0, len(config.Specialists)) + agentMap := make(map[string]bool, len(config.Specialists)+1) + for i := range config.Specialists { + specialist := config.Specialists[i] + + agentTools = append(agentTools, &schema.ToolInfo{ + Name: specialist.Name, + Desc: specialist.IntendedUse, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "reason": { + Type: schema.String, + Desc: "the reason to call this tool", + }, + }), + }) + + if err := addSpecialistAgent(specialist, g); err != nil { + return nil, err + } + + agentMap[specialist.Name] = true + } + + chatModel, err := agent.ChatModelWithTools(config.Host.ChatModel, config.Host.ToolCallingModel, agentTools) + if err != nil { + return nil, err + } + + if err = addHostAgent(chatModel, hostPrompt, g, hostKeyName); err != nil { + return nil, err + } + + const convertorName = "msg2MsgList" + if err = g.AddLambdaNode(convertorName, compose.ToList[*schema.Message](), compose.WithNodeName("converter")); err != nil { + return nil, err + } + + if err = addDirectAnswerBranch(convertorName, g, toolCallChecker); err != nil { + return nil, err + } + + if err = addMultiSpecialistsBranch(convertorName, agentMap, g); err != nil { + return nil, err + } + + if err = addSingleIntentAnswerNode(g); err != nil { + return nil, err + } + + if err = addMultiIntentsSummarizeNode(config.Summarizer, g); err != nil { + return nil, err + } + + if err = addAfterSpecialistsBranch(g); err != nil { + return nil, err + } + + compileOpts := []compose.GraphCompileOption{compose.WithNodeTriggerMode(compose.AnyPredecessor), compose.WithGraphName(name)} + r, err := g.Compile(ctx, compileOpts...) + if err != nil { + return nil, err + } + + return &MultiAgent{ + runnable: r, + graph: g, + graphAddNodeOpts: []compose.GraphAddNodeOpt{compose.WithGraphCompileOptions(compileOpts...)}, + }, nil +} + +func addSpecialistAgent(specialist *Specialist, g *compose.Graph[[]*schema.Message, *schema.Message]) error { + if specialist.Invokable != nil || specialist.Streamable != nil { + lambda, err := compose.AnyLambda(specialist.Invokable, specialist.Streamable, nil, nil, compose.WithLambdaType("Specialist")) + if err != nil { + return err + } + preHandler := func(_ context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) { + return state.msgs, nil // replace the tool call message with input msgs stored in state + } + if err := g.AddLambdaNode(specialist.Name, lambda, compose.WithStatePreHandler(preHandler), + compose.WithNodeName(specialist.Name), compose.WithOutputKey(specialist.Name)); err != nil { + return err + } + } else if specialist.ChatModel != nil { + preHandler := func(_ context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) { + if len(specialist.SystemPrompt) > 0 { + return append([]*schema.Message{{ + Role: schema.System, + Content: specialist.SystemPrompt, + }}, state.msgs...), nil + } + + return state.msgs, nil // replace the tool call message with input msgs stored in state + } + + if err := g.AddChatModelNode(specialist.Name, specialist.ChatModel, compose.WithStatePreHandler(preHandler), compose.WithNodeName(specialist.Name), compose.WithOutputKey(specialist.Name)); err != nil { + return err + } + } + + return g.AddEdge(specialist.Name, specialistsAnswersCollectorNodeKey) +} + +func addHostAgent(model model.BaseChatModel, prompt string, g *compose.Graph[[]*schema.Message, *schema.Message], hostNodeName string) error { + preHandler := func(_ context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) { + state.msgs = input + if len(prompt) == 0 { + return input, nil + } + return append([]*schema.Message{{ + Role: schema.System, + Content: prompt, + }}, input...), nil + } + if err := g.AddChatModelNode(defaultHostNodeKey, model, compose.WithStatePreHandler(preHandler), compose.WithNodeName(hostNodeName)); err != nil { + return err + } + + return g.AddEdge(compose.START, defaultHostNodeKey) +} + +func addDirectAnswerBranch(convertorName string, g *compose.Graph[[]*schema.Message, *schema.Message], + toolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error)) error { + // handles the case where the host agent returns a direct answer, instead of handling off to any specialist + branch := compose.NewStreamGraphBranch(func(ctx context.Context, sr *schema.StreamReader[*schema.Message]) (endNode string, err error) { + isToolCall, err := toolCallChecker(ctx, sr) + if err != nil { + return "", err + } + if isToolCall { + return convertorName, nil + } + return compose.END, nil + }, map[string]bool{convertorName: true, compose.END: true}) + + return g.AddBranch(defaultHostNodeKey, branch) +} + +func addMultiSpecialistsBranch(convertorName string, agentMap map[string]bool, g *compose.Graph[[]*schema.Message, *schema.Message]) error { + branch := compose.NewGraphMultiBranch(func(ctx context.Context, input []*schema.Message) (map[string]bool, error) { + if len(input) != 1 { + return nil, fmt.Errorf("host agent output %d messages, but expected 1", len(input)) + } + + results := map[string]bool{} + for _, toolCall := range input[0].ToolCalls { + results[toolCall.Function.Name] = true + } + + if len(results) > 1 { + _ = compose.ProcessState(ctx, func(_ context.Context, state *state) error { + state.isMultipleIntents = true + return nil + }) + } + + return results, nil + }, agentMap) + + return g.AddBranch(convertorName, branch) +} + +func addSingleIntentAnswerNode(g *compose.Graph[[]*schema.Message, *schema.Message]) error { + rc := func(ctx context.Context, input *schema.StreamReader[map[string]any]) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderWithConvert(input, func(msgs map[string]any) (*schema.Message, error) { + if len(msgs) != 1 { + return nil, fmt.Errorf("host agent output %d messages, but expected 1", len(msgs)) + } + for _, msg := range msgs { + return msg.(*schema.Message), nil + } + return nil, schema.ErrNoValue + }), nil + } + + _ = g.AddLambdaNode(singleIntentAnswerNodeKey, compose.TransformableLambda(rc)) + return g.AddEdge(singleIntentAnswerNodeKey, compose.END) +} + +func addAfterSpecialistsBranch(g *compose.Graph[[]*schema.Message, *schema.Message]) error { + ab := func(ctx context.Context, _ *schema.StreamReader[map[string]any]) (string, error) { + var isMultipleIntents bool + _ = compose.ProcessState(ctx, func(_ context.Context, state *state) error { + isMultipleIntents = state.isMultipleIntents + return nil + }) + + if !isMultipleIntents { + return singleIntentAnswerNodeKey, nil + } + + return map2ListConverterNodeKey, nil + } + + b := compose.NewStreamGraphBranch(ab, map[string]bool{ + singleIntentAnswerNodeKey: true, + map2ListConverterNodeKey: true, + }) + + return g.AddBranch(specialistsAnswersCollectorNodeKey, b) +} + +func addMultiIntentsSummarizeNode(summarizer *Summarizer, g *compose.Graph[[]*schema.Message, *schema.Message]) error { + map2list := func(ctx context.Context, input map[string]any) ([]*schema.Message, error) { + var output []*schema.Message + for k := range input { + output = append(output, input[k].(*schema.Message)) + } + return output, nil + } + + _ = g.AddLambdaNode(map2ListConverterNodeKey, compose.InvokableLambda(map2list)) + + if summarizer != nil { + _ = g.AddChatModelNode(multiIntentSummarizeNodeKey, summarizer.ChatModel, + compose.WithStatePreHandler(func(ctx context.Context, in []*schema.Message, state *state) ([]*schema.Message, error) { + var ( + out []*schema.Message + systemPrompt = defaultSummarizerPrompt + ) + if summarizer.SystemPrompt != "" { + systemPrompt = summarizer.SystemPrompt + } + out = append(out, &schema.Message{ + Role: schema.System, + Content: systemPrompt, + }) + + out = append(out, state.msgs...) + out = append(out, in...) + return out, nil + })) + _ = g.AddEdge(map2ListConverterNodeKey, multiIntentSummarizeNodeKey) + return g.AddEdge(multiIntentSummarizeNodeKey, compose.END) + } + + s := func(ctx context.Context, input []*schema.Message) (*schema.Message, error) { + output := &schema.Message{ + Role: schema.Assistant, + } + + for _, msg := range input { + output.Content += msg.Content + "\n" + } + + return output, nil + } + + _ = g.AddLambdaNode(multiIntentSummarizeNodeKey, compose.InvokableLambda(s)) + _ = g.AddEdge(map2ListConverterNodeKey, multiIntentSummarizeNodeKey) + return g.AddEdge(multiIntentSummarizeNodeKey, compose.END) +} diff --git a/flow/agent/multiagent/host/compose_test.go b/flow/agent/multiagent/host/compose_test.go new file mode 100644 index 0000000..c0afef4 --- /dev/null +++ b/flow/agent/multiagent/host/compose_test.go @@ -0,0 +1,691 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package host + +import ( + "context" + "io" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +func TestHostMultiAgent(t *testing.T) { + ctrl := gomock.NewController(t) + mockHostLLM := model.NewMockToolCallingChatModel(ctrl) + mockSpecialistLLM1 := model.NewMockChatModel(ctrl) + + specialist1 := &Specialist{ + ChatModel: mockSpecialistLLM1, + SystemPrompt: "You are a helpful assistant.", + AgentMeta: AgentMeta{ + Name: "specialist 1", + IntendedUse: "do stuff that works", + }, + } + + specialist2Msg1 := &schema.Message{ + Role: schema.Assistant, + Content: "specialist2", + } + specialist2Msg2 := &schema.Message{ + Role: schema.Assistant, + Content: " stream answer", + } + + specialist2 := &Specialist{ + Invokable: func(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) { + return &schema.Message{ + Role: schema.Assistant, + Content: "specialist2 invoke answer", + }, nil + }, + Streamable: func(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{specialist2Msg1, specialist2Msg2}), nil + }, + AgentMeta: AgentMeta{ + Name: "specialist 2", + IntendedUse: "do stuff that works too", + }, + } + + ctx := context.Background() + + mockHostLLM.EXPECT().WithTools(gomock.Any()).Return(mockHostLLM, nil).AnyTimes() + + hostMA, err := NewMultiAgent(ctx, &MultiAgentConfig{ + Host: Host{ + ToolCallingModel: mockHostLLM, + }, + Specialists: []*Specialist{ + specialist1, + specialist2, + }, + }) + + assert.NoError(t, err) + + t.Run("generate direct answer from host", func(t *testing.T) { + directAnswerMsg := &schema.Message{ + Role: schema.Assistant, + Content: "direct answer", + } + + mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(directAnswerMsg, nil).Times(1) + + mockCallback := newMockAgentCallback(0) + + out, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + assert.Equal(t, "direct answer", out.Content) + assert.Empty(t, mockCallback.infos) + }) + + t.Run("stream direct answer from host", func(t *testing.T) { + directAnswerMsg1 := &schema.Message{ + Role: schema.Assistant, + Content: "direct ", + } + + directAnswerMsg2 := &schema.Message{ + Role: schema.Assistant, + Content: "answer", + } + + sr, sw := schema.Pipe[*schema.Message](0) + go func() { + sw.Send(directAnswerMsg1, nil) + sw.Send(directAnswerMsg2, nil) + sw.Close() + }() + + mockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + + mockCallback := newMockAgentCallback(0) + outStream, err := hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + assert.Empty(t, mockCallback.infos) + + var msgs []*schema.Message + for { + msg, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + msgs = append(msgs, msg) + } + + outStream.Close() + + assert.Equal(t, directAnswerMsg1, msgs[0]) + assert.Equal(t, directAnswerMsg2, msgs[1]) + }) + + t.Run("generate hand off", func(t *testing.T) { + handOffMsg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{ + Name: specialist1.Name, + Arguments: `{"reason": "specialist 1 is the best"}`, + }, + }, + }, + } + + specialistMsg := &schema.Message{ + Role: schema.Assistant, + Content: "specialist 1 answer", + } + + mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1) + mockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(specialistMsg, nil).Times(1) + + mockCallback := newMockAgentCallback(1) + + out, err := hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + assert.Equal(t, "specialist 1 answer", out.Content) + mockCallback.wg.Wait() + assert.Equal(t, []*HandOffInfo{ + { + ToAgentName: specialist1.Name, + Argument: `{"reason": "specialist 1 is the best"}`, + }, + }, mockCallback.infos) + + handOffMsg.ToolCalls[0].Function.Name = specialist2.Name + handOffMsg.ToolCalls[0].Function.Arguments = `{"reason": "specialist 2 is even better"}` + mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1) + + mockCallback = newMockAgentCallback(1) + + out, err = hostMA.Generate(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + assert.Equal(t, "specialist2 invoke answer", out.Content) + mockCallback.wg.Wait() + assert.Equal(t, []*HandOffInfo{ + { + ToAgentName: specialist2.Name, + Argument: `{"reason": "specialist 2 is even better"}`, + }, + }, mockCallback.infos) + }) + + t.Run("stream hand off to chat model", func(t *testing.T) { + handOffMsg1 := &schema.Message{ + Role: schema.Assistant, + Content: "need to call function", + } + + handOffMsg2 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + }, + }, + } + + handOffMsg3 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{}, + }, + }, + } + + handOffMsg4 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{ + Name: specialist1.Name, + Arguments: `{"reason": "specialist 1 is the best"}`, + }, + }, + }, + } + + sr, sw := schema.Pipe[*schema.Message](0) + go func() { + sw.Send(handOffMsg1, nil) + sw.Send(handOffMsg2, nil) + sw.Send(handOffMsg3, nil) + sw.Send(handOffMsg4, nil) + sw.Close() + }() + + specialistMsg1 := &schema.Message{ + Role: schema.Assistant, + Content: "specialist ", + } + + specialistMsg2 := &schema.Message{ + Role: schema.Assistant, + Content: "1 answer", + } + + sr1, sw1 := schema.Pipe[*schema.Message](0) + go func() { + sw1.Send(specialistMsg1, nil) + sw1.Send(specialistMsg2, nil) + sw1.Close() + }() + + streamToolCallChecker := func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) { + defer modelOutput.Close() + + for { + msg, err := modelOutput.Recv() + if err != nil { + if err == io.EOF { + return false, nil + } + + return false, err + } + + if len(msg.ToolCalls) == 0 { + continue + } + + if len(msg.ToolCalls) > 0 { + return true, nil + } + } + } + + hostMA, err = NewMultiAgent(ctx, &MultiAgentConfig{ + Host: Host{ + ToolCallingModel: mockHostLLM, + }, + Specialists: []*Specialist{ + specialist1, + specialist2, + }, + StreamToolCallChecker: streamToolCallChecker, + }) + assert.NoError(t, err) + + mockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + mockSpecialistLLM1.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr1, nil).Times(1) + + mockCallback := newMockAgentCallback(1) + outStream, err := hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + + var msgs []*schema.Message + for { + msg, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + msgs = append(msgs, msg) + } + + outStream.Close() + + assert.Equal(t, specialistMsg1, msgs[0]) + assert.Equal(t, specialistMsg2, msgs[1]) + + mockCallback.wg.Wait() + + assert.Equal(t, []*HandOffInfo{ + { + ToAgentName: specialist1.Name, + Argument: `{"reason": "specialist 1 is the best"}`, + }, + }, mockCallback.infos) + + handOffMsg4.ToolCalls[0].Function.Name = specialist2.Name + handOffMsg4.ToolCalls[0].Function.Arguments = `{"reason": "specialist 2 is even better"}` + sr, sw = schema.Pipe[*schema.Message](0) + go func() { + sw.Send(handOffMsg1, nil) + sw.Send(handOffMsg2, nil) + sw.Send(handOffMsg3, nil) + sw.Send(handOffMsg4, nil) + sw.Close() + }() + + mockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + + mockCallback = newMockAgentCallback(1) + outStream, err = hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + + msgs = nil + for { + msg, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + msgs = append(msgs, msg) + } + + outStream.Close() + + assert.Equal(t, specialist2Msg1, msgs[0]) + assert.Equal(t, specialist2Msg2, msgs[1]) + + mockCallback.wg.Wait() + + assert.Equal(t, []*HandOffInfo{ + { + ToAgentName: specialist2.Name, + Argument: `{"reason": "specialist 2 is even better"}`, + }, + }, mockCallback.infos) + }) + + t.Run("multi-agent within graph", func(t *testing.T) { + handOffMsg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{ + Name: specialist1.Name, + Arguments: `{"reason": "specialist 1 is the best"}`, + }, + }, + }, + } + + specialistMsg := &schema.Message{ + Role: schema.Assistant, + Content: "Beijing", + } + + mockHostLLM.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(handOffMsg, nil).Times(1) + mockSpecialistLLM1.EXPECT().Generate(gomock.Any(), gomock.Any()).Return(specialistMsg, nil).Times(1) + + mockCallback := newMockAgentCallback(1) + + hostMA, err := NewMultiAgent(ctx, &MultiAgentConfig{ + Host: Host{ + ToolCallingModel: mockHostLLM, + }, + Specialists: []*Specialist{ + specialist1, + specialist2, + }, + }) + + assert.NoError(t, err) + + maGraph, opts := hostMA.ExportGraph() + + fullGraph, err := compose.NewChain[map[string]any, *schema.Message](). + AppendChatTemplate(prompt.FromMessages(schema.FString, schema.UserMessage("what's the capital city of {country_name}"))). + AppendGraph(maGraph, append(opts, compose.WithNodeKey("host_ma_node"))...). + Compile(ctx) + assert.NoError(t, err) + + out, err := fullGraph.Invoke(ctx, map[string]any{"country_name": "China"}, compose.WithCallbacks(ConvertCallbackHandlers(mockCallback)).DesignateNodeWithPath(compose.NewNodePath("host_ma_node", hostMA.HostNodeKey()))) + assert.NoError(t, err) + assert.Equal(t, "Beijing", out.Content) + + mockCallback.wg.Wait() + assert.Equal(t, []*HandOffInfo{ + { + ToAgentName: specialist1.Name, + Argument: `{"reason": "specialist 1 is the best"}`, + }, + }, mockCallback.infos) + }) + + t.Run("multiple intents", func(t *testing.T) { + handOffMsg1 := &schema.Message{ + Role: schema.Assistant, + Content: "need to call function", + } + + handOffMsg2 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + }, + }, + } + + handOffMsg3 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{}, + }, + }, + } + + handOffMsg4 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{ + Name: specialist1.Name, + Arguments: `{"reason": "specialist 1 is good"}`, + }, + }, { + Index: generic.PtrOf(1), + Function: schema.FunctionCall{ + Name: specialist2.Name, + Arguments: `{"reason": "specialist 2`, + }, + }, + }, + } + + handOffMsg5 := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(1), + Function: schema.FunctionCall{ + Arguments: ` is also good"}`, + }, + }, + }, + } + + sr := schema.StreamReaderFromArray([]*schema.Message{ + handOffMsg1, + handOffMsg2, + handOffMsg3, + handOffMsg4, + handOffMsg5, + }) + + specialist1Msg1 := &schema.Message{ + Role: schema.Assistant, + Content: "specialist ", + } + + specialist1Msg2 := &schema.Message{ + Role: schema.Assistant, + Content: "1 answer", + } + + sr1 := schema.StreamReaderFromArray([]*schema.Message{ + specialist1Msg1, + specialist1Msg2, + }) + + streamToolCallChecker := func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) { + defer modelOutput.Close() + + for { + msg, err := modelOutput.Recv() + if err != nil { + if err == io.EOF { + return false, nil + } + + return false, err + } + + if len(msg.ToolCalls) == 0 { + continue + } + + if len(msg.ToolCalls) > 0 { + return true, nil + } + } + } + + hostMA, err = NewMultiAgent(ctx, &MultiAgentConfig{ + Host: Host{ + ToolCallingModel: mockHostLLM, + }, + Specialists: []*Specialist{ + specialist1, + specialist2, + }, + StreamToolCallChecker: streamToolCallChecker, + }) + assert.NoError(t, err) + + mockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + mockSpecialistLLM1.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr1, nil).Times(1) + + mockCallback := newMockAgentCallback(2) + outStream, err := hostMA.Stream(ctx, nil, WithAgentCallbacks(mockCallback)) + assert.NoError(t, err) + + var msgs []*schema.Message + for { + msg, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + msgs = append(msgs, msg) + } + + outStream.Close() + + msg, err := schema.ConcatMessages(msgs) + assert.NoError(t, err) + if msg.Content != "specialist2 stream answer\nspecialist 1 answer\n" && + msg.Content != "specialist 1 answer\nspecialist2 stream answer\n" { + t.Errorf("Unexpected message content: %s", msg.Content) + } + + mockCallback.wg.Wait() + assert.Equal(t, []*HandOffInfo{ + { + ToAgentName: specialist1.Name, + Argument: `{"reason": "specialist 1 is good"}`, + }, + { + ToAgentName: specialist2.Name, + Argument: `{"reason": "specialist 2 is also good"}`, + }, + }, mockCallback.infos) + }) + + t.Run("summarize multiple intents", func(t *testing.T) { + + handOffMsg := &schema.Message{ + Role: schema.Assistant, + ToolCalls: []schema.ToolCall{ + { + Index: generic.PtrOf(0), + Function: schema.FunctionCall{ + Name: specialist1.Name, + Arguments: `{"reason": "specialist 1 is good"}`, + }, + }, { + Index: generic.PtrOf(1), + Function: schema.FunctionCall{ + Name: specialist2.Name, + Arguments: `{"reason": "specialist 2 is also good"}`, + }, + }, + }, + } + sr := schema.StreamReaderFromArray([]*schema.Message{ + handOffMsg, + }) + + specialist1Msg1 := &schema.Message{ + Role: schema.Assistant, + Content: "specialist 1 answer", + } + sr1 := schema.StreamReaderFromArray([]*schema.Message{ + specialist1Msg1, + }) + + const summaryContent = "summarized answer" + sr2 := schema.StreamReaderFromArray([]*schema.Message{ + { + Role: schema.Assistant, + Content: summaryContent, + }, + }) + + mockSumChatModel := model.NewMockChatModel(ctrl) + hostMA, err = NewMultiAgent(ctx, &MultiAgentConfig{ + Host: Host{ + ToolCallingModel: mockHostLLM, + }, + Specialists: []*Specialist{ + specialist1, + specialist2, + }, + Summarizer: &Summarizer{ + ChatModel: mockSumChatModel, + }, + }) + assert.NoError(t, err) + + mockHostLLM.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr, nil).Times(1) + mockSpecialistLLM1.EXPECT().Stream(gomock.Any(), gomock.Any()).Return(sr1, nil).Times(1) + mockSumChatModel.EXPECT().Stream(gomock.Any(), gomock.Cond(func(x any) bool { + return assert.Equal(t, defaultSummarizerPrompt, func() string { + if input := x.([]*schema.Message); len(input) > 0 { + return input[0].Content + } + return "" + }()) + })).Return(sr2, nil).Times(1) + + outStream, err := hostMA.Stream(ctx, nil) + assert.NoError(t, err) + + var msgs []*schema.Message + for { + msg, err := outStream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + msgs = append(msgs, msg) + } + outStream.Close() + + msg, err := schema.ConcatMessages(msgs) + assert.NoError(t, err) + if msg.Content != summaryContent { + t.Errorf("Unexpected message content: %s", msg.Content) + } + }, + ) +} + +type mockAgentCallback struct { + infos []*HandOffInfo + wg sync.WaitGroup +} + +func (m *mockAgentCallback) OnHandOff(ctx context.Context, info *HandOffInfo) context.Context { + m.infos = append(m.infos, info) + m.wg.Done() + return ctx +} + +func newMockAgentCallback(expects int) *mockAgentCallback { + m := &mockAgentCallback{ + infos: make([]*HandOffInfo, 0), + wg: sync.WaitGroup{}, + } + + m.wg.Add(expects) + return m +} diff --git a/flow/agent/multiagent/host/doc.go b/flow/agent/multiagent/host/doc.go new file mode 100644 index 0000000..3d205ce --- /dev/null +++ b/flow/agent/multiagent/host/doc.go @@ -0,0 +1,17 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package host diff --git a/flow/agent/multiagent/host/options.go b/flow/agent/multiagent/host/options.go new file mode 100644 index 0000000..55b5395 --- /dev/null +++ b/flow/agent/multiagent/host/options.go @@ -0,0 +1,30 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package host + +import "github.com/cloudwego/eino/flow/agent" + +type options struct { + agentCallbacks []MultiAgentCallback +} + +// WithAgentCallbacks registers callbacks to be invoked by the host multi-agent. +func WithAgentCallbacks(agentCallbacks ...MultiAgentCallback) agent.AgentOption { + return agent.WrapImplSpecificOptFn(func(opts *options) { + opts.agentCallbacks = append(opts.agentCallbacks, agentCallbacks...) + }) +} diff --git a/flow/agent/multiagent/host/types.go b/flow/agent/multiagent/host/types.go new file mode 100644 index 0000000..19d417a --- /dev/null +++ b/flow/agent/multiagent/host/types.go @@ -0,0 +1,204 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package host implements the host pattern for multi-agent system. +package host + +import ( + "context" + "errors" + "fmt" + "io" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + "github.com/cloudwego/eino/schema" +) + +// MultiAgent is a host multi-agent system. +// A host agent is responsible for deciding which specialist to 'hand off' the task to. +// One or more specialist agents are responsible for completing the task. +type MultiAgent struct { + runnable compose.Runnable[[]*schema.Message, *schema.Message] + graph *compose.Graph[[]*schema.Message, *schema.Message] + graphAddNodeOpts []compose.GraphAddNodeOpt +} + +// Generate runs the multi-agent synchronously and returns the final message. +func (ma *MultiAgent) Generate(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) { + composeOptions := agent.GetComposeOptions(opts...) + + handler := convertCallbacks(opts...) + if handler != nil { + composeOptions = append(composeOptions, compose.WithCallbacks(handler).DesignateNode(ma.HostNodeKey())) + } + + return ma.runnable.Invoke(ctx, input, composeOptions...) +} + +// Stream runs the multi-agent in streaming mode and returns a message stream. +func (ma *MultiAgent) Stream(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.StreamReader[*schema.Message], error) { + composeOptions := agent.GetComposeOptions(opts...) + + handler := convertCallbacks(opts...) + if handler != nil { + composeOptions = append(composeOptions, compose.WithCallbacks(handler).DesignateNode(ma.HostNodeKey())) + } + + return ma.runnable.Stream(ctx, input, composeOptions...) +} + +// ExportGraph exports the underlying graph from MultiAgent, along with the []compose.GraphAddNodeOpt to be used when adding this graph to another graph. +func (ma *MultiAgent) ExportGraph() (compose.AnyGraph, []compose.GraphAddNodeOpt) { + return ma.graph, ma.graphAddNodeOpts +} + +// HostNodeKey returns the graph node key used for the host agent. +func (ma *MultiAgent) HostNodeKey() string { + return defaultHostNodeKey +} + +// MultiAgentConfig is the config for host multi-agent system. +type MultiAgentConfig struct { + Host Host + Specialists []*Specialist + + Name string // the name of the host multi-agent + HostNodeName string // the name of the host node in the graph, default is "host" + // StreamToolCallChecker is a function to determine whether the model's streaming output contains tool calls. + // Different models have different ways of outputting tool calls in streaming mode: + // - Some models (like OpenAI) output tool calls directly + // - Others (like Claude) output text first, then tool calls + // This handler allows custom logic to check for tool calls in the stream. + // It should return: + // - true if the output contains tool calls and agent should continue processing + // - false if no tool calls and agent should stop + // Note: This field only needs to be configured when using streaming mode + // Note: The handler MUST close the modelOutput stream before returning + // Optional. By default, it checks if the first chunk contains tool calls. + // Note: The default implementation does not work well with Claude, which typically outputs tool calls after text content. + // Note: If your ChatModel doesn't output tool calls first, you can try adding prompts to constrain the model from generating extra text during the tool call. + StreamToolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) + + // Summarizer is the summarizer agent that will summarize the outputs of all the chosen specialist agents. + // Only when the Host agent picks multiple Specialist will this be called. + // If you do not provide a summarizer, a default summarizer that simply concatenates all the output messages into one message will be used. + // Note: the default summarizer do not support streaming. + Summarizer *Summarizer +} + +func (conf *MultiAgentConfig) validate() error { + if conf == nil { + return errors.New("host multi agent config is nil") + } + + if conf.Host.ChatModel == nil && conf.Host.ToolCallingModel == nil { + return errors.New("host multi agent host ChatModel is nil") + } + + if len(conf.Specialists) == 0 { + return errors.New("host multi agent specialists are empty") + } + + for _, s := range conf.Specialists { + if s.ChatModel == nil && s.Invokable == nil && s.Streamable == nil { + return fmt.Errorf("specialist %s has no chat model or Invokable or Streamable", s.Name) + } + + if err := s.AgentMeta.validate(); err != nil { + return err + } + } + + return nil +} + +// AgentMeta is the meta information of an agent within a multi-agent system. +type AgentMeta struct { + Name string // the name of the agent, should be unique within multi-agent system + IntendedUse string // the intended use-case of the agent, used as the reason for the multi-agent system to hand over control to this agent +} + +func (am AgentMeta) validate() error { + if len(am.Name) == 0 { + return errors.New("agent meta name is empty") + } + + if len(am.IntendedUse) == 0 { + return errors.New("agent meta intended use is empty") + } + + return nil +} + +// Host is the host agent within a multi-agent system. +// Currently, it can only be a model.ChatModel. +type Host struct { + ToolCallingModel model.ToolCallingChatModel + // Deprecated: ChatModel is deprecated, please use ToolCallingModel instead. + // This field will be removed in a future release. + ChatModel model.ChatModel + SystemPrompt string +} + +// Specialist is a specialist agent within a host multi-agent system. +// It can be a model.ChatModel or any Invokable and/or Streamable, such as react.Agent. +// ChatModel and (Invokable / Streamable) are mutually exclusive, only one should be provided. +// notice: SystemPrompt only effects when ChatModel has been set. +// If Invokable is provided but not Streamable, then the Specialist will be 'compose.InvokableLambda'. +// If Streamable is provided but not Invokable, then the Specialist will be 'compose.StreamableLambda'. +// if Both Invokable and Streamable is provided, then the Specialist will be 'compose.AnyLambda'. +type Specialist struct { + AgentMeta + + ChatModel model.BaseChatModel + SystemPrompt string + + Invokable compose.Invoke[[]*schema.Message, *schema.Message, agent.AgentOption] + Streamable compose.Stream[[]*schema.Message, *schema.Message, agent.AgentOption] +} + +// Summarizer defines a lightweight agent used to summarize +// conversations or tool outputs using a chat model and prompt. +type Summarizer struct { + ChatModel model.BaseChatModel + SystemPrompt string +} + +func firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) { + defer sr.Close() + + for { + msg, err := sr.Recv() + if err == io.EOF { + return false, nil + } + if err != nil { + return false, err + } + + if len(msg.ToolCalls) > 0 { + return true, nil + } + + if len(msg.Content) == 0 { // skip empty chunks at the front + continue + } + + return false, nil + } +} diff --git a/flow/agent/react/callback.go b/flow/agent/react/callback.go new file mode 100644 index 0000000..d2e4c76 --- /dev/null +++ b/flow/agent/react/callback.go @@ -0,0 +1,33 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package react provides helpers to build callback handlers for React agents. +package react + +import ( + "github.com/cloudwego/eino/callbacks" + template "github.com/cloudwego/eino/utils/callbacks" +) + +// BuildAgentCallback builds a callback handler for agent. +// e.g. +// +// callback := BuildAgentCallback(modelHandler, toolHandler) +// agent, err := react.NewAgent(ctx, &AgentConfig{}) +// agent.Generate(ctx, input, agent.WithComposeOptions(compose.WithCallbacks(callback))) +func BuildAgentCallback(modelHandler *template.ModelCallbackHandler, toolHandler *template.ToolCallbackHandler) callbacks.Handler { + return template.NewHandlerHelper().ChatModel(modelHandler).Tool(toolHandler).Handler() +} diff --git a/flow/agent/react/doc.go b/flow/agent/react/doc.go new file mode 100644 index 0000000..1e012c6 --- /dev/null +++ b/flow/agent/react/doc.go @@ -0,0 +1,17 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package react diff --git a/flow/agent/react/option.go b/flow/agent/react/option.go new file mode 100644 index 0000000..e6c9fb4 --- /dev/null +++ b/flow/agent/react/option.go @@ -0,0 +1,401 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package react + +import ( + "context" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/schema" + ub "github.com/cloudwego/eino/utils/callbacks" +) + +// WithToolOptions returns an agent option that specifies tool.Option for the tools in agent. +func WithToolOptions(opts ...tool.Option) agent.AgentOption { + return agent.WithComposeOptions(compose.WithToolsNodeOption(compose.WithToolOption(opts...))) +} + +// WithChatModelOptions returns an agent option that specifies model.Option for the chat model in agent. +func WithChatModelOptions(opts ...model.Option) agent.AgentOption { + return agent.WithComposeOptions(compose.WithChatModelOption(opts...)) +} + +// WithToolList returns an agent option that specifies compose.ToolsNodeOption for ToolsNode in agent. +// If you also need to pass ToolInfo to the chat model, use WithTools instead. +// Deprecated: This changes tool list for ToolsNode ONLY. +func WithToolList(tools ...tool.BaseTool) agent.AgentOption { + return agent.WithComposeOptions(compose.WithToolsNodeOption(compose.WithToolList(tools...))) +} + +// WithTools is a convenience function that configures a React agent with a list of tools. +// It performs two essential operations: +// 1. Extracts tool information for the chat model to understand available tools +// 2. Registers the actual tool implementations for execution +// +// Parameters: +// - ctx: The context for the operation, used when calling Info() on each tool +// - tools: A variadic list of tools that must implement either InvokableTool or StreamableTool interfaces +// +// Returns: +// - []agent.AgentOption: A slice containing exactly 2 agent options: +// - Option 1: Configures the chat model with tool schemas via model.WithTools(toolInfos) +// - Option 2: Registers the tool implementations via compose.WithToolList(tools...) +// - error: Returns an error if any tool's Info() method fails +// +// Usage Example: +// +// ctx := context.Background() +// agentOptions, err := WithTools(ctx, myTool1, myTool2, myTool3) +// if err != nil { +// return fmt.Errorf("failed to configure tools: %w", err) +// } +// +// agent, err := react.NewAgent(ctx, &react.AgentConfig{ +// ToolCallingModel: myModel, +// // other config... +// }) +// if err != nil { +// return fmt.Errorf("failed to create agent: %w", err) +// } +// +// // Use the tool options with Generate or Stream methods +// msg, err := agent.Generate(ctx, messages, agentOptions...) +// // or +// stream, err := agent.Stream(ctx, messages, agentOptions...) +// +// Comparison with Related Functions: +// - WithToolList: Only registers tool implementations, doesn't configure the chat model +// - WithTools: Comprehensive setup that handles both chat model configuration and tool registration +// +// Notes: +// - The function always returns exactly 2 options when successful +// - Both returned options should be applied to the agent for proper tool functionality +func WithTools(ctx context.Context, tools ...tool.BaseTool) ([]agent.AgentOption, error) { + toolInfos := make([]*schema.ToolInfo, 0, len(tools)) + for _, tl := range tools { + info, err := tl.Info(ctx) + if err != nil { + return nil, err + } + + toolInfos = append(toolInfos, info) + } + + opts := make([]agent.AgentOption, 2) + opts[0] = agent.WithComposeOptions(compose.WithChatModelOption(model.WithTools(toolInfos))) + opts[1] = agent.WithComposeOptions(compose.WithToolsNodeOption(compose.WithToolList(tools...))) + return opts, nil +} + +// Iterator provides a lightweight FIFO stream of values and errors +// produced during agent execution. +type Iterator[T any] struct { + ch *internal.UnboundedChan[item[T]] +} + +// Next retrieves the next value from the iterator. +// It returns the zero value and false when the stream is exhausted. +func (iter *Iterator[T]) Next() (T, bool, error) { + ch := iter.ch + if ch == nil { + var zero T + return zero, false, nil + } + + i, ok := ch.Receive() + if !ok { + var zero T + return zero, false, nil + } + + return i.v, true, i.err +} + +// MessageFuture exposes asynchronous accessors for messages produced +// by Generate and Stream calls. +type MessageFuture interface { + // GetMessages returns an iterator for retrieving messages generated during "agent.Generate" calls. + GetMessages() *Iterator[*schema.Message] + + // GetMessageStreams returns an iterator for retrieving streaming messages generated during "agent.Stream" calls. + GetMessageStreams() *Iterator[*schema.StreamReader[*schema.Message]] +} + +// WithMessageFuture returns an agent option and a MessageFuture interface instance. +// The option configures the agent to collect messages generated during execution, +// while the MessageFuture interface allows users to asynchronously retrieve these messages. +// +// This function works correctly both when the agent is used directly and when it is +// embedded as a subgraph within another graph. Graph callbacks are filtered using +// the address in the context: only callbacks whose address contains a runnable segment +// matching the agent's configured graph name are processed. +func WithMessageFuture() (agent.AgentOption, MessageFuture) { + h := &cbHandler{started: make(chan struct{}), graphName: GraphName} + + cmHandler := &ub.ModelCallbackHandler{ + OnEnd: h.onChatModelEnd, + OnEndWithStreamOutput: h.onChatModelEndWithStreamOutput, + } + createToolResultSender := func() toolResultSender { + return func(toolName, callID, result string) { + msg := schema.ToolMessage(result, callID, schema.WithToolName(toolName)) + h.sendMessage(msg) + } + } + createStreamToolResultSender := func() streamToolResultSender { + return func(toolName, callID string, resultStream *schema.StreamReader[string]) { + cvt := func(in string) (*schema.Message, error) { + return schema.ToolMessage(in, callID, schema.WithToolName(toolName)), nil + } + msgStream := schema.StreamReaderWithConvert(resultStream, cvt) + h.sendMessageStream(msgStream) + } + } + + createEnhancedToolResultSender := func() enhancedToolResultSender { + return func(toolName, callID string, result *schema.ToolResult) { + var err error + msg := schema.ToolMessage("", callID, schema.WithToolName(toolName)) + msg.UserInputMultiContent, err = result.ToMessageInputParts() + if err != nil { + return + } + h.sendMessage(msg) + } + } + + createEnhancedStreamToolResultSender := func() enhancedStreamToolResultSender { + return func(toolName, callID string, resultStream *schema.StreamReader[*schema.ToolResult]) { + cvt := func(result *schema.ToolResult) (*schema.Message, error) { + var err error + msg := schema.ToolMessage("", callID, schema.WithToolName(toolName)) + msg.UserInputMultiContent, err = result.ToMessageInputParts() + if err != nil { + return nil, err + } + return msg, nil + } + msgStream := schema.StreamReaderWithConvert(resultStream, cvt) + h.sendMessageStream(msgStream) + } + } + + graphHandler := callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + ctx = h.onGraphStart(ctx, info, input) + return setToolResultSendersToCtx(ctx, &toolResultSenders{ + sender: createToolResultSender(), + streamSender: createStreamToolResultSender(), + enhancedResultSender: createEnhancedToolResultSender(), + enhancedStreamToolResultSender: createEnhancedStreamToolResultSender(), + }) + }). + OnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + ctx = h.onGraphStartWithStreamInput(ctx, info, input) + return setToolResultSendersToCtx(ctx, &toolResultSenders{ + sender: createToolResultSender(), + streamSender: createStreamToolResultSender(), + enhancedResultSender: createEnhancedToolResultSender(), + enhancedStreamToolResultSender: createEnhancedStreamToolResultSender(), + }) + }). + OnEndFn(h.onGraphEnd). + OnEndWithStreamOutputFn(h.onGraphEndWithStreamOutput). + OnErrorFn(h.onGraphError).Build() + cb := ub.NewHandlerHelper().ChatModel(cmHandler).Graph(graphHandler).Handler() + option := agent.WithComposeOptions(compose.WithCallbacks(cb)) + + return option, h +} + +type item[T any] struct { + v T + err error +} + +type cbHandler struct { + graphName string + + ownAddress compose.Address + ownClaimed bool + + msgs *internal.UnboundedChan[item[*schema.Message]] + sMsgs *internal.UnboundedChan[item[*schema.StreamReader[*schema.Message]]] + + started chan struct{} +} + +func (h *cbHandler) GetMessages() *Iterator[*schema.Message] { + <-h.started + + return &Iterator[*schema.Message]{ch: h.msgs} +} + +func (h *cbHandler) GetMessageStreams() *Iterator[*schema.StreamReader[*schema.Message]] { + <-h.started + + return &Iterator[*schema.StreamReader[*schema.Message]]{ch: h.sMsgs} +} + +// isOwnGraph reports whether the callback is being invoked for the React agent's own graph. +// +// After the first onGraphStart call records h.ownAddress, subsequent calls compare the +// current context address against that exact recorded address. This handles nested React +// agents correctly: each cbHandler instance records its own unique full address path, so +// two React agents at different nesting depths (e.g., an outer agent whose tool invokes +// another React agent) will never interfere with each other even if they share the same +// graph name. +func (h *cbHandler) isOwnGraph(ctx context.Context) bool { + if !h.ownClaimed { + return false + } + return compose.GetCurrentAddress(ctx).Equals(h.ownAddress) +} + +// claimOwnership is called exclusively from onGraphStart / onGraphStartWithStreamInput +// to record this handler's address on first invocation. It returns true if the handler +// has not yet been initialised, records the current address, and returns false on all +// subsequent calls. +func (h *cbHandler) claimOwnership(ctx context.Context) bool { + if h.ownClaimed { + return false + } + h.ownAddress = compose.GetCurrentAddress(ctx) + h.ownClaimed = true + return true +} + +func (h *cbHandler) onChatModelEnd(ctx context.Context, + _ *callbacks.RunInfo, input *model.CallbackOutput) context.Context { + + h.sendMessage(input.Message) + + return ctx +} + +func (h *cbHandler) onChatModelEndWithStreamOutput(ctx context.Context, + _ *callbacks.RunInfo, input *schema.StreamReader[*model.CallbackOutput]) context.Context { + + c := func(output *model.CallbackOutput) (*schema.Message, error) { + return output.Message, nil + } + s := schema.StreamReaderWithConvert(input, c) + + h.sendMessageStream(s) + + return ctx +} + +func (h *cbHandler) onGraphError(ctx context.Context, + _ *callbacks.RunInfo, err error) context.Context { + + if !h.isOwnGraph(ctx) { + return ctx + } + + if h.msgs != nil { + h.msgs.Send(item[*schema.Message]{err: err}) + } else if h.sMsgs != nil { + h.sMsgs.Send(item[*schema.StreamReader[*schema.Message]]{err: err}) + } + + return ctx +} + +func (h *cbHandler) onGraphEnd(ctx context.Context, + _ *callbacks.RunInfo, _ callbacks.CallbackOutput) context.Context { + + if !h.isOwnGraph(ctx) { + return ctx + } + + if h.msgs != nil { + h.msgs.Close() + } + + return ctx +} + +func (h *cbHandler) onGraphEndWithStreamOutput(ctx context.Context, + _ *callbacks.RunInfo, _ *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + + if !h.isOwnGraph(ctx) { + return ctx + } + + if h.sMsgs != nil { + h.sMsgs.Close() + } + + return ctx +} + +func (h *cbHandler) onGraphStart(ctx context.Context, + _ *callbacks.RunInfo, _ callbacks.CallbackInput) context.Context { + + if !h.claimOwnership(ctx) { + return ctx + } + + h.msgs = internal.NewUnboundedChan[item[*schema.Message]]() + close(h.started) + + return ctx +} + +func (h *cbHandler) onGraphStartWithStreamInput(ctx context.Context, _ *callbacks.RunInfo, + input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + input.Close() + + if !h.claimOwnership(ctx) { + return ctx + } + + h.sMsgs = internal.NewUnboundedChan[item[*schema.StreamReader[*schema.Message]]]() + close(h.started) + + return ctx +} + +func (h *cbHandler) sendMessage(msg *schema.Message) { + if h.msgs != nil { + h.msgs.Send(item[*schema.Message]{v: msg}) + } else { + sMsg := schema.StreamReaderFromArray([]*schema.Message{msg}) + h.sMsgs.Send(item[*schema.StreamReader[*schema.Message]]{v: sMsg}) + } +} + +func (h *cbHandler) sendMessageStream(sMsg *schema.StreamReader[*schema.Message]) { + if h.sMsgs != nil { + h.sMsgs.Send(item[*schema.StreamReader[*schema.Message]]{v: sMsg}) + } else { + // concat + msg, err := schema.ConcatMessageStream(sMsg) + + if err != nil { + h.msgs.Send(item[*schema.Message]{err: err}) + } else { + h.msgs.Send(item[*schema.Message]{v: msg}) + } + } +} diff --git a/flow/agent/react/option_test.go b/flow/agent/react/option_test.go new file mode 100644 index 0000000..1449089 --- /dev/null +++ b/flow/agent/react/option_test.go @@ -0,0 +1,1388 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package react + +import ( + "context" + "io" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" +) + +func TestWithMessageFuture(t *testing.T) { + ctx := context.Background() + + // Test with tool calls + t.Run("test generate with tool calls", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Mock model response with tool call + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create agent with MessageFuture + option, future := WithMessageFuture() + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + // Generate response + response, err := a.Generate(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, option) + assert.Nil(t, err) + assert.Equal(t, "final response", response.Content) + + sIter := future.GetMessageStreams() + // Should be no messages + _, hasNext, err := sIter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + + iter := future.GetMessages() + // First message should be the assistant message for tool calling + msg1, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + // Second message should be the tool response + msg2, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Tool, msg2.Role) + + // Third message should be the final response + msg3, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, "final response", msg3.Content) + + // Should be no more messages + _, hasNext, err = iter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + // Test with streaming tool calls + t.Run("test generate with streaming tool calls", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeStreamToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Mock model response with tool call + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create agent with MessageFuture + option, future := WithMessageFuture() + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + // Generate response + response, err := a.Generate(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, option) + assert.Nil(t, err) + assert.Equal(t, "final response", response.Content) + + // Get messages from future + iter := future.GetMessages() + + // First message should be the assistant message for tool calling + msg1, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + // Second message should be the tool response + msg2, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Tool, msg2.Role) + + // Third message should be the final response + msg3, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, "final response", msg3.Content) + + // Should be no more messages + _, hasNext, err = iter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + + // Test with non-streaming tool but using agent's Stream interface + t.Run("test stream with tool calls", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Mock model response with tool call + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }, + })}), nil). + Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final response", nil)}), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create agent with MessageFuture + option, future := WithMessageFuture() + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + // Use Stream interface + stream, err := a.Stream(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, option) + assert.Nil(t, err) + + // Collect all chunks from stream + finalResponse, err := schema.ConcatMessageStream(stream) + assert.Nil(t, err) + assert.Equal(t, "final response", finalResponse.Content) + + iter := future.GetMessages() + // Should be no messages + _, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + + // Get message streams from future + sIter := future.GetMessageStreams() + + // First message should be the assistant message for tool calling + stream1, hasNext, err := sIter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.NotNil(t, stream1) + msg1, err := schema.ConcatMessageStream(stream1) + assert.Nil(t, err) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + // Second message should be the tool response + stream2, hasNext, err := sIter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.NotNil(t, stream2) + msg2, err := schema.ConcatMessageStream(stream2) + assert.Nil(t, err) + assert.Equal(t, schema.Tool, msg2.Role) + + // Third message should be the final response + stream3, hasNext, err := sIter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.NotNil(t, stream3) + msg3, err := schema.ConcatMessageStream(stream3) + assert.Nil(t, err) + assert.Equal(t, "final response", msg3.Content) + + // Should be no more messages + _, hasNext, err = sIter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + + t.Run("test stream with streaming tool calls and with concurrent goroutines", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeStreamToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + // Mock model response with tool call + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }, + })}), nil). + Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final response", nil)}), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + // Create agent with MessageFuture + option, future := WithMessageFuture() + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + + // Get message streams from future + sIter := future.GetMessageStreams() + + // First message should be the assistant message for tool calling + stream1, hasNext, err_ := sIter.Next() + assert.Nil(t, err_) + assert.True(t, hasNext) + assert.NotNil(t, stream1) + msg1, err_ := schema.ConcatMessageStream(stream1) + assert.Nil(t, err_) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + // Second message should be the tool response + stream2, hasNext, err_ := sIter.Next() + assert.Nil(t, err_) + assert.True(t, hasNext) + assert.NotNil(t, stream2) + msg2, err_ := schema.ConcatMessageStream(stream2) + assert.Nil(t, err_) + assert.Equal(t, schema.Tool, msg2.Role) + + // Third message should be the final response + stream3, hasNext, err_ := sIter.Next() + assert.Nil(t, err_) + assert.True(t, hasNext) + assert.NotNil(t, stream3) + msg3, err_ := schema.ConcatMessageStream(stream3) + assert.Nil(t, err_) + assert.Equal(t, "final response", msg3.Content) + + // Should be no more messages + _, hasNext, err_ = sIter.Next() + assert.Nil(t, err_) + assert.False(t, hasNext) + }() + + // Use Stream interface + stream, err := a.Stream(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, option) + assert.Nil(t, err) + + // Collect all chunks from stream + finalResponse, err := schema.ConcatMessageStream(stream) + assert.Nil(t, err) + assert.Equal(t, "final response", finalResponse.Content) + + wg.Wait() + }) +} + +func TestWithToolOptions(t *testing.T) { + type dummyOpt struct{ val string } + opt := tool.WrapImplSpecificOptFn(func(o *dummyOpt) { o.val = "mock" }) + agentOpt := WithToolOptions(opt) + assert.NotNil(t, agentOpt) + // The returned value should be an agent.AgentOption (function) + assert.IsType(t, agentOpt, agentOpt) +} + +func TestWithChatModelOptions(t *testing.T) { + opt := model.WithModel("mock-model") + agentOpt := WithChatModelOptions(opt) + assert.NotNil(t, agentOpt) + assert.IsType(t, agentOpt, agentOpt) +} + +// dummyBaseTool is a minimal implementation of tool.BaseTool for testing. +type dummyBaseTool struct{} + +func (d *dummyBaseTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "dummy"}, nil +} + +func (d *dummyBaseTool) InvokableRun(ctx context.Context, _ string, _ ...tool.Option) (string, error) { + return "dummy-response", nil +} + +type assertTool struct { + toolOptVal string + receivedToolOpt bool +} +type toolOpt struct{ val string } + +func (a *assertTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{Name: "assert_tool"}, nil +} +func (a *assertTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + opt := tool.GetImplSpecificOptions(&toolOpt{}, opts...) + if opt.val == a.toolOptVal { + a.receivedToolOpt = true + } + return "tool-response", nil +} + +func TestAgentWithAllOptions(t *testing.T) { + ctx := context.Background() + ctrl := gomock.NewController(t) + + // Prepare a tool that asserts it receives the tool option + toolOptVal := "tool-opt-value" + to := tool.WrapImplSpecificOptFn(func(o *toolOpt) { o.val = toolOptVal }) + at := &assertTool{toolOptVal: toolOptVal} + + // Prepare a mock chat model that asserts it receives the model option + cm := mockModel.NewMockToolCallingChatModel(ctrl) + modelOpt := model.WithModel("test-model") + modelOptReceived := false + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, _ []*schema.Message, opts ...model.Option) (*schema.Message, error) { + times++ + if times == 1 { + for _, o := range opts { + opt := model.GetCommonOptions(&model.Options{}, o) + if opt.Model != nil && *opt.Model == "test-model" { + modelOptReceived = true + } + } + + info, _ := at.Info(ctx) + return schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: "", + }, + }, + }), + nil + } + + return schema.AssistantMessage("ok", nil), nil + }, + ).AnyTimes() + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + agentOpt := WithToolOptions(to) + agentOpt2 := WithChatModelOptions(modelOpt) + agentOpt3, err := WithTools(context.Background(), at) + assert.NoError(t, err) + + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&dummyBaseTool{}}, + }, + MaxStep: 20, + }) + assert.NoError(t, err) + + _, err = a.Generate(ctx, []*schema.Message{ + schema.UserMessage("call the tool"), + }, agentOpt, agentOpt2, agentOpt3[0], agentOpt3[1]) + assert.NoError(t, err) + assert.True(t, modelOptReceived, "model option should be received by chat model") + assert.True(t, at.receivedToolOpt, "tool option should be received by tool") +} + +type simpleToolForMiddlewareTest struct { + name string + result string +} + +func (s *simpleToolForMiddlewareTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: s.name, + Desc: "simple tool for middleware test", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "input": { + Desc: "input", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (s *simpleToolForMiddlewareTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return s.result, nil +} + +func (s *simpleToolForMiddlewareTest) StreamableRun(_ context.Context, _ string, _ ...tool.Option) (*schema.StreamReader[string], error) { + return schema.StreamReaderFromArray([]string{s.result}), nil +} + +func TestMessageFuture_ToolResultMiddleware_EmitsFinalResult(t *testing.T) { + originalResult := "original_result" + modifiedResult := "modified_by_middleware" + + resultModifyingMiddleware := compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + output, err := next(ctx, input) + if err != nil { + return nil, err + } + output.Result = modifiedResult + return output, nil + } + }, + Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + output, err := next(ctx, input) + if err != nil { + return nil, err + } + output.Result = schema.StreamReaderFromArray([]string{modifiedResult}) + return output, nil + } + }, + } + + t.Run("Invoke", func(t *testing.T) { + ctx := context.Background() + testTool := &simpleToolForMiddlewareTest{name: "test_tool", result: originalResult} + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + info, err := testTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"input": "test"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + option, future := WithMessageFuture() + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + ToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware}, + }, + MaxStep: 3, + }) + assert.NoError(t, err) + + response, err := a.Generate(ctx, []*schema.Message{ + schema.UserMessage("call the tool"), + }, option) + assert.NoError(t, err) + assert.Equal(t, "final response", response.Content) + + iter := future.GetMessages() + + var allMsgs []*schema.Message + for { + msg, hasNext, err := iter.Next() + if err != nil || !hasNext { + break + } + allMsgs = append(allMsgs, msg) + } + + assert.GreaterOrEqual(t, len(allMsgs), 3, "should have at least 3 messages") + if len(allMsgs) >= 3 { + assert.Equal(t, schema.Assistant, allMsgs[0].Role) + assert.Equal(t, 1, len(allMsgs[0].ToolCalls)) + + assert.Equal(t, schema.Tool, allMsgs[1].Role) + assert.Equal(t, modifiedResult, allMsgs[1].Content, + "MessageFuture should receive the middleware-modified tool result") + assert.NotEqual(t, originalResult, allMsgs[1].Content, + "MessageFuture should NOT receive the original tool result") + + assert.Equal(t, "final response", allMsgs[2].Content) + } + }) + + t.Run("Stream", func(t *testing.T) { + ctx := context.Background() + testTool := &simpleToolForMiddlewareTest{name: "test_tool_stream", result: originalResult} + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + info, err := testTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("", []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"input": "test"}`, + }, + }, + }), + }), nil). + Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("final response", nil), + }), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + option, future := WithMessageFuture() + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{testTool}, + ToolCallMiddlewares: []compose.ToolMiddleware{resultModifyingMiddleware}, + }, + MaxStep: 3, + }) + assert.NoError(t, err) + + response, err := a.Stream(ctx, []*schema.Message{ + schema.UserMessage("call the tool"), + }, option) + assert.NoError(t, err) + + var msgs []*schema.Message + for { + msg, err := response.Recv() + if err != nil { + break + } + msgs = append(msgs, msg) + } + finalMsg, err := schema.ConcatMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, "final response", finalMsg.Content) + + iter := future.GetMessageStreams() + + var allMsgs []*schema.Message + for { + msgStream, hasNext, err := iter.Next() + if err != nil || !hasNext { + break + } + var streamMsgs []*schema.Message + for { + msg, err := msgStream.Recv() + if err != nil { + break + } + streamMsgs = append(streamMsgs, msg) + } + if len(streamMsgs) > 0 { + concated, err := schema.ConcatMessages(streamMsgs) + if err == nil { + allMsgs = append(allMsgs, concated) + } + } + } + + assert.GreaterOrEqual(t, len(allMsgs), 3, "should have at least 3 messages") + if len(allMsgs) >= 3 { + assert.Equal(t, schema.Assistant, allMsgs[0].Role) + assert.Equal(t, 1, len(allMsgs[0].ToolCalls)) + + assert.Equal(t, schema.Tool, allMsgs[1].Role) + assert.Equal(t, modifiedResult, allMsgs[1].Content, + "MessageFuture should receive the middleware-modified tool result") + assert.NotEqual(t, originalResult, allMsgs[1].Content, + "MessageFuture should NOT receive the original tool result") + + assert.Equal(t, "final response", allMsgs[2].Content) + } + }) +} + +func TestWithMessageFuture_NestedGraph(t *testing.T) { + ctx := context.Background() + + t.Run("agent in nested graph", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agentGraph, agentGraphOpts := a.ExportGraph() + + parentGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = parentGraph.AddGraphNode("agent", agentGraph, agentGraphOpts...) + assert.NoError(t, err) + err = parentGraph.AddEdge(compose.START, "agent") + assert.NoError(t, err) + err = parentGraph.AddEdge("agent", compose.END) + assert.NoError(t, err) + + runnable, err := parentGraph.Compile(ctx) + assert.NoError(t, err) + + option, future := WithMessageFuture() + + response, err := runnable.Invoke(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, agent.GetComposeOptions(option)...) + assert.Nil(t, err) + assert.Equal(t, "final response", response.Content) + + iter := future.GetMessages() + + msg1, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + msg2, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Tool, msg2.Role) + + msg3, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, "final response", msg3.Content) + + _, hasNext, err = iter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + + t.Run("agent in deeply nested graph", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test user"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("final response", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agentGraph, agentGraphOpts := a.ExportGraph() + + childGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = childGraph.AddGraphNode("agent", agentGraph, agentGraphOpts...) + assert.NoError(t, err) + err = childGraph.AddEdge(compose.START, "agent") + assert.NoError(t, err) + err = childGraph.AddEdge("agent", compose.END) + assert.NoError(t, err) + + parentGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = parentGraph.AddGraphNode("child", childGraph) + assert.NoError(t, err) + err = parentGraph.AddEdge(compose.START, "child") + assert.NoError(t, err) + err = parentGraph.AddEdge("child", compose.END) + assert.NoError(t, err) + + runnable, err := parentGraph.Compile(ctx) + assert.NoError(t, err) + + option, future := WithMessageFuture() + + response, err := runnable.Invoke(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, agent.GetComposeOptions(option)...) + assert.Nil(t, err) + assert.Equal(t, "final response", response.Content) + + iter := future.GetMessages() + + msg1, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + msg2, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Tool, msg2.Role) + + msg3, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, "final response", msg3.Content) + + _, hasNext, err = iter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + + t.Run("agent in nested graph with streaming", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "stream user"}`, + }, + }, + }), + }), nil + }). + Times(1) + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderFromArray([]*schema.Message{ + schema.AssistantMessage("streaming", nil), + schema.AssistantMessage(" final", nil), + schema.AssistantMessage(" response", nil), + }), nil + }). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agentGraph, agentGraphOpts := a.ExportGraph() + + parentGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = parentGraph.AddGraphNode("agent", agentGraph, agentGraphOpts...) + assert.NoError(t, err) + err = parentGraph.AddEdge(compose.START, "agent") + assert.NoError(t, err) + err = parentGraph.AddEdge("agent", compose.END) + assert.NoError(t, err) + + runnable, err := parentGraph.Compile(ctx) + assert.NoError(t, err) + + option, future := WithMessageFuture() + + streamReader, err := runnable.Stream(ctx, []*schema.Message{ + schema.UserMessage("use the greet tool"), + }, agent.GetComposeOptions(option)...) + assert.Nil(t, err) + + var finalContent string + for { + chunk, err := streamReader.Recv() + if err == io.EOF { + break + } + assert.Nil(t, err) + finalContent += chunk.Content + } + assert.Contains(t, finalContent, "final") + + siter := future.GetMessageStreams() + + stream1, hasNext, err := siter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + msg1, err := concatStreamMsg(stream1) + assert.Nil(t, err) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 1, len(msg1.ToolCalls)) + + stream2, hasNext, err := siter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + msg2, err := concatStreamMsg(stream2) + assert.Nil(t, err) + assert.Equal(t, schema.Tool, msg2.Role) + + stream3, hasNext, err := siter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + msg3, err := concatStreamMsg(stream3) + assert.Nil(t, err) + assert.Contains(t, msg3.Content, "final") + + _, hasNext, err = siter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + + t.Run("agent with multiple tool calls in nested graph", func(t *testing.T) { + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "user1"}`, + }, + }, + { + ID: "tool-call-2", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "user2"}`, + }, + }, + }), nil). + Times(1) + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("greeted both users", nil), nil). + Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agentGraph, agentGraphOpts := a.ExportGraph() + + parentGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = parentGraph.AddGraphNode("agent", agentGraph, agentGraphOpts...) + assert.NoError(t, err) + err = parentGraph.AddEdge(compose.START, "agent") + assert.NoError(t, err) + err = parentGraph.AddEdge("agent", compose.END) + assert.NoError(t, err) + + runnable, err := parentGraph.Compile(ctx) + assert.NoError(t, err) + + option, future := WithMessageFuture() + + response, err := runnable.Invoke(ctx, []*schema.Message{ + schema.UserMessage("greet multiple users"), + }, agent.GetComposeOptions(option)...) + assert.Nil(t, err) + assert.Equal(t, "greeted both users", response.Content) + + iter := future.GetMessages() + + msg1, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Assistant, msg1.Role) + assert.Equal(t, 2, len(msg1.ToolCalls)) + + msg2, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Tool, msg2.Role) + + msg3, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, schema.Tool, msg3.Role) + + msg4, hasNext, err := iter.Next() + assert.Nil(t, err) + assert.True(t, hasNext) + assert.Equal(t, "greeted both users", msg4.Content) + + _, hasNext, err = iter.Next() + assert.Nil(t, err) + assert.False(t, hasNext) + }) + + t.Run("direct agent invoke vs nested graph - same behavior", func(t *testing.T) { + ctrl := gomock.NewController(t) + fakeTool := &fakeToolGreetForTest{} + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil). + Times(1) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + + a1, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm1, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + option1, future1 := WithMessageFuture() + _, err = a1.Generate(ctx, []*schema.Message{ + schema.UserMessage("test direct"), + }, option1) + assert.Nil(t, err) + + directMsgs := collectAllMessages(future1.GetMessages()) + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test"}`, + }, + }, + }), nil). + Times(1) + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("done", nil), nil). + Times(1) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + + a2, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm2, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agentGraph, agentGraphOpts := a2.ExportGraph() + + parentGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = parentGraph.AddGraphNode("agent", agentGraph, agentGraphOpts...) + assert.NoError(t, err) + err = parentGraph.AddEdge(compose.START, "agent") + assert.NoError(t, err) + err = parentGraph.AddEdge("agent", compose.END) + assert.NoError(t, err) + + runnable, err := parentGraph.Compile(ctx) + assert.NoError(t, err) + + option2, future2 := WithMessageFuture() + _, err = runnable.Invoke(ctx, []*schema.Message{ + schema.UserMessage("test nested"), + }, agent.GetComposeOptions(option2)...) + assert.Nil(t, err) + + nestedMsgs := collectAllMessages(future2.GetMessages()) + + assert.Equal(t, len(directMsgs), len(nestedMsgs), "should have same number of messages") + for i := range directMsgs { + assert.Equal(t, directMsgs[i].Role, nestedMsgs[i].Role, "message %d role should match", i) + } + }) + + t.Run("multiple react agents in graph - agent calls tool that invokes another agent", func(t *testing.T) { + ctrl := gomock.NewController(t) + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("inner agent response", nil), nil). + Times(1) + + innerAgent, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm2, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&fakeToolGreetForTest{}}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + innerAgentTool := &agentAsTool{ + agent: innerAgent, + name: "inner_agent", + desc: "An inner agent that can greet users", + } + + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("", + []schema.ToolCall{ + { + ID: "tool-call-1", + Function: schema.FunctionCall{ + Name: "inner_agent", + Arguments: `{"query": "hello"}`, + }, + }, + }), nil). + Times(1) + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("outer agent final response", nil), nil). + Times(1) + + outerAgent, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm1, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{innerAgentTool}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + option, future := WithMessageFuture() + + response, err := outerAgent.Generate(ctx, []*schema.Message{ + schema.UserMessage("use the inner agent"), + }, option) + assert.Nil(t, err) + assert.Equal(t, "outer agent final response", response.Content) + + allMsgs := collectAllMessages(future.GetMessages()) + + assert.GreaterOrEqual(t, len(allMsgs), 3, "should have at least 3 messages") + + hasOuterToolCall := false + hasInnerResponse := false + hasOuterFinalResponse := false + for _, msg := range allMsgs { + if msg.Role == schema.Assistant && len(msg.ToolCalls) > 0 { + hasOuterToolCall = true + } + if msg.Content == "inner agent response" { + hasInnerResponse = true + } + if msg.Content == "outer agent final response" { + hasOuterFinalResponse = true + } + } + assert.True(t, hasOuterToolCall, "should have outer agent tool call") + assert.True(t, hasInnerResponse, "should have inner agent response") + assert.True(t, hasOuterFinalResponse, "should have outer agent final response") + }) + + t.Run("two sequential react agents in same graph", func(t *testing.T) { + ctrl := gomock.NewController(t) + + cm1 := mockModel.NewMockToolCallingChatModel(ctrl) + cm1.EXPECT().WithTools(gomock.Any()).Return(cm1, nil).AnyTimes() + cm1.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("agent1 response", nil), nil). + Times(1) + + cm2 := mockModel.NewMockToolCallingChatModel(ctrl) + cm2.EXPECT().WithTools(gomock.Any()).Return(cm2, nil).AnyTimes() + cm2.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + Return(schema.AssistantMessage("agent2 response", nil), nil). + Times(1) + + agent1, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm1, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&fakeToolGreetForTest{}}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agent2, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm2, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&fakeToolGreetForTest{}}, + }, + MaxStep: 3, + }) + assert.Nil(t, err) + + agent1Graph, agent1Opts := agent1.ExportGraph() + agent2Graph, agent2Opts := agent2.ExportGraph() + + parentGraph := compose.NewGraph[[]*schema.Message, *schema.Message]() + err = parentGraph.AddGraphNode("agent1", agent1Graph, agent1Opts...) + assert.NoError(t, err) + + err = parentGraph.AddLambdaNode("transform", compose.InvokableLambda(func(ctx context.Context, msg *schema.Message) ([]*schema.Message, error) { + return []*schema.Message{schema.UserMessage("agent2 input: " + msg.Content)}, nil + })) + assert.NoError(t, err) + + err = parentGraph.AddGraphNode("agent2", agent2Graph, agent2Opts...) + assert.NoError(t, err) + + err = parentGraph.AddEdge(compose.START, "agent1") + assert.NoError(t, err) + err = parentGraph.AddEdge("agent1", "transform") + assert.NoError(t, err) + err = parentGraph.AddEdge("transform", "agent2") + assert.NoError(t, err) + err = parentGraph.AddEdge("agent2", compose.END) + assert.NoError(t, err) + + runnable, err := parentGraph.Compile(ctx) + assert.NoError(t, err) + + option, future := WithMessageFuture() + + response, err := runnable.Invoke(ctx, []*schema.Message{ + schema.UserMessage("hello"), + }, agent.GetComposeOptions(option)...) + assert.Nil(t, err) + assert.Equal(t, "agent2 response", response.Content) + + allMsgs := collectAllMessages(future.GetMessages()) + assert.Equal(t, 2, len(allMsgs)) + + responseContents := make(map[string]bool) + for _, msg := range allMsgs { + responseContents[msg.Content] = true + } + assert.True(t, responseContents["agent1 response"]) + assert.True(t, responseContents["agent2 response"]) + }) +} + +type agentAsTool struct { + agent *Agent + name string + desc string +} + +func (a *agentAsTool) Info(ctx context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: a.name, + Desc: a.desc, + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": { + Type: "string", + Desc: "The query to send to the inner agent", + }, + }), + }, nil +} + +func (a *agentAsTool) InvokableRun(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) { + response, err := a.agent.Generate(ctx, []*schema.Message{ + schema.UserMessage(argumentsInJSON), + }) + if err != nil { + return "", err + } + return response.Content, nil +} + +func collectAllMessages(iter *Iterator[*schema.Message]) []*schema.Message { + var msgs []*schema.Message + for { + msg, hasNext, err := iter.Next() + if err != nil || !hasNext { + break + } + msgs = append(msgs, msg) + } + return msgs +} + +func concatStreamMsg(sr *schema.StreamReader[*schema.Message]) (*schema.Message, error) { + var result *schema.Message + for { + chunk, err := sr.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + if result == nil { + result = chunk + } else { + result.Content += chunk.Content + if len(chunk.ToolCalls) > 0 { + result.ToolCalls = append(result.ToolCalls, chunk.ToolCalls...) + } + } + } + return result, nil +} diff --git a/flow/agent/react/react.go b/flow/agent/react/react.go new file mode 100644 index 0000000..794a41f --- /dev/null +++ b/flow/agent/react/react.go @@ -0,0 +1,492 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package react + +import ( + "context" + "io" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + "github.com/cloudwego/eino/schema" +) + +type toolResultSender func(toolName, callID, result string) + +type enhancedToolResultSender func(toolName, callID string, result *schema.ToolResult) +type streamToolResultSender func(toolName, callID string, resultStream *schema.StreamReader[string]) +type enhancedStreamToolResultSender func(toolName, callID string, resultStream *schema.StreamReader[*schema.ToolResult]) +type toolResultSenders struct { + sender toolResultSender + streamSender streamToolResultSender + + enhancedResultSender enhancedToolResultSender + enhancedStreamToolResultSender enhancedStreamToolResultSender +} + +type toolResultSenderCtxKey struct{} + +func setToolResultSendersToCtx(ctx context.Context, senders *toolResultSenders) context.Context { + return context.WithValue(ctx, toolResultSenderCtxKey{}, senders) +} + +func getToolResultSendersFromCtx(ctx context.Context) *toolResultSenders { + v := ctx.Value(toolResultSenderCtxKey{}) + if v == nil { + return nil + } + return v.(*toolResultSenders) +} + +type state struct { + Messages []*schema.Message + ReturnDirectlyToolCallID string +} + +func init() { + schema.RegisterName[*state]("_eino_react_state") +} + +func newToolResultCollectorMiddleware() compose.ToolMiddleware { + return compose.ToolMiddleware{ + Invokable: func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + senders := getToolResultSendersFromCtx(ctx) + output, err := next(ctx, input) + if err != nil { + return nil, err + } + if senders != nil && senders.sender != nil { + senders.sender(input.Name, input.CallID, output.Result) + } + return output, nil + } + }, + Streamable: func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + senders := getToolResultSendersFromCtx(ctx) + output, err := next(ctx, input) + if err != nil { + return nil, err + } + if senders != nil && senders.streamSender != nil { + streams := output.Result.Copy(2) + senders.streamSender(input.Name, input.CallID, streams[0]) + output.Result = streams[1] + } + return output, nil + } + }, + EnhancedInvokable: func(next compose.EnhancedInvokableToolEndpoint) compose.EnhancedInvokableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedInvokableToolOutput, error) { + senders := getToolResultSendersFromCtx(ctx) + output, err := next(ctx, input) + if err != nil { + return nil, err + } + if senders != nil && senders.enhancedResultSender != nil { + senders.enhancedResultSender(input.Name, input.CallID, output.Result) + } + return output, nil + + } + }, + EnhancedStreamable: func(next compose.EnhancedStreamableToolEndpoint) compose.EnhancedStreamableToolEndpoint { + return func(ctx context.Context, input *compose.ToolInput) (*compose.EnhancedStreamableToolOutput, error) { + senders := getToolResultSendersFromCtx(ctx) + output, err := next(ctx, input) + if err != nil { + return nil, err + } + if senders != nil && senders.enhancedStreamToolResultSender != nil { + streams := output.Result.Copy(2) + senders.enhancedStreamToolResultSender(input.Name, input.CallID, streams[0]) + output.Result = streams[1] + } + return output, nil + } + }, + } +} + +const ( + nodeKeyTools = "tools" + nodeKeyModel = "chat" +) + +// MessageModifier modify the input messages before the model is called. +type MessageModifier func(ctx context.Context, input []*schema.Message) []*schema.Message + +// AgentConfig is the config for ReAct agent. +type AgentConfig struct { + // ToolCallingModel is the chat model to be used for handling user messages with tool calling capability. + // This is the recommended model field to use. + ToolCallingModel model.ToolCallingChatModel + + // Deprecated: Use ToolCallingModel instead. + Model model.ChatModel + + // ToolsConfig is the config for tools node. + ToolsConfig compose.ToolsNodeConfig + + // MessageModifier. + // modify the input messages before the model is called, it's useful when you want to add some system prompt or other messages. + MessageModifier MessageModifier + + // MessageRewriter modifies message in the state, before the ChatModel is called. + // It takes the messages stored accumulated in state, modify them, and put the modified version back into state. + // Useful for compressing message history to fit the model context window, + // or if you want to make changes to messages that take effect across multiple model calls. + // NOTE: if both MessageModifier and MessageRewriter are set, MessageRewriter will be called before MessageModifier. + MessageRewriter MessageModifier + + // MaxStep. + // default 12 of steps in pregel (node num + 10). + MaxStep int `json:"max_step"` + + // Tools that will make agent return directly when the tool is called. + // When multiple tools are called and more than one tool is in the return directly list, only the first one will be returned. + ToolReturnDirectly map[string]struct{} + + // StreamToolCallChecker is a function to determine whether the model's streaming output contains tool calls. + // Different models have different ways of outputting tool calls in streaming mode: + // - Some models (like OpenAI) output tool calls directly + // - Others (like Claude) output text first, then tool calls + // This handler allows custom logic to check for tool calls in the stream. + // It should return: + // - true if the output contains tool calls and agent should continue processing + // - false if no tool calls and agent should stop + // Note: This field only needs to be configured when using streaming mode + // Note: The handler MUST close the modelOutput stream before returning + // Optional. By default, it checks if the first chunk contains tool calls. + // Note: The default implementation does not work well with Claude, which typically outputs tool calls after text content. + // Note: If your ChatModel doesn't output tool calls first, you can try adding prompts to constrain the model from generating extra text during the tool call. + StreamToolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) + + // GraphName is the graph name of the ReAct Agent. + // Optional. Default `ReActAgent`. + GraphName string + // ModelNodeName is the node name of the model node in the ReAct Agent graph. + // Optional. Default `ChatModel`. + ModelNodeName string + // ToolsNodeName is the node name of the tools node in the ReAct Agent graph. + // Optional. Default `Tools`. + ToolsNodeName string +} + +// NewPersonaModifier returns a MessageModifier that adds a persona message to the input. +// example: +// +// persona := "You are an expert in golang." +// config := AgentConfig{ +// Model: model, +// MessageModifier: NewPersonaModifier(persona), +// } +// agent, err := NewAgent(ctx, config) +// if err != nil {return} +// msg, err := agent.Generate(ctx, []*schema.Message{{Role: schema.User, Content: "how to build agent with eino"}}) +// if err != nil {return} +// println(msg.Content) +// +// Deprecated: Prefer directly including the persona message in the +// input when calling Generate or Stream to avoid extra copying. +func NewPersonaModifier(persona string) MessageModifier { + return func(ctx context.Context, input []*schema.Message) []*schema.Message { + res := make([]*schema.Message, 0, len(input)+1) + + res = append(res, schema.SystemMessage(persona)) + res = append(res, input...) + return res + } +} + +func firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) { + defer sr.Close() + + for { + msg, err := sr.Recv() + if err == io.EOF { + return false, nil + } + if err != nil { + return false, err + } + + if len(msg.ToolCalls) > 0 { + return true, nil + } + + if len(msg.Content) == 0 { // skip empty chunks at the front + continue + } + + return false, nil + } +} + +// Default graph and node names for the ReAct agent. +const ( + GraphName = "ReActAgent" + ModelNodeName = "ChatModel" + ToolsNodeName = "Tools" +) + +// SetReturnDirectly is a helper function that can be called within a tool's execution. +// It signals the ReAct agent to stop further processing and return the result of the current tool call directly. +// This is useful when the tool's output is the final answer and no more steps are needed. +// Note: If multiple tools call this function in the same step, only the last call will take effect. +// This setting has a higher priority than the AgentConfig.ToolReturnDirectly. +func SetReturnDirectly(ctx context.Context) error { + return compose.ProcessState(ctx, func(ctx context.Context, s *state) error { + s.ReturnDirectlyToolCallID = compose.GetToolCallID(ctx) + return nil + }) +} + +// Agent is the ReAct agent. +// ReAct agent is a simple agent that handles user messages with a chat model and tools. +// ReAct will call the chat model, if the message contains tool calls, it will call the tools. +// if the tool is configured to return directly, ReAct will return directly. +// otherwise, ReAct will continue to call the chat model until the message contains no tool calls. +// e.g. +// +// agent, err := ReAct.NewAgent(ctx, &react.AgentConfig{}) +// if err != nil {...} +// msg, err := agent.Generate(ctx, []*schema.Message{{Role: schema.User, Content: "how to build agent with eino"}}) +// if err != nil {...} +// println(msg.Content) +type Agent struct { + runnable compose.Runnable[[]*schema.Message, *schema.Message] + graph *compose.Graph[[]*schema.Message, *schema.Message] + graphAddNodeOpts []compose.GraphAddNodeOpt +} + +// NewAgent creates a ReAct agent that feeds tool response into next round of Chat Model generation. +// +// IMPORTANT!! For models that don't output tool calls in the first streaming chunk (e.g. Claude) +// the default StreamToolCallChecker may not work properly since it only checks the first chunk for tool calls. +// In such cases, you need to implement a custom StreamToolCallChecker that can properly detect tool calls. +func NewAgent(ctx context.Context, config *AgentConfig) (_ *Agent, err error) { + var ( + chatModel model.BaseChatModel + toolsNode *compose.ToolsNode + toolInfos []*schema.ToolInfo + toolCallChecker = config.StreamToolCallChecker + messageModifier = config.MessageModifier + ) + + graphName := GraphName + if config.GraphName != "" { + graphName = config.GraphName + } + + modelNodeName := ModelNodeName + if config.ModelNodeName != "" { + modelNodeName = config.ModelNodeName + } + + toolsNodeName := ToolsNodeName + if config.ToolsNodeName != "" { + toolsNodeName = config.ToolsNodeName + } + + if toolCallChecker == nil { + toolCallChecker = firstChunkStreamToolCallChecker + } + + if toolInfos, err = genToolInfos(ctx, config.ToolsConfig); err != nil { + return nil, err + } + + if chatModel, err = agent.ChatModelWithTools(config.Model, config.ToolCallingModel, toolInfos); err != nil { + return nil, err + } + + config.ToolsConfig.ToolCallMiddlewares = append( + []compose.ToolMiddleware{newToolResultCollectorMiddleware()}, + config.ToolsConfig.ToolCallMiddlewares..., + ) + + if toolsNode, err = compose.NewToolNode(ctx, &config.ToolsConfig); err != nil { + return nil, err + } + + graph := compose.NewGraph[[]*schema.Message, *schema.Message](compose.WithGenLocalState(func(ctx context.Context) *state { + return &state{Messages: make([]*schema.Message, 0, config.MaxStep+1)} + })) + + modelPreHandle := func(ctx context.Context, input []*schema.Message, state *state) ([]*schema.Message, error) { + state.Messages = append(state.Messages, input...) + + if config.MessageRewriter != nil { + state.Messages = config.MessageRewriter(ctx, state.Messages) + } + + if messageModifier == nil { + return state.Messages, nil + } + + modifiedInput := make([]*schema.Message, len(state.Messages)) + copy(modifiedInput, state.Messages) + return messageModifier(ctx, modifiedInput), nil + } + + if err = graph.AddChatModelNode(nodeKeyModel, chatModel, compose.WithStatePreHandler(modelPreHandle), compose.WithNodeName(modelNodeName)); err != nil { + return nil, err + } + + if err = graph.AddEdge(compose.START, nodeKeyModel); err != nil { + return nil, err + } + + toolsNodePreHandle := func(ctx context.Context, input *schema.Message, state *state) (*schema.Message, error) { + if input == nil { + return state.Messages[len(state.Messages)-1], nil // used for rerun interrupt resume + } + state.Messages = append(state.Messages, input) + state.ReturnDirectlyToolCallID = getReturnDirectlyToolCallID(input, config.ToolReturnDirectly) + return input, nil + } + if err = graph.AddToolsNode(nodeKeyTools, toolsNode, compose.WithStatePreHandler(toolsNodePreHandle), compose.WithNodeName(toolsNodeName)); err != nil { + return nil, err + } + + modelPostBranchCondition := func(ctx context.Context, sr *schema.StreamReader[*schema.Message]) (endNode string, err error) { + if isToolCall, err := toolCallChecker(ctx, sr); err != nil { + return "", err + } else if isToolCall { + return nodeKeyTools, nil + } + return compose.END, nil + } + + if err = graph.AddBranch(nodeKeyModel, compose.NewStreamGraphBranch(modelPostBranchCondition, map[string]bool{nodeKeyTools: true, compose.END: true})); err != nil { + return nil, err + } + + if err = buildReturnDirectly(graph); err != nil { + return nil, err + } + + compileOpts := []compose.GraphCompileOption{compose.WithMaxRunSteps(config.MaxStep), compose.WithNodeTriggerMode(compose.AnyPredecessor), compose.WithGraphName(graphName)} + runnable, err := graph.Compile(ctx, compileOpts...) + if err != nil { + return nil, err + } + + return &Agent{ + runnable: runnable, + graph: graph, + graphAddNodeOpts: []compose.GraphAddNodeOpt{compose.WithGraphCompileOptions(compileOpts...)}, + }, nil +} + +func buildReturnDirectly(graph *compose.Graph[[]*schema.Message, *schema.Message]) (err error) { + directReturn := func(ctx context.Context, msgs *schema.StreamReader[[]*schema.Message]) (*schema.StreamReader[*schema.Message], error) { + return schema.StreamReaderWithConvert(msgs, func(msgs []*schema.Message) (*schema.Message, error) { + var msg *schema.Message + err = compose.ProcessState[*state](ctx, func(_ context.Context, state *state) error { + for i := range msgs { + if msgs[i] != nil && msgs[i].ToolCallID == state.ReturnDirectlyToolCallID { + msg = msgs[i] + return nil + } + } + return nil + }) + if err != nil { + return nil, err + } + if msg == nil { + return nil, schema.ErrNoValue + } + return msg, nil + }), nil + } + + nodeKeyDirectReturn := "direct_return" + if err = graph.AddLambdaNode(nodeKeyDirectReturn, compose.TransformableLambda(directReturn)); err != nil { + return err + } + + // this branch checks if the tool called should return directly. It either leads to END or back to ChatModel + err = graph.AddBranch(nodeKeyTools, compose.NewStreamGraphBranch(func(ctx context.Context, msgsStream *schema.StreamReader[[]*schema.Message]) (endNode string, err error) { + msgsStream.Close() + + err = compose.ProcessState[*state](ctx, func(_ context.Context, state *state) error { + if len(state.ReturnDirectlyToolCallID) > 0 { + endNode = nodeKeyDirectReturn + } else { + endNode = nodeKeyModel + } + return nil + }) + if err != nil { + return "", err + } + return endNode, nil + }, map[string]bool{nodeKeyModel: true, nodeKeyDirectReturn: true})) + if err != nil { + return err + } + + return graph.AddEdge(nodeKeyDirectReturn, compose.END) +} + +func genToolInfos(ctx context.Context, config compose.ToolsNodeConfig) ([]*schema.ToolInfo, error) { + toolInfos := make([]*schema.ToolInfo, 0, len(config.Tools)) + for _, t := range config.Tools { + tl, err := t.Info(ctx) + if err != nil { + return nil, err + } + + toolInfos = append(toolInfos, tl) + } + + return toolInfos, nil +} + +func getReturnDirectlyToolCallID(input *schema.Message, toolReturnDirectly map[string]struct{}) string { + if len(toolReturnDirectly) == 0 { + return "" + } + + for _, toolCall := range input.ToolCalls { + if _, ok := toolReturnDirectly[toolCall.Function.Name]; ok { + return toolCall.ID + } + } + + return "" +} + +// Generate generates a response from the agent. +func (r *Agent) Generate(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (*schema.Message, error) { + return r.runnable.Invoke(ctx, input, agent.GetComposeOptions(opts...)...) +} + +// Stream calls the agent and returns a stream response. +func (r *Agent) Stream(ctx context.Context, input []*schema.Message, opts ...agent.AgentOption) (output *schema.StreamReader[*schema.Message], err error) { + return r.runnable.Stream(ctx, input, agent.GetComposeOptions(opts...)...) +} + +// ExportGraph exports the underlying graph from Agent, along with the []compose.GraphAddNodeOpt to be used when adding this graph to another graph. +func (r *Agent) ExportGraph() (compose.AnyGraph, []compose.GraphAddNodeOpt) { + return r.graph, r.graphAddNodeOpts +} diff --git a/flow/agent/react/react_test.go b/flow/agent/react/react_test.go new file mode 100644 index 0000000..175cedc --- /dev/null +++ b/flow/agent/react/react_test.go @@ -0,0 +1,822 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package react + +import ( + "context" + "errors" + "fmt" + "io" + "math/rand" + "testing" + + "github.com/bytedance/sonic" + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/agent" + mockModel "github.com/cloudwego/eino/internal/mock/components/model" + "github.com/cloudwego/eino/schema" + template "github.com/cloudwego/eino/utils/callbacks" +) + +func TestReact(t *testing.T) { + ctx := context.Background() + + fakeTool := &fakeToolGreetForTest{ + tarCount: 3, + } + + info, err := fakeTool.Info(ctx) + assert.NoError(t, err) + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockChatModel(ctrl) + + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + times++ + if times <= 2 { + info, _ := fakeTool.Info(ctx) + + return schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()), + }, + }, + }), + nil + } + + return schema.AssistantMessage("bye", nil), nil + }).AnyTimes() + cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message { + assert.Equal(t, len(input), times*2+1) + return input + }, + MaxStep: 40, + }) + assert.Nil(t, err) + + out, err := a.Generate(ctx, []*schema.Message{ + { + Role: schema.User, + Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!", + }, + }, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest))) + assert.Nil(t, err) + + if out != nil { + t.Log(out.Content) + } + + // test return directly + times = 0 + a, err = NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message { + assert.Equal(t, len(input), times*2+1) + return input + }, + MaxStep: 40, + ToolReturnDirectly: map[string]struct{}{info.Name: {}}, + }) + assert.Nil(t, err) + + out, err = a.Generate(ctx, []*schema.Message{ + { + Role: schema.User, + Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!", + }, + }, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest))) + assert.Nil(t, err) + + if out != nil { + t.Log(out.Content) + } +} + +func TestReactWithMessageRewriterAndModifier(t *testing.T) { + ctx := context.Background() + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + // This test simulates a single Generate call with a long history. + // The MessageRewriter should shorten the history. + // The MessageModifier should add a system prompt. + + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + // Check messages passed to the model. + // Expected: [system prompt, user: "message 2", assistant: "response 2"] + assert.Len(t, input, 3) + assert.Equal(t, schema.System, input[0].Role) + assert.Equal(t, "system prompt", input[0].Content) + assert.Equal(t, schema.User, input[1].Role) + assert.Equal(t, "message 2", input[1].Content) + assert.Equal(t, schema.Assistant, input[2].Role) + assert.Equal(t, "response 2", input[2].Content) + return schema.AssistantMessage("final response", nil), nil + }).Times(1) + cm.EXPECT().WithTools(gomock.Any()).Return(cm, nil).AnyTimes() + + ra, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + MessageRewriter: func(ctx context.Context, messages []*schema.Message) []*schema.Message { + // Keep only the last 2 messages if history is longer. + assert.Len(t, messages, 4) // user1, assistant1, user2, assistant2 + if len(messages) > 2 { + return messages[len(messages)-2:] + } + return messages + }, + MessageModifier: func(ctx context.Context, messages []*schema.Message) []*schema.Message { + // messages should be the result from rewriter + assert.Len(t, messages, 2) // user2, assistant2 + + // Add a system prompt + res := make([]*schema.Message, 0, len(messages)+1) + res = append(res, schema.SystemMessage("system prompt")) + res = append(res, messages...) + return res + }, + }) + assert.NoError(t, err) + + // Simulate a conversation history + history := []*schema.Message{ + schema.UserMessage("message 1"), + schema.AssistantMessage("response 1", nil), + schema.UserMessage("message 2"), + schema.AssistantMessage("response 2", nil), + } + + // Run the react agent + finalMsg, err := ra.Generate(ctx, history) + assert.NoError(t, err) + assert.Equal(t, "final response", finalMsg.Content) +} + +func TestReactStream(t *testing.T) { + ctx := context.Background() + + fakeTool := &fakeToolGreetForTest{ + tarCount: 20, + } + + fakeStreamTool := &fakeStreamToolGreetForTest{ + tarCount: 20, + } + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockChatModel(ctrl) + + times := 0 + cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes() + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) ( + *schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + defer sw.Close() + + info, _ := fakeTool.Info(ctx) + streamInfo, _ := fakeStreamTool.Info(ctx) + + times++ + if times <= 2 { + sw.Send(schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "tool"}`, randStr()), + }, + }, + }), + nil) + return sr, nil + } else if times == 3 { + sw.Send(schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: streamInfo.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "stream tool"}`, randStr()), + }, + }, + }), + nil) + return sr, nil + } else if times == 4 { // parallel tool call + sw.Send(schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "tool"}`, randStr()), + }, + }, + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: streamInfo.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "stream tool"}`, randStr()), + }, + }, + }), + nil) + return sr, nil + } + + sw.Send(schema.AssistantMessage("bye", nil), nil) + return sr, nil + }).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool, fakeStreamTool}, + }, + + MaxStep: 40, + }) + assert.Nil(t, err) + + out, err := a.Stream(ctx, []*schema.Message{ + { + Role: schema.User, + Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!", + }, + }, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest))) + if err != nil { + t.Fatal(err) + } + + defer out.Close() + + msgs := make([]*schema.Message, 0) + for { + msg, err := out.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + + t.Fatal(err) + } + + msgs = append(msgs, msg) + } + + assert.Equal(t, 1, len(msgs)) + + msg, err := schema.ConcatMessages(msgs) + if err != nil { + t.Fatal(err) + } + + t.Log(msg.Content) + + info, err := fakeStreamTool.Info(ctx) + assert.NoError(t, err) + + // test return directly + a, err = NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool, fakeStreamTool}, + }, + + MaxStep: 40, + ToolReturnDirectly: map[string]struct{}{info.Name: {}}, // one of the two tools is return directly + }) + assert.Nil(t, err) + + times = 0 + out, err = a.Stream(ctx, []*schema.Message{ + { + Role: schema.User, + Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!", + }, + }, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest))) + if err != nil { + t.Fatal(err) + } + + defer out.Close() + + msgs = make([]*schema.Message, 0) + for { + msg, err := out.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + + t.Fatal(err) + } + + msgs = append(msgs, msg) + } + + assert.Equal(t, 1, len(msgs)) + + msg, err = schema.ConcatMessages(msgs) + if err != nil { + t.Fatal(err) + } + + t.Log(msg.Content) + + // return directly tool call within parallel tool calls + out, err = a.Stream(ctx, []*schema.Message{ + { + Role: schema.User, + Content: "Use greet tool to continuously say hello until you get a bye response, greet names in the following order: max, bob, alice, john, marry, joe, ken, lily, please start directly! please start directly! please start directly!", + }, + }, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest))) + assert.NoError(t, err) + + defer out.Close() + + msgs = make([]*schema.Message, 0) + for { + msg, err := out.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + assert.NoError(t, err) + } + + msgs = append(msgs, msg) + } + + assert.Equal(t, 1, len(msgs)) + + msg, err = schema.ConcatMessages(msgs) + assert.NoError(t, err) + + t.Log("parallel tool call with return directly: ", msg.Content) +} + +func TestReactWithModifier(t *testing.T) { + ctx := context.Background() + + fakeTool := &fakeToolGreetForTest{} + ctrl := gomock.NewController(t) + cm := mockModel.NewMockChatModel(ctrl) + + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + times++ + if times <= 2 { + info, _ := fakeTool.Info(ctx) + + return schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()), + }, + }, + }), + nil + } + + return schema.AssistantMessage("bye", nil), nil + }).AnyTimes() + cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool}, + }, + MessageModifier: func(ctx context.Context, input []*schema.Message) []*schema.Message { + res := make([]*schema.Message, 0, len(input)+1) + + res = append(res, schema.SystemMessage("you are a helpful assistant")) + res = append(res, input...) + return res + }, + + MaxStep: 40, + }) + + assert.Nil(t, err) + + out, err := a.Generate(ctx, []*schema.Message{ + { + Role: schema.User, + Content: "hello", + }, + }, agent.WithComposeOptions(compose.WithCallbacks(callbackForTest))) + if err != nil { + t.Fatal(err) + } + + if out != nil { + t.Log(out.Content) + } +} + +func TestAgentInGraph(t *testing.T) { + t.Run("agent generate in chain", func(t *testing.T) { + ctx := context.Background() + + fakeTool := &fakeToolGreetForTest{} + ctrl := gomock.NewController(t) + cm := mockModel.NewMockChatModel(ctrl) + + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + + times += 1 + if times <= 2 { + info, _ := fakeTool.Info(ctx) + + return schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()), + }, + }, + }), + nil + } + + return schema.AssistantMessage("bye", nil), nil + + }).Times(3) + cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{fakeTool, &fakeStreamToolGreetForTest{}}, + }, + + MaxStep: 40, + }) + assert.Nil(t, err) + + chain := compose.NewChain[[]*schema.Message, string]() + agentLambda, err := compose.AnyLambda(a.Generate, a.Stream, nil, nil) + assert.Nil(t, err) + + chain. + AppendLambda(agentLambda). + AppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) { + t.Log("got agent response: ", input.Content) + return input.Content, nil + })) + r, err := chain.Compile(ctx) + assert.Nil(t, err) + + res, err := r.Invoke(ctx, []*schema.Message{{Role: schema.User, Content: "hello"}}, + compose.WithCallbacks(callbackForTest)) + assert.Nil(t, err) + + t.Log(res) + }) + + t.Run("agent stream in chain", func(t *testing.T) { + + fakeStreamTool := &fakeStreamToolGreetForTest{} + ctx := context.Background() + ctrl := gomock.NewController(t) + cm := mockModel.NewMockChatModel(ctrl) + + times := 0 + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) ( + *schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + defer sw.Close() + + times += 1 + if times <= 2 { + info, _ := fakeStreamTool.Info(ctx) + sw.Send(schema.AssistantMessage("hello max", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: fmt.Sprintf(`{"name": "%s", "hh": "123"}`, randStr()), + }, + }, + }), + nil) + return sr, nil + } + + sw.Send(schema.AssistantMessage("bye", nil), nil) + return sr, nil + }).Times(3) + cm.EXPECT().BindTools(gomock.Any()).Return(nil).AnyTimes() + + a, err := NewAgent(ctx, &AgentConfig{ + Model: cm, + ToolsConfig: compose.ToolsNodeConfig{ + Tools: []tool.BaseTool{&fakeToolGreetForTest{}, fakeStreamTool}, + }, + + MaxStep: 40, + }) + assert.Nil(t, err) + + chain := compose.NewChain[[]*schema.Message, string]() + agentGraph, opts := a.ExportGraph() + assert.Nil(t, err) + + chain. + AppendGraph(agentGraph, opts...). + AppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) { + t.Log("got agent response: ", input.Content) + return input.Content, nil + })) + r, err := chain.Compile(ctx) + assert.Nil(t, err) + + outStream, err := r.Stream(ctx, []*schema.Message{{Role: schema.User, Content: "hello"}}, + compose.WithCallbacks(callbackForTest)) + if err != nil { + t.Fatal(err) + } + + defer outStream.Close() + + msg := "" + for { + msgItem, err := outStream.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + + t.Fatal(err) + } + + msg += msgItem + } + + t.Log(msg) + }) + +} + +func TestWithTools(t *testing.T) { + ctx := context.Background() + + fakeTool := &fakeToolGreetForTest{ + tarCount: 2, + } + fakeStreamTool := &fakeStreamToolGreetForTest{ + tarCount: 2, + } + + ctrl := gomock.NewController(t) + cm := mockModel.NewMockToolCallingChatModel(ctrl) + + times := 0 + cm.EXPECT().Generate(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + times++ + if times <= 1 { + info, _ := fakeTool.Info(ctx) + return schema.AssistantMessage("calling tool", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test"}`, + }, + }, + }), + nil + } + return schema.AssistantMessage("done", nil), nil + }).AnyTimes() + + cm.EXPECT().Stream(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, input []*schema.Message, opts ...model.Option) ( + *schema.StreamReader[*schema.Message], error) { + sr, sw := schema.Pipe[*schema.Message](1) + defer sw.Close() + + times++ + if times <= 2 { + info, _ := fakeStreamTool.Info(ctx) + sw.Send(schema.AssistantMessage("calling stream tool", + []schema.ToolCall{ + { + ID: randStr(), + Function: schema.FunctionCall{ + Name: info.Name, + Arguments: `{"name": "test"}`, + }, + }, + }), + nil) + return sr, nil + } + + sw.Send(schema.AssistantMessage("stream done", nil), nil) + return sr, nil + }).AnyTimes() + + // Test WithTools function + toolOptions, err := WithTools(ctx, fakeTool, fakeStreamTool) + assert.NoError(t, err) + assert.Len(t, toolOptions, 2, "WithTools should return exactly 2 options") + + // Create agent without tools in config + a, err := NewAgent(ctx, &AgentConfig{ + ToolCallingModel: cm, + MaxStep: 10, + }) + assert.NoError(t, err) + + // Test Generate with WithTools options + times = 0 + msg, err := a.Generate(ctx, []*schema.Message{ + schema.UserMessage("test generate with tools"), + }, toolOptions...) + assert.NoError(t, err) + assert.Equal(t, "done", msg.Content) + + // Test Stream with WithTools options + times = 0 + stream, err := a.Stream(ctx, []*schema.Message{ + schema.UserMessage("test stream with tools"), + }, toolOptions...) + assert.NoError(t, err) + + defer stream.Close() + msgs := make([]*schema.Message, 0) + for { + msg, err := stream.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + assert.NoError(t, err) + } + msgs = append(msgs, msg) + } + + assert.Len(t, msgs, 1) + concatMsg, err := schema.ConcatMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, "stream done", concatMsg.Content) + + // Test error case - tool Info() returns error + errorTool := &errorToolForTest{} + _, err = WithTools(ctx, errorTool) + assert.Error(t, err) + assert.Contains(t, err.Error(), "info error") +} + +// Helper tool for testing error cases +type errorToolForTest struct{} + +func (t *errorToolForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return nil, errors.New("info error") +} + +func (t *errorToolForTest) InvokableRun(_ context.Context, _ string, _ ...tool.Option) (string, error) { + return "", nil +} + +type fakeStreamToolGreetForTest struct { + tarCount int + curCount int +} + +func (t *fakeStreamToolGreetForTest) StreamableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) ( + *schema.StreamReader[string], error) { + p := &fakeToolInput{} + err := sonic.UnmarshalString(argumentsInJSON, p) + if err != nil { + return nil, err + } + + if t.curCount >= t.tarCount { + s := schema.StreamReaderFromArray([]string{`{"say": "bye"}`}) + return s, nil + } + t.curCount++ + s := schema.StreamReaderFromArray([]string{fmt.Sprintf(`{"say": "hello %v"}`, p.Name)}) + return s, nil +} + +type fakeToolGreetForTest struct { + tarCount int + curCount int +} + +func (t *fakeToolGreetForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "greet", + Desc: "greet with name", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Desc: "user name who to greet", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (t *fakeStreamToolGreetForTest) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "greet in stream", + Desc: "greet with name in stream", + ParamsOneOf: schema.NewParamsOneOfByParams( + map[string]*schema.ParameterInfo{ + "name": { + Desc: "user name who to greet", + Required: true, + Type: schema.String, + }, + }), + }, nil +} + +func (t *fakeToolGreetForTest) InvokableRun(_ context.Context, argumentsInJSON string, _ ...tool.Option) (string, error) { + p := &fakeToolInput{} + err := sonic.UnmarshalString(argumentsInJSON, p) + if err != nil { + return "", err + } + + if t.curCount >= t.tarCount { + return `{"say": "bye"}`, nil + } + + t.curCount++ + return fmt.Sprintf(`{"say": "hello %v"}`, p.Name), nil +} + +type fakeToolInput struct { + Name string `json:"name"` +} + +func randStr() string { + seeds := []rune("this is a seed") + b := make([]rune, 8) + for i := range b { + b[i] = seeds[rand.Intn(len(seeds))] + } + return string(b) +} + +var callbackForTest = BuildAgentCallback(&template.ModelCallbackHandler{}, &template.ToolCallbackHandler{}) diff --git a/flow/agent/utils.go b/flow/agent/utils.go new file mode 100644 index 0000000..d8100bb --- /dev/null +++ b/flow/agent/utils.go @@ -0,0 +1,52 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package agent + +import ( + "errors" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +// ChatModelWithTools returns a chat model configured with tool schemas. +// If a ToolCallingChatModel is provided, it is used directly (and optionally +// configured with tools). Otherwise, a plain ChatModel is bound with tools. +func ChatModelWithTools(cm model.ChatModel, toolCallingModel model.ToolCallingChatModel, toolInfos []*schema.ToolInfo) ( + model.BaseChatModel, error) { + + if toolCallingModel != nil { + if len(toolInfos) == 0 { + return toolCallingModel, nil + } + return toolCallingModel.WithTools(toolInfos) + } + + if cm != nil { + if len(toolInfos) == 0 { + return cm, nil + } + err := cm.BindTools(toolInfos) + if err != nil { + return nil, err + } + + return cm, nil + } + + return nil, errors.New("no chat model provided") +} diff --git a/flow/indexer/parent/parent.go b/flow/indexer/parent/parent.go new file mode 100644 index 0000000..aff63b7 --- /dev/null +++ b/flow/indexer/parent/parent.go @@ -0,0 +1,162 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package parent provides an indexer that assigns stable IDs to sub-documents +// and preserves relationships to their original parent document. +package parent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/schema" +) + +// Config configures the parent indexer that assigns IDs to sub-documents. +type Config struct { + // Indexer is the underlying indexer implementation that handles the actual document indexing. + // For example: a vector database indexer like Milvus, or a full-text search indexer like Elasticsearch. + Indexer indexer.Indexer + + // Transformer processes documents before indexing, typically splitting them into smaller chunks. + // Each sub-document generated by the transformer must retain its parent document's ID. + // For example: if a document with ID "doc_1" is split into 3 chunks, all chunks will initially + // have ID "doc_1". These IDs will later be modified by the SubIDGenerator. + // + // Example transformations: + // - A text splitter that breaks down large documents into paragraphs + // - A code splitter that separates code files into functions + Transformer document.Transformer + + // ParentIDKey specifies the metadata key used to store the original document's ID in each sub-document. + // For example: if ParentIDKey is "parent_id", each sub-document will have metadata like: + // {"parent_id": "original_doc_123"} + ParentIDKey string + + // SubIDGenerator generates unique IDs for sub-documents based on their parent document ID. + // For example: if parent ID is "doc_1" and we need 3 sub-document IDs, it might generate: + // ["doc_1_chunk_1", "doc_1_chunk_2", "doc_1_chunk_3"] + // + // Parameters: + // - ctx: context for the operation + // - parentID: the ID of the parent document + // - num: number of sub-document IDs needed + // Returns: + // - []string: slice of generated sub-document IDs + // - error: any error encountered during ID generation + SubIDGenerator func(ctx context.Context, parentID string, num int) ([]string, error) +} + +// NewIndexer creates a new parent indexer that handles document splitting and sub-document management. +// +// Parameters: +// - ctx: context for the operation +// - config: configuration for the parent indexer +// +// Example usage: +// +// indexer, err := NewIndexer(ctx, &Config{ +// Indexer: milvusIndexer, +// Transformer: textSplitter, +// ParentIDKey: "source_doc_id", +// SubIDGenerator: func(ctx context.Context, parentID string, num int) ([]string, error) { +// ids := make([]string, num) +// for i := 0; i < num; i++ { +// ids[i] = fmt.Sprintf("%s_chunk_%d", parentID, i+1) +// } +// return ids, nil +// }, +// }) +// +// Returns: +// - indexer.Indexer: the created parent indexer +// - error: any error encountered during creation +func NewIndexer(ctx context.Context, config *Config) (indexer.Indexer, error) { + if config.Indexer == nil { + return nil, fmt.Errorf("indexer is empty") + } + if config.Transformer == nil { + return nil, fmt.Errorf("transformer is empty") + } + if config.SubIDGenerator == nil { + return nil, fmt.Errorf("sub id generator is empty") + } + + return &parentIndexer{ + indexer: config.Indexer, + transformer: config.Transformer, + parentIDKey: config.ParentIDKey, + subIDGenerator: config.SubIDGenerator, + }, nil +} + +type parentIndexer struct { + indexer indexer.Indexer + transformer document.Transformer + parentIDKey string + subIDGenerator func(ctx context.Context, parentID string, num int) ([]string, error) +} + +func (p *parentIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) ([]string, error) { + subDocs, err := p.transformer.Transform(ctx, docs) + if err != nil { + return nil, fmt.Errorf("transform docs fail: %w", err) + } + if len(subDocs) == 0 { + return nil, fmt.Errorf("doc transformer returned no documents") + } + currentID := subDocs[0].ID + startIdx := 0 + for i, subDoc := range subDocs { + if subDoc.MetaData == nil { + subDoc.MetaData = make(map[string]any) + } + subDoc.MetaData[p.parentIDKey] = subDoc.ID + + if subDoc.ID == currentID { + continue + } + + // generate new doc id + subIDs, err_ := p.subIDGenerator(ctx, subDocs[startIdx].ID, i-startIdx) + if err_ != nil { + return nil, err_ + } + if len(subIDs) != i-startIdx { + return nil, fmt.Errorf("generated sub IDs' num is unexpected") + } + for j := startIdx; j < i; j++ { + subDocs[j].ID = subIDs[j-startIdx] + } + startIdx = i + currentID = subDoc.ID + } + // generate new doc id + subIDs, err := p.subIDGenerator(ctx, subDocs[startIdx].ID, len(subDocs)-startIdx) + if err != nil { + return nil, err + } + if len(subIDs) != len(subDocs)-startIdx { + return nil, fmt.Errorf("generated sub IDs' num is unexpected") + } + for j := startIdx; j < len(subDocs); j++ { + subDocs[j].ID = subIDs[j-startIdx] + } + + return p.indexer.Store(ctx, subDocs, opts...) +} diff --git a/flow/indexer/parent/parent_test.go b/flow/indexer/parent/parent_test.go new file mode 100644 index 0000000..82b3621 --- /dev/null +++ b/flow/indexer/parent/parent_test.go @@ -0,0 +1,121 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parent + +import ( + "context" + "fmt" + "reflect" + "strconv" + "strings" + "testing" + + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/schema" +) + +type testIndexer struct{} + +func (t *testIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) (ids []string, err error) { + ret := make([]string, len(docs)) + for i, d := range docs { + ret[i] = d.ID + if !strings.HasPrefix(d.ID, d.MetaData["parent"].(string)) { + return nil, fmt.Errorf("invalid parent key") + } + } + return ret, nil +} + +type testTransformer struct { +} + +func (t *testTransformer) Transform(ctx context.Context, src []*schema.Document, opts ...document.TransformerOption) ([]*schema.Document, error) { + var ret []*schema.Document + for _, d := range src { + ret = append(ret, &schema.Document{ + ID: d.ID, + Content: d.Content[:len(d.Content)/2], + MetaData: deepCopyMap(d.MetaData), + }, &schema.Document{ + ID: d.ID, + Content: d.Content[len(d.Content)/2:], + MetaData: deepCopyMap(d.MetaData), + }) + } + return ret, nil +} + +func TestParentIndexer(t *testing.T) { + tests := []struct { + name string + config *Config + input []*schema.Document + want []string + }{ + { + name: "success", + config: &Config{ + Indexer: &testIndexer{}, + Transformer: &testTransformer{}, + ParentIDKey: "parent", + SubIDGenerator: func(ctx context.Context, parentID string, num int) ([]string, error) { + ret := make([]string, num) + for i := range ret { + ret[i] = parentID + strconv.Itoa(i) + } + return ret, nil + }, + }, + input: []*schema.Document{{ + ID: "id", + Content: "1234567890", + MetaData: map[string]interface{}{}, + }, { + ID: "ID", + Content: "0987654321", + MetaData: map[string]interface{}{}, + }}, + want: []string{"id0", "id1", "ID0", "ID1"}, + }, + } + ctx := context.Background() + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + index, err := NewIndexer(ctx, tt.config) + if err != nil { + t.Fatal(err) + } + ret, err := index.Store(ctx, tt.input) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(ret, tt.want) { + t.Errorf("NewHeaderSplitter() got = %v, want %v", ret, tt.want) + } + }) + } +} + +func deepCopyMap(in map[string]interface{}) map[string]interface{} { + out := make(map[string]interface{}) + for k, v := range in { + out[k] = v + } + return out +} diff --git a/flow/retriever/multiquery/multi_query.go b/flow/retriever/multiquery/multi_query.go new file mode 100644 index 0000000..1d064cd --- /dev/null +++ b/flow/retriever/multiquery/multi_query.go @@ -0,0 +1,211 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package multiquery implements a query-rewriting retriever that expands +// user queries into multiple variants to improve recall. +package multiquery + +import ( + "context" + "fmt" + "strings" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/retriever/utils" + "github.com/cloudwego/eino/schema" +) + +const ( + defaultRewritePrompt = `You are an helpful assistant. + Your role is to create three different versions of the user query to retrieve relevant documents from store. + Your goal is to improve the performance of similarity search by generating text from different perspectives based on the user query. + Only provide the generated queries and separate them by newlines. + user query: {{query}}` + defaultQueryVariable = "query" + defaultMaxQueriesNum = 5 +) + +var deduplicateFusion = func(ctx context.Context, docs [][]*schema.Document) ([]*schema.Document, error) { + m := map[string]bool{} + var ret []*schema.Document + for i := range docs { + for j := range docs[i] { + if _, ok := m[docs[i][j].ID]; !ok { + m[docs[i][j].ID] = true + ret = append(ret, docs[i][j]) + } + } + } + return ret, nil +} + +// NewRetriever creates a multi-query retriever. +// multi-query retriever is useful when you want to retrieve documents from multiple retrievers with different queries. +// e.g. +// +// multiRetriever := multiquery.NewRetriever(ctx, &multiquery.Config{}) +// docs, err := multiRetriever.Retrieve(ctx, "how to build agent with eino") +// if err != nil { +// ... +// } +// println(docs) +func NewRetriever(ctx context.Context, config *Config) (retriever.Retriever, error) { + var err error + + // config validate + if config.OrigRetriever == nil { + return nil, fmt.Errorf("OrigRetriever is required") + } + if config.RewriteHandler == nil && config.RewriteLLM == nil { + return nil, fmt.Errorf("at least one of RewriteHandler and RewriteLLM must not be empty") + } + + // construct rewrite chain + rewriteChain := compose.NewChain[string, []string]() + if config.RewriteHandler != nil { + rewriteChain.AppendLambda(compose.InvokableLambda(config.RewriteHandler), compose.WithNodeName("CustomQueryRewriter")) + } else { + tpl := config.RewriteTemplate + variable := config.QueryVar + parser := config.LLMOutputParser + if tpl == nil { + tpl = prompt.FromMessages(schema.Jinja2, schema.UserMessage(defaultRewritePrompt)) + variable = defaultQueryVariable + } + if parser == nil { + parser = func(ctx context.Context, message *schema.Message) ([]string, error) { + return strings.Split(message.Content, "\n"), nil + } + } + + rewriteChain. + AppendLambda(compose.InvokableLambda(func(ctx context.Context, input string) (output map[string]any, err error) { + return map[string]any{variable: input}, nil + }), compose.WithNodeName("Converter")). + AppendChatTemplate(tpl). + AppendChatModel(config.RewriteLLM). + AppendLambda(compose.InvokableLambda(parser), compose.WithNodeName("OutputParser")) + } + rewriteRunner, err := rewriteChain.Compile(ctx, compose.WithGraphName("QueryRewrite")) + if err != nil { + return nil, err + } + + maxQueriesNum := config.MaxQueriesNum + if maxQueriesNum == 0 { + maxQueriesNum = defaultMaxQueriesNum + } + + fusionFunc := config.FusionFunc + if fusionFunc == nil { + fusionFunc = deduplicateFusion + } + + return &multiQueryRetriever{ + queryRunner: rewriteRunner, + maxQueriesNum: maxQueriesNum, + origRetriever: config.OrigRetriever, + fusionFunc: fusionFunc, + }, nil +} + +// Config is the config for multi-query retriever. +type Config struct { + // Rewrite + // 1. set the following fields to use llm to generate multi queries + // a. chat model, required + RewriteLLM model.ChatModel + // b. prompt llm to generate multi queries, we provide default template so you can leave this field blank + RewriteTemplate prompt.ChatTemplate + // c. origin query variable of your custom template, it can be empty if you use default template + QueryVar string + // d. parser llm output to queries, split content using "\n" by default + LLMOutputParser func(context.Context, *schema.Message) ([]string, error) + // 2. set RewriteHandler to provide custom query generation logic, possibly without a ChatModel. If this field is set, it takes precedence over other configurations above + RewriteHandler func(ctx context.Context, query string) ([]string, error) + // limit max queries num that Rewrite generates, and excess queries will be truncated, 5 by default + MaxQueriesNum int + + // Origin Retriever + OrigRetriever retriever.Retriever + + // fusion docs recalled from multi retrievers, remove dup based on document id by default + FusionFunc func(ctx context.Context, docs [][]*schema.Document) ([]*schema.Document, error) +} + +type multiQueryRetriever struct { + queryRunner compose.Runnable[string, []string] + maxQueriesNum int + origRetriever retriever.Retriever + fusionFunc func(ctx context.Context, docs [][]*schema.Document) ([]*schema.Document, error) +} + +// Retrieve retrieves documents from the multi-query retriever. +func (m *multiQueryRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + // generate queries + queries, err := m.queryRunner.Invoke(ctx, query) + if err != nil { + return nil, err + } + if len(queries) > m.maxQueriesNum { + queries = queries[:m.maxQueriesNum] + } + + // retrieve + tasks := make([]*utils.RetrieveTask, len(queries)) + for i := range queries { + tasks[i] = &utils.RetrieveTask{Retriever: m.origRetriever, Query: queries[i]} + } + utils.ConcurrentRetrieveWithCallback(ctx, tasks) + result := make([][]*schema.Document, len(queries)) + for i, task := range tasks { + if task.Err != nil { + return nil, task.Err + } + result[i] = task.Result + } + + // fusion + ctx = ctxWithFusionRunInfo(ctx) + ctx = callbacks.OnStart(ctx, result) + fusionDocs, err := m.fusionFunc(ctx, result) + if err != nil { + callbacks.OnError(ctx, err) + return nil, err + } + callbacks.OnEnd(ctx, fusionDocs) + return fusionDocs, nil +} + +// GetType returns the type of the retriever (MultiQuery). +func (m *multiQueryRetriever) GetType() string { + return "MultiQuery" +} + +func ctxWithFusionRunInfo(ctx context.Context) context.Context { + runInfo := &callbacks.RunInfo{ + Component: compose.ComponentOfLambda, + Type: "FusionFunc", + } + + runInfo.Name = runInfo.Type + string(runInfo.Component) + + return callbacks.ReuseHandlers(ctx, runInfo) +} diff --git a/flow/retriever/multiquery/multi_query_test.go b/flow/retriever/multiquery/multi_query_test.go new file mode 100644 index 0000000..84baf06 --- /dev/null +++ b/flow/retriever/multiquery/multi_query_test.go @@ -0,0 +1,118 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package multiquery + +import ( + "context" + "strings" + "testing" + + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +type mockRetriever struct { +} + +func (m *mockRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + var ret []*schema.Document + if strings.Contains(query, "1") { + ret = append(ret, &schema.Document{ID: "1"}) + } + if strings.Contains(query, "2") { + ret = append(ret, &schema.Document{ID: "2"}) + } + if strings.Contains(query, "3") { + ret = append(ret, &schema.Document{ID: "3"}) + } + if strings.Contains(query, "4") { + ret = append(ret, &schema.Document{ID: "4"}) + } + if strings.Contains(query, "5") { + ret = append(ret, &schema.Document{ID: "5"}) + } + return ret, nil +} + +type mockModel struct { +} + +func (m *mockModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + return &schema.Message{ + Content: "12\n23\n34\n14\n23\n45", + }, nil +} + +func (m *mockModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + panic("implement me") +} + +func (m *mockModel) BindTools(tools []*schema.ToolInfo) error { + panic("implement me") +} + +func TestMultiQueryRetriever(t *testing.T) { + ctx := context.Background() + + // use default llm + mqr, err := NewRetriever(ctx, &Config{ + RewriteLLM: &mockModel{}, + OrigRetriever: &mockRetriever{}, + }) + if err != nil { + t.Fatal(err) + } + c := compose.NewChain[string, []*schema.Document]() + cr, err := c.AppendRetriever(mqr).Compile(ctx) + if err != nil { + t.Fatal(err) + } + + result, err := cr.Invoke(ctx, "query") + if err != nil { + t.Fatal(err) + } + if len(result) != 4 { + t.Fatal("default llm retrieve result is unexpected") + } + + // use custom + mqr, err = NewRetriever(ctx, &Config{ + RewriteHandler: func(ctx context.Context, query string) ([]string, error) { + return []string{"1", "3", "5"}, nil + }, + OrigRetriever: &mockRetriever{}, + }) + if err != nil { + t.Fatal(err) + } + c = compose.NewChain[string, []*schema.Document]() + cr, err = c.AppendRetriever(mqr).Compile(ctx) + if err != nil { + t.Fatal(err) + } + + result, err = cr.Invoke(ctx, "query") + if err != nil { + t.Fatal(err) + } + if len(result) != 3 { + t.Fatal("default llm retrieve result is unexpected") + } +} diff --git a/flow/retriever/parent/doc.go b/flow/retriever/parent/doc.go new file mode 100644 index 0000000..1b5937c --- /dev/null +++ b/flow/retriever/parent/doc.go @@ -0,0 +1,19 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package parent provides a retriever that maps sub-document results +// back to their original parent documents. +package parent diff --git a/flow/retriever/parent/parent.go b/flow/retriever/parent/parent.go new file mode 100644 index 0000000..979afeb --- /dev/null +++ b/flow/retriever/parent/parent.go @@ -0,0 +1,113 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parent + +import ( + "context" + "fmt" + + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +// Config configures the parent retriever. +type Config struct { + // Retriever specifies the original retriever used to retrieve documents. + // For example: a vector database retriever like Milvus, or a full-text search retriever like Elasticsearch. + Retriever retriever.Retriever + // ParentIDKey specifies the key used in the sub-document metadata to store the parent document ID. + // Documents without this key will be removed from the recall results. + // For example: if ParentIDKey is "parent_id", it will look for metadata like: + // {"parent_id": "original_doc_123"} + ParentIDKey string + // OrigDocGetter specifies the method for getting original documents by ids from the sub-document metadata. + // Parameters: + // - ctx: context for the operation + // - ids: slice of parent document IDs to retrieve + // Returns: + // - []*schema.Document: slice of retrieved parent documents + // - error: any error encountered during retrieval + // + // For example: if sub-documents with parent IDs ["doc_1", "doc_2"] are retrieved, + // OrigDocGetter will be called to fetch the original documents with these IDs. + OrigDocGetter func(ctx context.Context, ids []string) ([]*schema.Document, error) +} + +// NewRetriever creates a new parent retriever that handles retrieving original documents +// based on sub-document search results. +// +// Parameters: +// - ctx: context for the operation +// - config: configuration for the parent retriever +// +// Example usage: +// +// retriever, err := NewRetriever(ctx, &Config{ +// Retriever: milvusRetriever, +// ParentIDKey: "source_doc_id", +// OrigDocGetter: func(ctx context.Context, ids []string) ([]*schema.Document, error) { +// return documentStore.GetByIDs(ctx, ids) +// }, +// }) +// +// Returns: +// - retriever.Retriever: the created parent retriever +// - error: any error encountered during creation +func NewRetriever(ctx context.Context, config *Config) (retriever.Retriever, error) { + if config.Retriever == nil { + return nil, fmt.Errorf("retriever is required") + } + if config.OrigDocGetter == nil { + return nil, fmt.Errorf("orig doc getter is required") + } + return &parentRetriever{ + retriever: config.Retriever, + parentIDKey: config.ParentIDKey, + origDocGetter: config.OrigDocGetter, + }, nil +} + +type parentRetriever struct { + retriever retriever.Retriever + parentIDKey string + origDocGetter func(ctx context.Context, ids []string) ([]*schema.Document, error) +} + +func (p *parentRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + subDocs, err := p.retriever.Retrieve(ctx, query, opts...) + if err != nil { + return nil, err + } + ids := make([]string, 0, len(subDocs)) + for _, subDoc := range subDocs { + if k, ok := subDoc.MetaData[p.parentIDKey]; ok { + if s, okk := k.(string); okk && !inList(s, ids) { + ids = append(ids, s) + } + } + } + return p.origDocGetter(ctx, ids) +} + +func inList(elem string, list []string) bool { + for _, v := range list { + if v == elem { + return true + } + } + return false +} diff --git a/flow/retriever/parent/parent_test.go b/flow/retriever/parent/parent_test.go new file mode 100644 index 0000000..6487e25 --- /dev/null +++ b/flow/retriever/parent/parent_test.go @@ -0,0 +1,88 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package parent + +import ( + "context" + "reflect" + "testing" + + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +type testRetriever struct{} + +func (t *testRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + ret := make([]*schema.Document, 0) + for i := range query { + ret = append(ret, &schema.Document{ + ID: "", + Content: "", + MetaData: map[string]interface{}{ + "parent": query[i : i+1], + }, + }) + } + return ret, nil +} + +func TestParentRetriever(t *testing.T) { + tests := []struct { + name string + config *Config + input string + want []*schema.Document + }{ + { + name: "success", + config: &Config{ + Retriever: &testRetriever{}, + ParentIDKey: "parent", + OrigDocGetter: func(ctx context.Context, ids []string) ([]*schema.Document, error) { + var ret []*schema.Document + for i := range ids { + ret = append(ret, &schema.Document{ID: ids[i]}) + } + return ret, nil + }, + }, + input: "123233", + want: []*schema.Document{ + {ID: "1"}, + {ID: "2"}, + {ID: "3"}, + }, + }, + } + ctx := context.Background() + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r, err := NewRetriever(ctx, tt.config) + if err != nil { + t.Fatal(err) + } + ret, err := r.Retrieve(ctx, tt.input) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(ret, tt.want) { + t.Errorf("got %v, want %v", ret, tt.want) + } + }) + } +} diff --git a/flow/retriever/router/router.go b/flow/retriever/router/router.go new file mode 100644 index 0000000..f036338 --- /dev/null +++ b/flow/retriever/router/router.go @@ -0,0 +1,195 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package router provides retrieval routing helpers that merge results +// from multiple retrievers and apply ranking strategies. +package router + +import ( + "context" + "fmt" + "sort" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/flow/retriever/utils" + "github.com/cloudwego/eino/schema" +) + +var rrf = func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error) { + if len(result) < 1 { + return nil, fmt.Errorf("no documents") + } + if len(result) == 1 { + for _, docs := range result { + return docs, nil + } + } + + docRankMap := make(map[string]float64) + docMap := make(map[string]*schema.Document) + for _, v := range result { + for i := range v { + docMap[v[i].ID] = v[i] + if _, ok := docRankMap[v[i].ID]; !ok { + docRankMap[v[i].ID] = 1.0 / float64(i+60) + } else { + docRankMap[v[i].ID] += 1.0 / float64(i+60) + } + } + } + docList := make([]*schema.Document, 0, len(docMap)) + for id := range docMap { + docList = append(docList, docMap[id]) + } + + sort.Slice(docList, func(i, j int) bool { + return docRankMap[docList[i].ID] > docRankMap[docList[j].ID] + }) + + return docList, nil +} + +// NewRetriever creates a router retriever. +// router retriever is useful when you want to retrieve documents from multiple retrievers with different queries. +// eg. +// +// routerRetriever := router.NewRetriever(ctx, &router.Config{}) +// docs, err := routerRetriever.Retrieve(ctx, "how to build agent with eino") +// if err != nil { +// ... +// } +// println(docs) +func NewRetriever(ctx context.Context, config *Config) (retriever.Retriever, error) { + if len(config.Retrievers) == 0 { + return nil, fmt.Errorf("retrievers is empty") + } + + router := config.Router + if router == nil { + var retrieverSet []string + for k := range config.Retrievers { + retrieverSet = append(retrieverSet, k) + } + router = func(ctx context.Context, query string) ([]string, error) { + return retrieverSet, nil + } + } + + fusion := config.FusionFunc + if fusion == nil { + fusion = rrf + } + + return &routerRetriever{ + retrievers: config.Retrievers, + router: config.Router, + fusionFunc: fusion, + }, nil +} + +// Config is the config for router retriever. +type Config struct { + // Retrievers is the retrievers to be used. + Retrievers map[string]retriever.Retriever + // Router is the function to route the query to the retrievers. + Router func(ctx context.Context, query string) ([]string, error) + // FusionFunc is the function to fuse the documents from the retrievers. + FusionFunc func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error) +} + +type routerRetriever struct { + retrievers map[string]retriever.Retriever + router func(ctx context.Context, query string) ([]string, error) + fusionFunc func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error) +} + +// Retrieve retrieves documents from the router retriever. +func (e *routerRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + routeCtx := ctxWithRouterRunInfo(ctx) + routeCtx = callbacks.OnStart(routeCtx, query) + retrieverNames, err := e.router(routeCtx, query) + if err != nil { + callbacks.OnError(routeCtx, err) + return nil, err + } + if len(retrieverNames) == 0 { + err = fmt.Errorf("no retriever has been selected") + callbacks.OnError(routeCtx, err) + return nil, err + } + callbacks.OnEnd(routeCtx, retrieverNames) + + // retrieve + tasks := make([]*utils.RetrieveTask, len(retrieverNames)) + for i := range retrieverNames { + r, ok := e.retrievers[retrieverNames[i]] + if !ok { + return nil, fmt.Errorf("router output[%s] has not registered", retrieverNames[i]) + } + tasks[i] = &utils.RetrieveTask{ + Name: retrieverNames[i], + Retriever: r, + Query: query, + RetrieveOptions: opts, + } + } + utils.ConcurrentRetrieveWithCallback(ctx, tasks) + result := make(map[string][]*schema.Document) + for i := range tasks { + if tasks[i].Err != nil { + return nil, tasks[i].Err + } + result[tasks[i].Name] = tasks[i].Result + } + + // fusion + fusionCtx := ctxWithFusionRunInfo(ctx) + fusionCtx = callbacks.OnStart(fusionCtx, result) + fusionDocs, err := e.fusionFunc(fusionCtx, result) + if err != nil { + callbacks.OnError(fusionCtx, err) + return nil, err + } + callbacks.OnEnd(fusionCtx, fusionDocs) + return fusionDocs, nil +} + +// GetType returns the type of the retriever (Router). +func (e *routerRetriever) GetType() string { return "Router" } + +func ctxWithRouterRunInfo(ctx context.Context) context.Context { + runInfo := &callbacks.RunInfo{ + Component: compose.ComponentOfLambda, + Type: "Router", + } + + runInfo.Name = runInfo.Type + string(runInfo.Component) + + return callbacks.ReuseHandlers(ctx, runInfo) +} + +func ctxWithFusionRunInfo(ctx context.Context) context.Context { + runInfo := &callbacks.RunInfo{ + Component: compose.ComponentOfLambda, + Type: "FusionFunc", + } + + runInfo.Name = runInfo.Type + string(runInfo.Component) + + return callbacks.ReuseHandlers(ctx, runInfo) +} diff --git a/flow/retriever/router/router_test.go b/flow/retriever/router/router_test.go new file mode 100644 index 0000000..edf1af1 --- /dev/null +++ b/flow/retriever/router/router_test.go @@ -0,0 +1,134 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package router + +import ( + "context" + "reflect" + "strings" + "testing" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +type mockRetriever struct { +} + +func (m *mockRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + var ret []*schema.Document + if strings.Contains(query, "1") { + ret = append(ret, &schema.Document{ID: "1"}) + } + if strings.Contains(query, "2") { + ret = append(ret, &schema.Document{ID: "2"}) + } + if strings.Contains(query, "3") { + ret = append(ret, &schema.Document{ID: "3"}) + } + if strings.Contains(query, "4") { + ret = append(ret, &schema.Document{ID: "4"}) + } + if strings.Contains(query, "5") { + ret = append(ret, &schema.Document{ID: "5"}) + } + return ret, nil +} + +func (m *mockRetriever) GetType() string { + return "Mock" +} + +func TestRouterRetriever(t *testing.T) { + ctx := context.Background() + r, err := NewRetriever(ctx, &Config{ + Retrievers: map[string]retriever.Retriever{ + "1": &mockRetriever{}, + "2": &mockRetriever{}, + "3": &mockRetriever{}, + }, + Router: func(ctx context.Context, query string) ([]string, error) { + return []string{"2", "3"}, nil + }, + FusionFunc: func(ctx context.Context, result map[string][]*schema.Document) ([]*schema.Document, error) { + var ret []*schema.Document + for _, v := range result { + ret = append(ret, v...) + } + return ret, nil + }, + }) + if err != nil { + t.Fatal(err) + } + + handler := callbacks.NewHandlerBuilder(). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + switch info.Name { + case "FusionFuncLambda": + if _, ok := output.([]*schema.Document); !ok { + t.Fatal("FusionFuncLambda output is not a []*schema.Document") + } + case "RouterLambda": + if _, ok := output.([]string); !ok { + t.Fatal("RouterLambda output is not a []string") + } + case "MockRetriever": + if _, ok := output.([]*schema.Document); !ok { + t.Fatal("MockRetriever output is not a []string") + } + default: + t.Fatalf("unknown name: %s", info.Name) + } + return ctx + }). + OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + t.Fatal(err) + return ctx + }).Build() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{}, handler) + result, err := r.Retrieve(ctx, "3") + if err != nil { + t.Fatal(err) + } + if len(result) != 2 { + t.Fatal("expected 2 results") + } +} + +func TestRRF(t *testing.T) { + doc1 := &schema.Document{ID: "1"} + doc2 := &schema.Document{ID: "2"} + doc3 := &schema.Document{ID: "3"} + doc4 := &schema.Document{ID: "4"} + doc5 := &schema.Document{ID: "5"} + + input := map[string][]*schema.Document{ + "1": {doc1, doc2, doc3, doc4, doc5}, + "2": {doc2, doc3, doc4, doc5, doc1}, + "3": {doc3, doc4, doc5, doc1, doc2}, + } + + result, err := rrf(context.Background(), input) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(result, []*schema.Document{doc3, doc2, doc4, doc1, doc5}) { + t.Fatal("rrf fail") + } +} diff --git a/flow/retriever/utils/utils.go b/flow/retriever/utils/utils.go new file mode 100644 index 0000000..4fab29c --- /dev/null +++ b/flow/retriever/utils/utils.go @@ -0,0 +1,86 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package utils provides helper utilities for retriever flows, including +// concurrent retrieval with callback instrumentation. +package utils + +import ( + "context" + "fmt" + "sync" + + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/schema" +) + +// RetrieveTask is a task for retrieving documents. +// RetrieveTask represents a single retrieval job with its result or error. +type RetrieveTask struct { + Name string + Retriever retriever.Retriever + Query string + RetrieveOptions []retriever.Option + Result []*schema.Document + Err error +} + +// ConcurrentRetrieveWithCallback concurrently retrieves documents with callback. +func ConcurrentRetrieveWithCallback(ctx context.Context, tasks []*RetrieveTask) { + wg := sync.WaitGroup{} + for i := range tasks { + wg.Add(1) + go func(ctx context.Context, t *RetrieveTask) { + ctx = ctxWithRetrieverRunInfo(ctx, t.Retriever) + + defer func() { + if e := recover(); e != nil { + t.Err = fmt.Errorf("retrieve panic, query: %s, error: %v", t.Query, e) + ctx = callbacks.OnError(ctx, t.Err) + } + wg.Done() + }() + + ctx = callbacks.OnStart(ctx, t.Query) + docs, err := t.Retriever.Retrieve(ctx, t.Query, t.RetrieveOptions...) + if err != nil { + callbacks.OnError(ctx, err) + t.Err = err + return + } + + callbacks.OnEnd(ctx, docs) + t.Result = docs + }(ctx, tasks[i]) + } + wg.Wait() +} + +func ctxWithRetrieverRunInfo(ctx context.Context, r retriever.Retriever) context.Context { + runInfo := &callbacks.RunInfo{ + Component: components.ComponentOfRetriever, + } + + if typ, okk := components.GetType(r); okk { + runInfo.Type = typ + } + + runInfo.Name = runInfo.Type + string(runInfo.Component) + + return callbacks.ReuseHandlers(ctx, runInfo) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..0b87a6c --- /dev/null +++ b/go.mod @@ -0,0 +1,47 @@ +module github.com/cloudwego/eino + +go 1.18 + +require ( + github.com/bmatcuk/doublestar/v4 v4.10.0 + github.com/bytedance/sonic v1.15.0 + github.com/eino-contrib/jsonschema v1.0.3 + github.com/google/uuid v1.6.0 + github.com/nikolalohinski/gonja v1.5.3 + github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f + github.com/smartystreets/goconvey v1.8.1 + github.com/stretchr/testify v1.10.0 + github.com/wk8/go-ordered-map/v2 v2.1.8 + go.uber.org/mock v0.4.0 + gopkg.in/yaml.v3 v3.0.1 +) + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic/loader v0.5.0 // indirect + github.com/cloudwego/base64x v0.1.6 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/goph/emperror v0.17.2 // indirect + github.com/gopherjs/gopherjs v1.17.2 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/jtolds/gls v4.20.0+incompatible // indirect + github.com/klauspost/cpuid/v2 v2.2.9 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/pelletier/go-toml/v2 v2.0.9 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/sirupsen/logrus v1.9.3 // indirect + github.com/smarty/assertions v1.15.0 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/yargevad/filepathx v1.0.0 // indirect + golang.org/x/arch v0.11.0 // indirect + golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect + golang.org/x/sys v0.29.0 // indirect + golang.org/x/term v0.28.0 // indirect + gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..5813766 --- /dev/null +++ b/go.sum @@ -0,0 +1,133 @@ +github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmatcuk/doublestar/v4 v4.10.0 h1:zU9WiOla1YA122oLM6i4EXvGW62DvKZVxIe6TYWexEs= +github.com/bmatcuk/doublestar/v4 v4.10.0/go.mod h1:xBQ8jztBU6kakFMg+8WGxn0c6z1fTSPVIjEY1Wr7jzc= +github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8= +github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE= +github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= +github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= +github.com/bytedance/sonic v1.15.0 h1:/PXeWFaR5ElNcVE84U0dOHjiMHQOwNIx3K4ymzh/uSE= +github.com/bytedance/sonic v1.15.0/go.mod h1:tFkWrPz0/CUCLEF4ri4UkHekCIcdnkqXw9VduqpJh0k= +github.com/bytedance/sonic/loader v0.5.0 h1:gXH3KVnatgY7loH5/TkeVyXPfESoqSBSBEiDd5VjlgE= +github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= +github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= +github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/eino-contrib/jsonschema v1.0.3 h1:2Kfsm1xlMV0ssY2nuxshS4AwbLFuqmPmzIjLVJ1Fsp0= +github.com/eino-contrib/jsonschema v1.0.3/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI= +github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18= +github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/cpuid/v2 v2.2.9 h1:66ze0taIn2H33fBvCkXuv9BmCwDfafmiIVpKV9kKGuY= +github.com/klauspost/cpuid/v2 v2.2.9/go.mod h1:rqkxqrZ1EhYM9G+hXH7YdowN5R5RGN6NK4QwQ3WMXF8= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-colorable v0.1.2 h1:/bC9yWikZXAL9uJdulbSfyVNIR3n3trXl+v8+1sx8mU= +github.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= +github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/pelletier/go-toml/v2 v2.0.9 h1:uH2qQXheeefCCkuBBSLi7jCiSmj3VRh2+Goq2N7Xxu0= +github.com/pelletier/go-toml/v2 v2.0.9/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ= +github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg= +github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc= +github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= +go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU= +go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc= +golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4= +golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= +golang.org/x/crypto v0.11.0 h1:6Ewdq3tDic1mg5xRO4milcWCfMVQhI4NkqWWvqejpuA= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU= +golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.28.0 h1:/Ts8HFuMR2E6IP/jlo7QVLZHggjKQbhu/7H0LJFr3Gg= +golang.org/x/term v0.28.0/go.mod h1:Sw/lC2IAUZ92udQNf3WodGtn4k/XoLyZoh8v/8uiwek= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/callbacks/inject.go b/internal/callbacks/inject.go new file mode 100644 index 0000000..0ee1579 --- /dev/null +++ b/internal/callbacks/inject.go @@ -0,0 +1,203 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/schema" +) + +func InitCallbacks(ctx context.Context, info *RunInfo, handlers ...Handler) context.Context { + mgr, ok := newManager(info, handlers...) + if ok { + return ctxWithManager(ctx, mgr) + } + + return ctxWithManager(ctx, nil) +} + +func EnsureRunInfo(ctx context.Context, typ string, comp components.Component) context.Context { + cbm, ok := managerFromCtx(ctx) + if !ok { + return InitCallbacks(ctx, &RunInfo{ + Type: typ, + Component: comp, + }) + } + if cbm.runInfo == nil { + return ReuseHandlers(ctx, &RunInfo{ + Type: typ, + Component: comp, + }) + } + return ctx +} + +func ReuseHandlers(ctx context.Context, info *RunInfo) context.Context { + cbm, ok := managerFromCtx(ctx) + if !ok { + return InitCallbacks(ctx, info) + } + return ctxWithManager(ctx, cbm.withRunInfo(info)) +} + +func AppendHandlers(ctx context.Context, info *RunInfo, handlers ...Handler) context.Context { + cbm, ok := managerFromCtx(ctx) + if !ok { + return InitCallbacks(ctx, info, handlers...) + } + nh := make([]Handler, len(cbm.handlers)+len(handlers)) + copy(nh[:len(cbm.handlers)], cbm.handlers) + copy(nh[len(cbm.handlers):], handlers) + return InitCallbacks(ctx, info, nh...) +} + +type Handle[T any] func(context.Context, T, *RunInfo, []Handler) (context.Context, T) + +func On[T any](ctx context.Context, inOut T, handle Handle[T], timing CallbackTiming, start bool) (context.Context, T) { + mgr, ok := managerFromCtx(ctx) + if !ok { + return ctx, inOut + } + nMgr := *mgr + + var info *RunInfo + if start { + info = nMgr.runInfo + nMgr.runInfo = nil + ctx = context.WithValue(ctx, CtxRunInfoKey{}, info) + } else { + if nMgr.runInfo != nil { + info = nMgr.runInfo + } else { + info, _ = ctx.Value(CtxRunInfoKey{}).(*RunInfo) + } + } + + hs := make([]Handler, 0, len(nMgr.handlers)+len(nMgr.globalHandlers)) + for _, handler := range append(nMgr.handlers, nMgr.globalHandlers...) { + timingChecker, ok_ := handler.(TimingChecker) + if !ok_ || timingChecker.Needed(ctx, info, timing) { + hs = append(hs, handler) + } + } + + var out T + ctx, out = handle(ctx, inOut, info, hs) + return ctxWithManager(ctx, &nMgr), out +} + +func OnStartHandle[T any](ctx context.Context, input T, + runInfo *RunInfo, handlers []Handler) (context.Context, T) { + + for i := len(handlers) - 1; i >= 0; i-- { + ctx = handlers[i].OnStart(ctx, runInfo, input) + } + + return ctx, input +} + +func OnEndHandle[T any](ctx context.Context, output T, + runInfo *RunInfo, handlers []Handler) (context.Context, T) { + + for _, handler := range handlers { + ctx = handler.OnEnd(ctx, runInfo, output) + } + + return ctx, output +} + +func BuildOnEndHandleWithCopy[T any](copyFn func(T, int) []T) Handle[T] { + return func(ctx context.Context, output T, runInfo *RunInfo, handlers []Handler) (context.Context, T) { + if len(handlers) == 0 { + return ctx, output + } + + copies := copyFn(output, len(handlers)) + + for i, handler := range handlers { + ctx = handler.OnEnd(ctx, runInfo, copies[i]) + } + + return ctx, output + } +} + +func OnWithStreamHandle[S any]( + ctx context.Context, + inOut S, + handlers []Handler, + cpy func(int) []S, + handle func(context.Context, Handler, S) context.Context) (context.Context, S) { + + if len(handlers) == 0 { + return ctx, inOut + } + + inOuts := cpy(len(handlers) + 1) + + for i, handler := range handlers { + ctx = handle(ctx, handler, inOuts[i]) + } + + return ctx, inOuts[len(inOuts)-1] +} + +func OnStartWithStreamInputHandle[T any](ctx context.Context, input *schema.StreamReader[T], + runInfo *RunInfo, handlers []Handler) (context.Context, *schema.StreamReader[T]) { + + handlers = generic.Reverse(handlers) + + cpy := input.Copy + + handle := func(ctx context.Context, handler Handler, in *schema.StreamReader[T]) context.Context { + in_ := schema.StreamReaderWithConvert(in, func(i T) (CallbackInput, error) { + return i, nil + }) + return handler.OnStartWithStreamInput(ctx, runInfo, in_) + } + + return OnWithStreamHandle(ctx, input, handlers, cpy, handle) +} + +func OnEndWithStreamOutputHandle[T any](ctx context.Context, output *schema.StreamReader[T], + runInfo *RunInfo, handlers []Handler) (context.Context, *schema.StreamReader[T]) { + + cpy := output.Copy + + handle := func(ctx context.Context, handler Handler, out *schema.StreamReader[T]) context.Context { + out_ := schema.StreamReaderWithConvert(out, func(i T) (CallbackOutput, error) { + return i, nil + }) + return handler.OnEndWithStreamOutput(ctx, runInfo, out_) + } + + return OnWithStreamHandle(ctx, output, handlers, cpy, handle) +} + +func OnErrorHandle(ctx context.Context, err error, + runInfo *RunInfo, handlers []Handler) (context.Context, error) { + + for _, handler := range handlers { + ctx = handler.OnError(ctx, runInfo, err) + } + + return ctx, err +} diff --git a/internal/callbacks/interface.go b/internal/callbacks/interface.go new file mode 100644 index 0000000..9e7665c --- /dev/null +++ b/internal/callbacks/interface.go @@ -0,0 +1,54 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/schema" +) + +type RunInfo struct { + // Name is the graph node name for display purposes, not unique. + // Passed from compose.WithNodeName(). + Name string + Type string + Component components.Component +} + +type CallbackInput any + +type CallbackOutput any + +type Handler interface { + OnStart(ctx context.Context, info *RunInfo, input CallbackInput) context.Context + OnEnd(ctx context.Context, info *RunInfo, output CallbackOutput) context.Context + + OnError(ctx context.Context, info *RunInfo, err error) context.Context + + OnStartWithStreamInput(ctx context.Context, info *RunInfo, + input *schema.StreamReader[CallbackInput]) context.Context + OnEndWithStreamOutput(ctx context.Context, info *RunInfo, + output *schema.StreamReader[CallbackOutput]) context.Context +} + +type CallbackTiming uint8 + +type TimingChecker interface { + Needed(ctx context.Context, info *RunInfo, timing CallbackTiming) bool +} diff --git a/internal/callbacks/manager.go b/internal/callbacks/manager.go new file mode 100644 index 0000000..6daf82e --- /dev/null +++ b/internal/callbacks/manager.go @@ -0,0 +1,70 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import "context" + +type CtxManagerKey struct{} +type CtxRunInfoKey struct{} + +type manager struct { + globalHandlers []Handler + handlers []Handler + runInfo *RunInfo +} + +var GlobalHandlers []Handler + +func newManager(runInfo *RunInfo, handlers ...Handler) (*manager, bool) { + if len(handlers)+len(GlobalHandlers) == 0 { + return nil, false + } + + hs := make([]Handler, len(GlobalHandlers)) + copy(hs, GlobalHandlers) + + return &manager{ + globalHandlers: hs, + handlers: handlers, + runInfo: runInfo, + }, true +} + +func ctxWithManager(ctx context.Context, manager *manager) context.Context { + return context.WithValue(ctx, CtxManagerKey{}, manager) +} + +func (m *manager) withRunInfo(runInfo *RunInfo) *manager { + if m == nil { + return nil + } + + n := *m + n.runInfo = runInfo + return &n +} + +func managerFromCtx(ctx context.Context) (*manager, bool) { + v := ctx.Value(CtxManagerKey{}) + m, ok := v.(*manager) + if ok && m != nil { + n := *m + return &n, true + } + + return nil, false +} diff --git a/internal/channel.go b/internal/channel.go new file mode 100644 index 0000000..fa42153 --- /dev/null +++ b/internal/channel.go @@ -0,0 +1,94 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import "sync" + +// UnboundedChan represents a channel with unlimited capacity +type UnboundedChan[T any] struct { + buffer []T // Internal buffer to store data + mutex sync.Mutex // Mutex to protect buffer access + notEmpty *sync.Cond // Condition variable to wait for data + closed bool // Indicates if the channel has been closed +} + +// NewUnboundedChan initializes and returns an UnboundedChan +func NewUnboundedChan[T any]() *UnboundedChan[T] { + ch := &UnboundedChan[T]{} + ch.notEmpty = sync.NewCond(&ch.mutex) + return ch +} + +// Send puts an item into the channel +func (ch *UnboundedChan[T]) Send(value T) { + ch.mutex.Lock() + defer ch.mutex.Unlock() + + if ch.closed { + panic("send on closed channel") + } + + ch.buffer = append(ch.buffer, value) + ch.notEmpty.Signal() // Wake up one goroutine waiting to receive +} + +// TrySend attempts to put an item into the channel. +// Returns false if the channel is closed, true otherwise. +func (ch *UnboundedChan[T]) TrySend(value T) bool { + ch.mutex.Lock() + defer ch.mutex.Unlock() + + if ch.closed { + return false + } + + ch.buffer = append(ch.buffer, value) + ch.notEmpty.Signal() + return true +} + +// Receive gets an item from the channel (blocks if empty). +// Returns (value, true) if an item was received. +// Returns (zero, false) if the channel was closed with no data remaining. +func (ch *UnboundedChan[T]) Receive() (T, bool) { + ch.mutex.Lock() + defer ch.mutex.Unlock() + + for len(ch.buffer) == 0 && !ch.closed { + ch.notEmpty.Wait() + } + + if len(ch.buffer) == 0 { + var zero T + return zero, false + } + + val := ch.buffer[0] + ch.buffer = ch.buffer[1:] + return val, true +} + +// Close marks the channel as closed +func (ch *UnboundedChan[T]) Close() { + ch.mutex.Lock() + defer ch.mutex.Unlock() + + if !ch.closed { + ch.closed = true + ch.notEmpty.Broadcast() + } +} diff --git a/internal/channel_test.go b/internal/channel_test.go new file mode 100644 index 0000000..736a274 --- /dev/null +++ b/internal/channel_test.go @@ -0,0 +1,221 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "sync" + "testing" + "time" +) + +func TestUnboundedChan_Send(t *testing.T) { + ch := NewUnboundedChan[string]() + + // Test sending a value + ch.Send("test") + if len(ch.buffer) != 1 { + t.Errorf("buffer length should be 1, got %d", len(ch.buffer)) + } + if ch.buffer[0] != "test" { + t.Errorf("expected 'test', got '%s'", ch.buffer[0]) + } + + // Test sending multiple values + ch.Send("test2") + ch.Send("test3") + if len(ch.buffer) != 3 { + t.Errorf("buffer length should be 3, got %d", len(ch.buffer)) + } +} + +func TestUnboundedChan_SendPanic(t *testing.T) { + ch := NewUnboundedChan[int]() + ch.Close() + + // Test sending to closed channel should panic + defer func() { + if r := recover(); r == nil { + t.Error("sending to closed channel should panic") + } + }() + + ch.Send(1) +} + +func TestUnboundedChan_Receive(t *testing.T) { + ch := NewUnboundedChan[int]() + + // Send values + ch.Send(1) + ch.Send(2) + + // Test receiving values + val, ok := ch.Receive() + if !ok { + t.Error("receive should succeed") + } + if val != 1 { + t.Errorf("expected 1, got %d", val) + } + + val, ok = ch.Receive() + if !ok { + t.Error("receive should succeed") + } + if val != 2 { + t.Errorf("expected 2, got %d", val) + } +} + +func TestUnboundedChan_ReceiveFromClosed(t *testing.T) { + ch := NewUnboundedChan[int]() + ch.Close() + + // Test receiving from closed, empty channel + val, ok := ch.Receive() + if ok { + t.Error("receive from closed, empty channel should return ok=false") + } + if val != 0 { + t.Errorf("expected zero value, got %d", val) + } + + // Test receiving from closed channel with values + ch = NewUnboundedChan[int]() + ch.Send(42) + ch.Close() + + val, ok = ch.Receive() + if !ok { + t.Error("receive should succeed") + } + if val != 42 { + t.Errorf("expected 42, got %d", val) + } + + // After consuming all values + val, ok = ch.Receive() + if ok { + t.Error("receive from closed, empty channel should return ok=false") + } +} + +func TestUnboundedChan_Close(t *testing.T) { + ch := NewUnboundedChan[int]() + + // Test closing + ch.Close() + if !ch.closed { + t.Error("channel should be marked as closed") + } + + // Test double closing (should not panic) + ch.Close() +} + +func TestUnboundedChan_Concurrency(t *testing.T) { + ch := NewUnboundedChan[int]() + const numSenders = 5 + const numReceivers = 3 + const messagesPerSender = 100 + + var rwg, swg sync.WaitGroup + rwg.Add(numReceivers) + swg.Add(numSenders) + + // Start senders + for i := 0; i < numSenders; i++ { + go func(id int) { + defer swg.Done() + for j := 0; j < messagesPerSender; j++ { + ch.Send(id*messagesPerSender + j) + time.Sleep(time.Microsecond) // Small delay to increase concurrency chance + } + }(i) + } + + // Start receivers + received := make([]int, 0, numSenders*messagesPerSender) + var mu sync.Mutex + + for i := 0; i < numReceivers; i++ { + go func() { + defer rwg.Done() + for { + val, ok := ch.Receive() + if !ok { + return + } + mu.Lock() + received = append(received, val) + mu.Unlock() + } + }() + } + + // Wait for senders to finish + swg.Wait() + ch.Close() + + // Wait for all goroutines to finish + rwg.Wait() + + // Verify we received all messages + if len(received) != numSenders*messagesPerSender { + t.Errorf("expected %d messages, got %d", numSenders*messagesPerSender, len(received)) + } + + // Create a map to check for duplicates and missing values + receivedMap := make(map[int]bool) + for _, val := range received { + receivedMap[val] = true + } + + if len(receivedMap) != numSenders*messagesPerSender { + t.Error("duplicate or missing messages detected") + } +} + +func TestUnboundedChan_BlockingReceive(t *testing.T) { + ch := NewUnboundedChan[int]() + + // Test that Receive blocks when channel is empty + receiveDone := make(chan bool) + go func() { + ch.Receive() + receiveDone <- true + }() + + // Check that receive is blocked + select { + case <-receiveDone: + t.Error("Receive should block on empty channel") + case <-time.After(50 * time.Millisecond): + // This is expected + } + + // Send a value to unblock + ch.Send(1) + + // Now receive should complete + select { + case <-receiveDone: + // This is expected + case <-time.After(50 * time.Millisecond): + t.Error("Receive should have unblocked") + } +} diff --git a/internal/concat.go b/internal/concat.go new file mode 100644 index 0000000..fd9b8ab --- /dev/null +++ b/internal/concat.go @@ -0,0 +1,224 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "fmt" + "reflect" + "strings" + "time" + + "github.com/cloudwego/eino/internal/generic" +) + +var ( + concatFuncs = map[reflect.Type]any{ + generic.TypeOf[string](): concatStrings, + generic.TypeOf[int8](): useLast[int8], + generic.TypeOf[int16](): useLast[int16], + generic.TypeOf[int32](): useLast[int32], + generic.TypeOf[int64](): useLast[int64], + generic.TypeOf[int](): useLast[int], + generic.TypeOf[uint8](): useLast[uint8], + generic.TypeOf[uint16](): useLast[uint16], + generic.TypeOf[uint32](): useLast[uint32], + generic.TypeOf[uint64](): useLast[uint64], + generic.TypeOf[uint](): useLast[uint], + generic.TypeOf[bool](): useLast[bool], + generic.TypeOf[float32](): useLast[float32], + generic.TypeOf[float64](): useLast[float64], + generic.TypeOf[time.Time](): useLast[time.Time], + generic.TypeOf[time.Duration](): useLast[time.Duration], + } +) + +func useLast[T any](s []T) (T, error) { + return s[len(s)-1], nil +} + +func concatStrings(ss []string) (string, error) { + var n int + for _, s := range ss { + n += len(s) + } + + var b strings.Builder + b.Grow(n) + for _, s := range ss { + _, err := b.WriteString(s) + if err != nil { + return "", err + } + } + + return b.String(), nil +} + +func RegisterStreamChunkConcatFunc[T any](fn func([]T) (T, error)) { + concatFuncs[generic.TypeOf[T]()] = fn +} + +func GetConcatFunc(typ reflect.Type) func(reflect.Value) (reflect.Value, error) { + if fn, ok := concatFuncs[typ]; ok { + return func(a reflect.Value) (reflect.Value, error) { + rvs := reflect.ValueOf(fn).Call([]reflect.Value{a}) + var err error + if !rvs[1].IsNil() { + err = rvs[1].Interface().(error) + } + return rvs[0], err + } + } + + return nil +} + +// ConcatItems the caller should ensure len(items) > 1 +func ConcatItems[T any](items []T) (T, error) { + typ := generic.TypeOf[T]() + v := reflect.ValueOf(items) + + var cv reflect.Value + var err error + + // handle map kind + if typ.Kind() == reflect.Map { + cv, err = concatMaps(v) + } else { + cv, err = ConcatSliceValue(v) + } + + if err != nil { + var t T + return t, err + } + + return cv.Interface().(T), nil +} + +func concatMaps(ms reflect.Value) (reflect.Value, error) { + typ := ms.Type().Elem() + + rms := reflect.MakeMap(reflect.MapOf(typ.Key(), generic.TypeOf[[]any]())) + ret := reflect.MakeMap(typ) + + n := ms.Len() + for i := 0; i < n; i++ { + m := ms.Index(i) + + for _, key := range m.MapKeys() { + vals := rms.MapIndex(key) + if !vals.IsValid() { + var s []any + vals = reflect.ValueOf(s) + } + + val := m.MapIndex(key) + vals = reflect.Append(vals, val) + rms.SetMapIndex(key, vals) + } + } + + for _, key := range rms.MapKeys() { + vals := rms.MapIndex(key) + + anyVals := vals.Interface().([]any) + if len(anyVals) == 1 { + ele := anyVals[0] + if ele == nil { // we cannot SetMapIndex with nil because it will delete the key + ret.SetMapIndex(key, reflect.Zero(typ.Elem())) + continue + } + + ret.SetMapIndex(key, reflect.ValueOf(ele)) + continue + } + + v, err := toSliceValue(anyVals) + if err != nil { + return reflect.Value{}, err + } + + var cv reflect.Value + + if v.Type().Elem().Kind() == reflect.Map { + cv, err = concatMaps(v) + } else { + cv, err = ConcatSliceValue(v) + } + + if err != nil { + return reflect.Value{}, err + } + + ret.SetMapIndex(key, cv) + } + + return ret, nil +} + +func ConcatSliceValue(val reflect.Value) (reflect.Value, error) { + elmType := val.Type().Elem() + + if val.Len() == 1 { + return val.Index(0), nil + } + + f := GetConcatFunc(elmType) + if f != nil { + return f(val) + } + + // if all elements in the slice are empty, return an empty value + // if there is exactly one non-empty element in the slice, return that non-empty element + // otherwise, throw an error. + var filtered reflect.Value + for i := 0; i < val.Len(); i++ { + oneVal := val.Index(i) + if !oneVal.IsZero() { + if filtered.IsValid() { + return reflect.Value{}, fmt.Errorf("cannot concat multiple non-zero value of type %s", elmType) + } + + filtered = oneVal + } + } + if !filtered.IsValid() { + filtered = reflect.New(elmType).Elem() + } + + return filtered, nil +} + +func toSliceValue(vs []any) (reflect.Value, error) { + typ := reflect.TypeOf(vs[0]) + + ret := reflect.MakeSlice(reflect.SliceOf(typ), len(vs), len(vs)) + ret.Index(0).Set(reflect.ValueOf(vs[0])) + + for i := 1; i < len(vs); i++ { + v := vs[i] + vt := reflect.TypeOf(v) + if typ != vt { + return reflect.Value{}, fmt.Errorf("unexpected slice element type. Got %v, expected %v", typ, vt) + } + + ret.Index(i).Set(reflect.ValueOf(v)) + } + + return ret, nil +} diff --git a/internal/concat_test.go b/internal/concat_test.go new file mode 100644 index 0000000..12b4837 --- /dev/null +++ b/internal/concat_test.go @@ -0,0 +1,52 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestConcat(t *testing.T) { + t.Run("concat map chunks with nil value", func(t *testing.T) { + c1 := map[string]any{ + "a": map[string]any{ + "b": map[string]any{ + "c1": nil, + }, + }, + } + c2 := map[string]any{ + "a": map[string]any{ + "b": map[string]any{ + "c2": "c2", + }, + }, + } + m, err := ConcatItems([]map[string]any{c1, c2}) + assert.Nil(t, err) + assert.Equal(t, map[string]any{ + "a": map[string]any{ + "b": map[string]any{ + "c1": nil, + "c2": "c2", + }, + }, + }, m) + }) +} diff --git a/internal/core/address.go b/internal/core/address.go new file mode 100644 index 0000000..bb2400a --- /dev/null +++ b/internal/core/address.go @@ -0,0 +1,338 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package core + +import ( + "context" + "fmt" + "strings" + "sync" + + "github.com/cloudwego/eino/internal/generic" +) + +// AddressSegmentType defines the type of a segment in an execution address. +type AddressSegmentType string + +// Address represents a full, hierarchical address to a point in the execution structure. +type Address []AddressSegment + +// String converts an Address into its unique string representation. +func (p Address) String() string { + if p == nil { + return "" + } + var sb strings.Builder + for i, s := range p { + sb.WriteString(string(s.Type)) + sb.WriteString(":") + sb.WriteString(s.ID) + if s.SubID != "" { + sb.WriteString(":") + sb.WriteString(s.SubID) + } + if i != len(p)-1 { + sb.WriteString(";") + } + } + return sb.String() +} + +func (p Address) Equals(other Address) bool { + if len(p) != len(other) { + return false + } + for i := range p { + if p[i].Type != other[i].Type || p[i].ID != other[i].ID || p[i].SubID != other[i].SubID { + return false + } + } + return true +} + +// AddressSegment represents a single segment in the hierarchical address of an execution point. +// A sequence of AddressSegments uniquely identifies a location within a potentially nested structure. +type AddressSegment struct { + // ID is the unique identifier for this segment, e.g., the node's key or the tool's name. + ID string + // Type indicates whether this address segment is a graph node, a tool call, an agent, etc. + Type AddressSegmentType + // In some cases, ID alone are not unique enough, we need this SubID to guarantee uniqueness. + // e.g. parallel tool calls with the same name but different tool call IDs. + SubID string +} + +type addrCtxKey struct{} + +type addrCtx struct { + addr Address + interruptState *InterruptState + isResumeTarget bool + resumeData any +} + +type globalResumeInfoKey struct{} + +type globalResumeInfo struct { + mu sync.RWMutex + id2ResumeData map[string]any + id2ResumeDataUsed map[string]bool + id2State map[string]InterruptState + id2StateUsed map[string]bool + id2Addr map[string]Address +} + +// GetCurrentAddress returns the hierarchical address of the currently executing component. +// The address is a sequence of segments, each identifying a structural part of the execution +// like an agent, a graph node, or a tool call. This can be useful for logging or debugging. +func GetCurrentAddress(ctx context.Context) Address { + if p, ok := ctx.Value(addrCtxKey{}).(*addrCtx); ok { + return p.addr + } + + return nil +} + +// AppendAddressSegment creates a new execution context for a sub-component (e.g., a graph node or a tool call). +// +// It extends the current context's address with a new segment and populates the new context with the +// appropriate interrupt state and resume data for that specific sub-address. +// +// - ctx: The parent context, typically the one passed into the component's Invoke/Stream method. +// - segType: The type of the new address segment (e.g., "node", "tool"). +// - segID: The unique ID for the new address segment. +func AppendAddressSegment(ctx context.Context, segType AddressSegmentType, segID string, + subID string) context.Context { + // get current address + currentAddress := GetCurrentAddress(ctx) + if len(currentAddress) == 0 { + currentAddress = []AddressSegment{ + { + Type: segType, + ID: segID, + SubID: subID, + }, + } + } else { + newAddress := make([]AddressSegment, len(currentAddress)+1) + copy(newAddress, currentAddress) + newAddress[len(newAddress)-1] = AddressSegment{ + Type: segType, + ID: segID, + SubID: subID, + } + currentAddress = newAddress + } + + runCtx := &addrCtx{ + addr: currentAddress, + } + + rInfo, hasRInfo := getResumeInfo(ctx) + if !hasRInfo { + return context.WithValue(ctx, addrCtxKey{}, runCtx) + } + + rInfo.mu.Lock() + defer rInfo.mu.Unlock() + + var id string + for id_, addr := range rInfo.id2Addr { + if addr.Equals(currentAddress) { + if used, ok := rInfo.id2StateUsed[id_]; !ok || !used { + runCtx.interruptState = generic.PtrOf(rInfo.id2State[id_]) + rInfo.id2StateUsed[id_] = true + id = id_ + break + } + } + } + + used := rInfo.id2ResumeDataUsed[id] + if !used { + rData, existed := rInfo.id2ResumeData[id] + if existed { + rInfo.id2ResumeDataUsed[id] = true + runCtx.resumeData = rData + runCtx.isResumeTarget = true + } + } + + if !runCtx.isResumeTarget { + for id_, addr := range rInfo.id2Addr { + if len(addr) > len(currentAddress) && addr[:len(currentAddress)].Equals(currentAddress) { + if !rInfo.id2ResumeDataUsed[id_] { + runCtx.isResumeTarget = true + break + } + } + } + } + + return context.WithValue(ctx, addrCtxKey{}, runCtx) +} + +// GetNextResumptionPoints finds the immediate child resumption points for a given parent address. +func GetNextResumptionPoints(ctx context.Context) (map[string]bool, error) { + parentAddr := GetCurrentAddress(ctx) + + rInfo, exists := getResumeInfo(ctx) + if !exists { + return nil, fmt.Errorf("GetNextResumptionPoints: failed to get resume info from context") + } + + rInfo.mu.RLock() + defer rInfo.mu.RUnlock() + + nextPoints := make(map[string]bool) + parentAddrLen := len(parentAddr) + + for _, addr := range rInfo.id2Addr { + // Check if addr is a potential child (must be longer than parent) + if len(addr) <= parentAddrLen { + continue + } + + // Check if it has the parent address as a prefix + var isPrefix bool + if parentAddrLen == 0 { + isPrefix = true + } else { + isPrefix = addr[:parentAddrLen].Equals(parentAddr) + } + + if !isPrefix { + continue + } + + // We are looking for immediate children. + // The address of an immediate child should be one segment longer. + childAddr := addr[parentAddrLen : parentAddrLen+1] + childID := childAddr[0].ID + + // Avoid adding duplicates. + if _, ok := nextPoints[childID]; !ok { + nextPoints[childID] = true + } + } + + return nextPoints, nil +} + +// BatchResumeWithData is the core function for preparing a resume context. It injects a map +// of resume targets and their corresponding data into the context. +// +// The `resumeData` map should contain the interrupt IDs (which are the string form of addresses) of the +// components to be resumed as keys. The value can be the resume data for that component, or `nil` +// if no data is needed (equivalent to using `Resume`). +// +// This function is the foundation for the "Explicit Targeted Resume" strategy. Components whose interrupt IDs +// are present as keys in the map will receive `isResumeFlow = true` when they call `GetResumeContext`. +func BatchResumeWithData(ctx context.Context, resumeData map[string]any) context.Context { + rInfo, ok := ctx.Value(globalResumeInfoKey{}).(*globalResumeInfo) + if !ok { + // Create a new globalResumeInfo and copy the map to prevent external mutation. + newMap := make(map[string]any, len(resumeData)) + for k, v := range resumeData { + newMap[k] = v + } + return context.WithValue(ctx, globalResumeInfoKey{}, &globalResumeInfo{ + id2ResumeData: newMap, + id2ResumeDataUsed: make(map[string]bool), + id2StateUsed: make(map[string]bool), + }) + } + + rInfo.mu.Lock() + defer rInfo.mu.Unlock() + if rInfo.id2ResumeData == nil { + rInfo.id2ResumeData = make(map[string]any) + } + for id, data := range resumeData { + rInfo.id2ResumeData[id] = data + } + return ctx +} + +func PopulateInterruptState(ctx context.Context, id2Addr map[string]Address, + id2State map[string]InterruptState) context.Context { + rInfo, ok := ctx.Value(globalResumeInfoKey{}).(*globalResumeInfo) + if ok { + rInfo.mu.Lock() + defer rInfo.mu.Unlock() + + if rInfo.id2Addr == nil { + rInfo.id2Addr = make(map[string]Address) + } + for id, addr := range id2Addr { + rInfo.id2Addr[id] = addr + } + if rInfo.id2State == nil { + rInfo.id2State = make(map[string]InterruptState) + } + for id, state := range id2State { + rInfo.id2State[id] = state + } + } else { + rInfo = &globalResumeInfo{ + id2Addr: id2Addr, + id2State: id2State, + id2StateUsed: make(map[string]bool), + id2ResumeDataUsed: make(map[string]bool), + } + ctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo) + } + + runCtx, ok := getRunCtx(ctx) + if ok { + for id_, addr := range id2Addr { + if addr.Equals(runCtx.addr) { + if used, ok := rInfo.id2StateUsed[id_]; !ok || !used { + runCtx.interruptState = generic.PtrOf(rInfo.id2State[id_]) + rInfo.id2StateUsed[id_] = true + } + + if used, ok := rInfo.id2ResumeDataUsed[id_]; !ok || !used { + runCtx.isResumeTarget = true + runCtx.resumeData = rInfo.id2ResumeData[id_] + rInfo.id2ResumeDataUsed[id_] = true + } + + break + } + } + } + + return ctx +} + +func getResumeInfo(ctx context.Context) (*globalResumeInfo, bool) { + info, ok := ctx.Value(globalResumeInfoKey{}).(*globalResumeInfo) + return info, ok +} + +type InterruptInfo struct { + Info any + IsRootCause bool +} + +func (i *InterruptInfo) String() string { + if i == nil { + return "" + } + return fmt.Sprintf("interrupt info: Info=%v, IsRootCause=%v", i.Info, i.IsRootCause) +} diff --git a/internal/core/interrupt.go b/internal/core/interrupt.go new file mode 100644 index 0000000..38ddbda --- /dev/null +++ b/internal/core/interrupt.go @@ -0,0 +1,349 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package core + +import ( + "context" + "fmt" + "reflect" + + "github.com/google/uuid" +) + +type CheckPointStore interface { + Get(ctx context.Context, checkPointID string) ([]byte, bool, error) + Set(ctx context.Context, checkPointID string, checkPoint []byte) error +} + +// CheckPointDeleter is an optional interface that CheckPointStore implementations +// can implement to support explicit checkpoint deletion. +// +// If the Store does not implement this interface, stale checkpoints will NOT be +// automatically cleaned up. The store owner is responsible for managing checkpoint +// lifecycle in that case (e.g., via TTL, external cleanup, or implementing this +// interface). +type CheckPointDeleter interface { + Delete(ctx context.Context, checkPointID string) error +} + +type InterruptSignal struct { + ID string + Address + InterruptInfo + InterruptState + Subs []*InterruptSignal +} + +func (is *InterruptSignal) Error() string { + return fmt.Sprintf("interrupt signal: ID=%s, Addr=%s, Info=%s, State=%s, SubsLen=%d", + is.ID, is.Address.String(), is.InterruptInfo.String(), is.InterruptState.String(), len(is.Subs)) +} + +type InterruptState struct { + State any + LayerSpecificPayload any +} + +func (is *InterruptState) String() string { + if is == nil { + return "" + } + return fmt.Sprintf("interrupt state: State=%v, LayerSpecificPayload=%v", is.State, is.LayerSpecificPayload) +} + +// InterruptConfig holds optional parameters for creating an interrupt. +type InterruptConfig struct { + LayerPayload any +} + +// InterruptOption is a function that configures an InterruptConfig. +type InterruptOption func(*InterruptConfig) + +// WithLayerPayload creates an option to attach layer-specific metadata +// to the interrupt's state. +func WithLayerPayload(payload any) InterruptOption { + return func(c *InterruptConfig) { + c.LayerPayload = payload + } +} + +func Interrupt(ctx context.Context, info any, state any, subContexts []*InterruptSignal, opts ...InterruptOption) ( + *InterruptSignal, error) { + addr := GetCurrentAddress(ctx) + + // Apply options to get config + config := &InterruptConfig{} + for _, opt := range opts { + opt(config) + } + + myPoint := InterruptInfo{ + Info: info, + } + + if len(subContexts) == 0 { + myPoint.IsRootCause = true + return &InterruptSignal{ + ID: uuid.NewString(), + Address: addr, + InterruptInfo: myPoint, + InterruptState: InterruptState{ + State: state, + LayerSpecificPayload: config.LayerPayload, + }, + }, nil + } + + return &InterruptSignal{ + ID: uuid.NewString(), + Address: addr, + InterruptInfo: myPoint, + InterruptState: InterruptState{ + State: state, + LayerSpecificPayload: config.LayerPayload, + }, + Subs: subContexts, + }, nil +} + +// InterruptCtx provides a complete, user-facing context for a single, resumable interrupt point. +type InterruptCtx struct { + // ID is the unique, fully-qualified address of the interrupt point. + // It is constructed by joining the individual Address segments, e.g., "agent:A;node:graph_a;tool:tool_call_123". + // This ID should be used when providing resume data via ResumeWithData. + ID string + // Address is the structured sequence of AddressSegment segments that leads to the interrupt point. + Address Address + // Info is the user-facing information associated with the interrupt, provided by the component that triggered it. + Info any + // IsRootCause indicates whether the interrupt point is the exact root cause for an interruption. + IsRootCause bool + // Parent points to the context of the parent component in the interrupt chain. + // It is nil for the top-level interrupt. + Parent *InterruptCtx +} + +func (ic *InterruptCtx) EqualsWithoutID(other *InterruptCtx) bool { + if ic == nil && other == nil { + return true + } + + if ic == nil || other == nil { + return false + } + + if !ic.Address.Equals(other.Address) { + return false + } + + if ic.IsRootCause != other.IsRootCause { + return false + } + + if ic.Info != nil || other.Info != nil { + if ic.Info == nil || other.Info == nil { + return false + } + + if !reflect.DeepEqual(ic.Info, other.Info) { + return false + } + } + + if ic.Parent != nil || other.Parent != nil { + if ic.Parent == nil || other.Parent == nil { + return false + } + + if !ic.Parent.EqualsWithoutID(other.Parent) { + return false + } + } + + return true +} + +// InterruptContextsProvider is an interface for errors that contain interrupt contexts. +// This allows different packages to check for and extract interrupt contexts from errors +// without needing to know the concrete error type. +type InterruptContextsProvider interface { + GetInterruptContexts() []*InterruptCtx +} + +// FromInterruptContexts converts a list of user-facing InterruptCtx objects into an +// internal InterruptSignal tree. It correctly handles common ancestors and ensures +// that the resulting tree is consistent with the original interrupt chain. +// +// This method is primarily used by components that bridge different execution environments. +// For example, an `adk.AgentTool` might catch an `adk.InterruptInfo`, extract the +// `adk.InterruptCtx` objects from it, and then call this method on each one. The resulting +// error signals are then typically aggregated into a single error using `compose.CompositeInterrupt` +// to be returned from the tool's `InvokableRun` method. +// FromInterruptContexts reconstructs a single InterruptSignal tree from a list of +// user-facing InterruptCtx objects. It correctly merges common ancestors. +func FromInterruptContexts(contexts []*InterruptCtx) *InterruptSignal { + if len(contexts) == 0 { + return nil + } + + signalMap := make(map[string]*InterruptSignal) + var rootSignal *InterruptSignal + + // getOrCreateSignal is a recursive helper that builds the tree bottom-up. + var getOrCreateSignal func(*InterruptCtx) *InterruptSignal + getOrCreateSignal = func(ctx *InterruptCtx) *InterruptSignal { + if ctx == nil { + return nil + } + // If we've already created a signal for this context, return it. + if signal, exists := signalMap[ctx.ID]; exists { + return signal + } + + // Create the signal for the current context. + newSignal := &InterruptSignal{ + ID: ctx.ID, + Address: ctx.Address, + InterruptInfo: InterruptInfo{ + Info: ctx.Info, + IsRootCause: ctx.IsRootCause, + }, + } + signalMap[ctx.ID] = newSignal // Cache it immediately. + + // Recursively ensure the parent exists. If it doesn't, this is the root. + if parentSignal := getOrCreateSignal(ctx.Parent); parentSignal != nil { + parentSignal.Subs = append(parentSignal.Subs, newSignal) + } else { + rootSignal = newSignal + } + return newSignal + } + + // Process all contexts to ensure all branches of the tree are built. + for _, ctx := range contexts { + _ = getOrCreateSignal(ctx) + } + + return rootSignal +} + +// ToInterruptContexts converts the internal InterruptSignal tree into a list of +// user-facing InterruptCtx objects for the root causes of the interruption. +// Each returned context has its Parent field populated (if it has a parent), +// allowing traversal up the interrupt chain. +// +// If allowedSegmentTypes is nil, all segment types are kept and addresses are unchanged. +// If allowedSegmentTypes is provided, it: +// 1. Filters the parent chain to only keep contexts whose leaf segment type is allowed +// 2. Strips non-allowed segment types from all addresses +func ToInterruptContexts(is *InterruptSignal, allowedSegmentTypes []AddressSegmentType) []*InterruptCtx { + if is == nil { + return nil + } + var rootCauseContexts []*InterruptCtx + + var buildContexts func(*InterruptSignal, *InterruptCtx) + buildContexts = func(signal *InterruptSignal, parentCtx *InterruptCtx) { + currentCtx := &InterruptCtx{ + ID: signal.ID, + Address: signal.Address, + Info: signal.InterruptInfo.Info, + IsRootCause: signal.InterruptInfo.IsRootCause, + Parent: parentCtx, + } + + if currentCtx.IsRootCause { + rootCauseContexts = append(rootCauseContexts, currentCtx) + } + + for _, subSignal := range signal.Subs { + buildContexts(subSignal, currentCtx) + } + } + + buildContexts(is, nil) + + if len(allowedSegmentTypes) > 0 { + allowedSet := make(map[AddressSegmentType]bool, len(allowedSegmentTypes)) + for _, t := range allowedSegmentTypes { + allowedSet[t] = true + } + + for _, ctx := range rootCauseContexts { + filterParentChain(ctx, allowedSet) + encapsulateContextAddresses(ctx, allowedSet) + } + } + + return rootCauseContexts +} + +func filterParentChain(ctx *InterruptCtx, allowedSet map[AddressSegmentType]bool) { + if ctx == nil { + return + } + + parent := ctx.Parent + for parent != nil { + if len(parent.Address) > 0 && allowedSet[parent.Address[len(parent.Address)-1].Type] { + break + } + parent = parent.Parent + } + + ctx.Parent = parent + + filterParentChain(parent, allowedSet) +} + +func encapsulateContextAddresses(ctx *InterruptCtx, allowedSet map[AddressSegmentType]bool) { + for c := ctx; c != nil; c = c.Parent { + newAddr := make(Address, 0, len(c.Address)) + for _, seg := range c.Address { + if allowedSet[seg.Type] { + newAddr = append(newAddr, seg) + } + } + c.Address = newAddr + } +} + +// SignalToPersistenceMaps flattens an InterruptSignal tree into two maps suitable for persistence in a checkpoint. +func SignalToPersistenceMaps(is *InterruptSignal) (map[string]Address, map[string]InterruptState) { + id2addr := make(map[string]Address) + id2state := make(map[string]InterruptState) + + if is == nil { + return id2addr, id2state + } + + var traverse func(*InterruptSignal) + traverse = func(signal *InterruptSignal) { + // Add current signal's data to the maps. + id2addr[signal.ID] = signal.Address + id2state[signal.ID] = signal.InterruptState // The embedded struct + + // Recurse into children. + for _, sub := range signal.Subs { + traverse(sub) + } + } + + traverse(is) + return id2addr, id2state +} diff --git a/internal/core/interrupt_test.go b/internal/core/interrupt_test.go new file mode 100644 index 0000000..86d1149 --- /dev/null +++ b/internal/core/interrupt_test.go @@ -0,0 +1,1128 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package core + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +// Define AddressSegmentType constants locally to avoid dependency cycles +const ( + AddressSegmentAgent AddressSegmentType = "agent" + AddressSegmentTool AddressSegmentType = "tool" + AddressSegmentNode AddressSegmentType = "node" +) + +func TestInterruptConversion(t *testing.T) { + // Test Case 1: Simple Chain (A -> B -> C) + t.Run("SimpleChain", func(t *testing.T) { + // Manually construct the user-facing contexts with parent pointers + ctxA := &InterruptCtx{ID: "A", IsRootCause: false} + ctxB := &InterruptCtx{ID: "B", Parent: ctxA, IsRootCause: false} + ctxC := &InterruptCtx{ID: "C", Parent: ctxB, IsRootCause: true} + + // The input to FromInterruptContexts is just the root cause leaf node + contexts := []*InterruptCtx{ctxC} + + // Convert from user-facing contexts to internal signal tree + signal := FromInterruptContexts(contexts) + + // Assertions for the signal tree structure + assert.NotNil(t, signal) + assert.Equal(t, "A", signal.ID) + assert.Len(t, signal.Subs, 1) + assert.Equal(t, "B", signal.Subs[0].ID) + assert.Len(t, signal.Subs[0].Subs, 1) + assert.Equal(t, "C", signal.Subs[0].Subs[0].ID) + assert.True(t, signal.Subs[0].Subs[0].IsRootCause) + + // Convert back from the signal tree to user-facing contexts + finalContexts := ToInterruptContexts(signal, nil) + + // Assertions for the final list of contexts + assert.Len(t, finalContexts, 1) + finalC := finalContexts[0] + assert.Equal(t, "C", finalC.ID) + assert.True(t, finalC.IsRootCause) + assert.NotNil(t, finalC.Parent) + assert.Equal(t, "B", finalC.Parent.ID) + assert.NotNil(t, finalC.Parent.Parent) + assert.Equal(t, "A", finalC.Parent.Parent.ID) + assert.Nil(t, finalC.Parent.Parent.Parent) + }) + + // Test Case 2: Multiple Root Causes with Shared Parent (B -> D, C -> D) + t.Run("MultipleRootsSharedParent", func(t *testing.T) { + // Manually construct the contexts + ctxD := &InterruptCtx{ID: "D", IsRootCause: false} + ctxB := &InterruptCtx{ID: "B", Parent: ctxD, IsRootCause: true} + ctxC := &InterruptCtx{ID: "C", Parent: ctxD, IsRootCause: true} + + // The input contains both root cause leaves + contexts := []*InterruptCtx{ctxB, ctxC} + + // Convert to signal tree + signal := FromInterruptContexts(contexts) + + // Assertions for the signal tree structure (should merge at D) + assert.NotNil(t, signal) + assert.Equal(t, "D", signal.ID) + assert.Len(t, signal.Subs, 2) + // Order of subs is not guaranteed, so we check for presence + subIDs := []string{signal.Subs[0].ID, signal.Subs[1].ID} + assert.Contains(t, subIDs, "B") + assert.Contains(t, subIDs, "C") + + // Convert back to user-facing contexts + finalContexts := ToInterruptContexts(signal, nil) + + // Assertions for the final list of contexts + assert.Len(t, finalContexts, 2) + finalIDs := []string{finalContexts[0].ID, finalContexts[1].ID} + assert.Contains(t, finalIDs, "B") + assert.Contains(t, finalIDs, "C") + + // Check parent linking for one of the branches + var finalB *InterruptCtx + if finalContexts[0].ID == "B" { + finalB = finalContexts[0] + } else { + finalB = finalContexts[1] + } + assert.NotNil(t, finalB.Parent) + assert.Equal(t, "D", finalB.Parent.ID) + assert.Nil(t, finalB.Parent.Parent) + }) + + // Test Case 3: Nil and Empty Inputs + t.Run("NilAndEmpty", func(t *testing.T) { + assert.Nil(t, FromInterruptContexts(nil)) + assert.Nil(t, FromInterruptContexts([]*InterruptCtx{})) + assert.Nil(t, ToInterruptContexts(nil, nil)) + }) +} + +func TestSignalToPersistenceMaps(t *testing.T) { + // Test Case 1: Nil Signal + t.Run("NilSignal", func(t *testing.T) { + id2addr, id2state := SignalToPersistenceMaps(nil) + assert.NotNil(t, id2addr) + assert.NotNil(t, id2state) + assert.Empty(t, id2addr) + assert.Empty(t, id2state) + }) + + // Test Case 2: Single Node Signal + t.Run("SingleNode", func(t *testing.T) { + signal := &InterruptSignal{ + ID: "node1", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + }, + InterruptState: InterruptState{ + State: "test state", + LayerSpecificPayload: "test payload", + }, + } + + id2addr, id2state := SignalToPersistenceMaps(signal) + + assert.Len(t, id2addr, 1) + assert.Len(t, id2state, 1) + + assert.Equal(t, signal.Address, id2addr["node1"]) + assert.Equal(t, signal.InterruptState, id2state["node1"]) + }) + + // Test Case 3: Simple Tree Structure + t.Run("SimpleTree", func(t *testing.T) { + child1 := &InterruptSignal{ + ID: "child1", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + }, + InterruptState: InterruptState{ + State: "child1 state", + }, + } + + child2 := &InterruptSignal{ + ID: "child2", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool2"}, + }, + InterruptState: InterruptState{ + State: "child2 state", + }, + } + + parent := &InterruptSignal{ + ID: "parent", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + }, + InterruptState: InterruptState{ + State: "parent state", + }, + Subs: []*InterruptSignal{child1, child2}, + } + + id2addr, id2state := SignalToPersistenceMaps(parent) + + // Should contain all 3 nodes + assert.Len(t, id2addr, 3) + assert.Len(t, id2state, 3) + + // Check parent node + assert.Equal(t, parent.Address, id2addr["parent"]) + assert.Equal(t, parent.InterruptState, id2state["parent"]) + + // Check child nodes + assert.Equal(t, child1.Address, id2addr["child1"]) + assert.Equal(t, child1.InterruptState, id2state["child1"]) + assert.Equal(t, child2.Address, id2addr["child2"]) + assert.Equal(t, child2.InterruptState, id2state["child2"]) + }) + + // Test Case 4: Deeply Nested Tree + t.Run("DeeplyNestedTree", func(t *testing.T) { + leaf1 := &InterruptSignal{ + ID: "leaf1", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + {Type: AddressSegmentNode, ID: "node1"}, + }, + InterruptState: InterruptState{ + State: "leaf1 state", + }, + } + + leaf2 := &InterruptSignal{ + ID: "leaf2", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + {Type: AddressSegmentNode, ID: "node2"}, + }, + InterruptState: InterruptState{ + State: "leaf2 state", + }, + } + + middle := &InterruptSignal{ + ID: "middle", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + }, + InterruptState: InterruptState{ + State: "middle state", + }, + Subs: []*InterruptSignal{leaf1, leaf2}, + } + + root := &InterruptSignal{ + ID: "root", + Address: Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + }, + InterruptState: InterruptState{ + State: "root state", + }, + Subs: []*InterruptSignal{middle}, + } + + id2addr, id2state := SignalToPersistenceMaps(root) + + // Should contain all 4 nodes + assert.Len(t, id2addr, 4) + assert.Len(t, id2state, 4) + + // Verify all nodes are present + assert.Equal(t, root.Address, id2addr["root"]) + assert.Equal(t, root.InterruptState, id2state["root"]) + assert.Equal(t, middle.Address, id2addr["middle"]) + assert.Equal(t, middle.InterruptState, id2state["middle"]) + assert.Equal(t, leaf1.Address, id2addr["leaf1"]) + assert.Equal(t, leaf1.InterruptState, id2state["leaf1"]) + assert.Equal(t, leaf2.Address, id2addr["leaf2"]) + assert.Equal(t, leaf2.InterruptState, id2state["leaf2"]) + }) + + // Test Case 5: Complex Tree with Multiple Branches + t.Run("ComplexTree", func(t *testing.T) { + // Create a complex tree structure with multiple branches + branch1Leaf1 := &InterruptSignal{ID: "b1l1", Address: Address{{Type: AddressSegmentAgent, ID: "a1"}}, InterruptState: InterruptState{State: "b1l1"}} + branch1Leaf2 := &InterruptSignal{ID: "b1l2", Address: Address{{Type: AddressSegmentAgent, ID: "a1"}}, InterruptState: InterruptState{State: "b1l2"}} + branch1 := &InterruptSignal{ID: "b1", Address: Address{{Type: AddressSegmentAgent, ID: "a1"}}, InterruptState: InterruptState{State: "b1"}, Subs: []*InterruptSignal{branch1Leaf1, branch1Leaf2}} + + branch2Leaf1 := &InterruptSignal{ID: "b2l1", Address: Address{{Type: AddressSegmentAgent, ID: "a1"}}, InterruptState: InterruptState{State: "b2l1"}} + branch2 := &InterruptSignal{ID: "b2", Address: Address{{Type: AddressSegmentAgent, ID: "a1"}}, InterruptState: InterruptState{State: "b2"}, Subs: []*InterruptSignal{branch2Leaf1}} + + root := &InterruptSignal{ID: "root", Address: Address{{Type: AddressSegmentAgent, ID: "a1"}}, InterruptState: InterruptState{State: "root"}, Subs: []*InterruptSignal{branch1, branch2}} + + id2addr, id2state := SignalToPersistenceMaps(root) + + // Should contain all 6 nodes + assert.Len(t, id2addr, 6) + assert.Len(t, id2state, 6) + + // Verify all nodes are present + expectedNodes := []string{"root", "b1", "b2", "b1l1", "b1l2", "b2l1"} + for _, nodeID := range expectedNodes { + assert.Contains(t, id2addr, nodeID) + assert.Contains(t, id2state, nodeID) + } + }) + + // Test Case 6: Empty InterruptState Values + t.Run("EmptyInterruptState", func(t *testing.T) { + signal := &InterruptSignal{ + ID: "node1", + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + InterruptState: InterruptState{ + // Empty state values + }, + } + + id2addr, id2state := SignalToPersistenceMaps(signal) + + assert.Len(t, id2addr, 1) + assert.Len(t, id2state, 1) + assert.Equal(t, signal.Address, id2addr["node1"]) + assert.Equal(t, signal.InterruptState, id2state["node1"]) + }) +} + +func TestGetCurrentAddress(t *testing.T) { + // Test Case 1: No Address in Context + t.Run("NoAddressInContext", func(t *testing.T) { + ctx := context.Background() + addr := GetCurrentAddress(ctx) + assert.Nil(t, addr) + }) + + // Test Case 2: Address in Context + t.Run("AddressInContext", func(t *testing.T) { + ctx := context.Background() + expectedAddr := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + } + + // Create a context with address using internal addrCtx + runCtx := &addrCtx{ + addr: expectedAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + addr := GetCurrentAddress(ctx) + assert.Equal(t, expectedAddr, addr) + }) +} + +func TestGetNextResumptionPoints(t *testing.T) { + // Test Case 1: No Resume Info in Context + t.Run("NoResumeInfo", func(t *testing.T) { + ctx := context.Background() + _, err := GetNextResumptionPoints(ctx) + assert.Error(t, err) + assert.Contains(t, err.Error(), "failed to get resume info") + }) + + // Test Case 2: Empty Resume Info + t.Run("EmptyResumeInfo", func(t *testing.T) { + ctx := context.Background() + rInfo := &globalResumeInfo{ + id2Addr: make(map[string]Address), + } + ctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo) + + points, err := GetNextResumptionPoints(ctx) + assert.NoError(t, err) + assert.Empty(t, points) + }) + + // Test Case 3: Valid Resume Points + t.Run("ValidResumePoints", func(t *testing.T) { + ctx := context.Background() + + // Set up current address + currentAddr := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + } + runCtx := &addrCtx{ + addr: currentAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + // Set up resume info with child addresses + rInfo := &globalResumeInfo{ + id2Addr: map[string]Address{ + "child1": { + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + }, + "child2": { + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool2"}, + }, + "unrelated": { + {Type: AddressSegmentAgent, ID: "agent2"}, + }, + }, + } + ctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo) + + points, err := GetNextResumptionPoints(ctx) + assert.NoError(t, err) + assert.Len(t, points, 2) + assert.True(t, points["tool1"]) + assert.True(t, points["tool2"]) + }) + + // Test Case 4: Root Address (Empty Parent) + t.Run("RootAddress", func(t *testing.T) { + ctx := context.Background() + + // Empty current address (root) + runCtx := &addrCtx{ + addr: Address{}, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + // Set up resume info with various addresses + rInfo := &globalResumeInfo{ + id2Addr: map[string]Address{ + "agent1": { + {Type: AddressSegmentAgent, ID: "agent1"}, + }, + "agent2": { + {Type: AddressSegmentAgent, ID: "agent2"}, + }, + }, + } + ctx = context.WithValue(ctx, globalResumeInfoKey{}, rInfo) + + points, err := GetNextResumptionPoints(ctx) + assert.NoError(t, err) + assert.Len(t, points, 2) + assert.True(t, points["agent1"]) + assert.True(t, points["agent2"]) + }) +} + +func TestBatchResumeWithData(t *testing.T) { + // Test Case 1: New Resume Data + t.Run("NewResumeData", func(t *testing.T) { + ctx := context.Background() + resumeData := map[string]any{ + "id1": "data1", + "id2": "data2", + } + + newCtx := BatchResumeWithData(ctx, resumeData) + + // Verify the data was set correctly + rInfo, ok := newCtx.Value(globalResumeInfoKey{}).(*globalResumeInfo) + assert.True(t, ok) + assert.NotNil(t, rInfo) + assert.Equal(t, "data1", rInfo.id2ResumeData["id1"]) + assert.Equal(t, "data2", rInfo.id2ResumeData["id2"]) + }) + + // Test Case 2: Merge with Existing Resume Data + t.Run("MergeWithExisting", func(t *testing.T) { + ctx := context.Background() + + // First call with initial data + initialData := map[string]any{ + "id1": "initial", + } + ctx = BatchResumeWithData(ctx, initialData) + + // Second call with additional data + additionalData := map[string]any{ + "id2": "additional", + } + newCtx := BatchResumeWithData(ctx, additionalData) + + // Verify both data sets are present + rInfo, ok := newCtx.Value(globalResumeInfoKey{}).(*globalResumeInfo) + assert.True(t, ok) + assert.NotNil(t, rInfo) + assert.Equal(t, "initial", rInfo.id2ResumeData["id1"]) + assert.Equal(t, "additional", rInfo.id2ResumeData["id2"]) + }) + + // Test Case 3: Empty Resume Data + t.Run("EmptyResumeData", func(t *testing.T) { + ctx := context.Background() + newCtx := BatchResumeWithData(ctx, map[string]any{}) + + rInfo, ok := newCtx.Value(globalResumeInfoKey{}).(*globalResumeInfo) + assert.True(t, ok) + assert.NotNil(t, rInfo) + assert.Empty(t, rInfo.id2ResumeData) + }) +} + +func TestGetInterruptState(t *testing.T) { + // Test Case 1: No Interrupt State + t.Run("NoInterruptState", func(t *testing.T) { + ctx := context.Background() + wasInterrupted, hasState, state := GetInterruptState[string](ctx) + assert.False(t, wasInterrupted) + assert.False(t, hasState) + assert.Equal(t, "", state) + }) + + // Test Case 2: With Interrupt State + t.Run("WithInterruptState", func(t *testing.T) { + ctx := context.Background() + + // Create a context with interrupt state + runCtx := &addrCtx{ + interruptState: &InterruptState{ + State: "test state", + }, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + wasInterrupted, hasState, state := GetInterruptState[string](ctx) + assert.True(t, wasInterrupted) + assert.True(t, hasState) + assert.Equal(t, "test state", state) + }) + + // Test Case 3: Wrong Type for Interrupt State + t.Run("WrongType", func(t *testing.T) { + ctx := context.Background() + + // Create a context with interrupt state of wrong type + runCtx := &addrCtx{ + interruptState: &InterruptState{ + State: 123, // int instead of string + }, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + wasInterrupted, hasState, state := GetInterruptState[string](ctx) + assert.True(t, wasInterrupted) + assert.False(t, hasState) // Should be false due to type mismatch + assert.Equal(t, "", state) + }) + + // Test Case 4: Nil Interrupt State + t.Run("NilInterruptState", func(t *testing.T) { + ctx := context.Background() + + // Create a context with nil interrupt state + runCtx := &addrCtx{ + interruptState: nil, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + wasInterrupted, hasState, state := GetInterruptState[string](ctx) + assert.False(t, wasInterrupted) // Should be false because interruptState is nil + assert.False(t, hasState) // Should be false because state is nil + assert.Equal(t, "", state) + }) +} + +func TestGetResumeContext(t *testing.T) { + // Test Case 1: Not Resume Target + t.Run("NotResumeTarget", func(t *testing.T) { + ctx := context.Background() + isResumeTarget, hasData, data := GetResumeContext[string](ctx) + assert.False(t, isResumeTarget) + assert.False(t, hasData) + assert.Equal(t, "", data) + }) + + // Test Case 2: Resume Target with Data + t.Run("ResumeTargetWithData", func(t *testing.T) { + ctx := context.Background() + + // Create a context as resume target with data + runCtx := &addrCtx{ + isResumeTarget: true, + resumeData: "resume data", + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + isResumeTarget, hasData, data := GetResumeContext[string](ctx) + assert.True(t, isResumeTarget) + assert.True(t, hasData) + assert.Equal(t, "resume data", data) + }) + + // Test Case 3: Resume Target without Data + t.Run("ResumeTargetWithoutData", func(t *testing.T) { + ctx := context.Background() + + // Create a context as resume target without data + runCtx := &addrCtx{ + isResumeTarget: true, + resumeData: nil, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + isResumeTarget, hasData, data := GetResumeContext[string](ctx) + assert.True(t, isResumeTarget) + assert.False(t, hasData) + assert.Equal(t, "", data) + }) + + // Test Case 4: Wrong Type for Resume Data + t.Run("WrongType", func(t *testing.T) { + ctx := context.Background() + + // Create a context with resume data of wrong type + runCtx := &addrCtx{ + isResumeTarget: true, + resumeData: 123, // int instead of string + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + isResumeTarget, hasData, data := GetResumeContext[string](ctx) + assert.True(t, isResumeTarget) + assert.False(t, hasData) // Should be false due to type mismatch + assert.Equal(t, "", data) + }) +} + +func TestWithLayerPayload(t *testing.T) { + // Test Case 1: Basic Usage + t.Run("BasicUsage", func(t *testing.T) { + config := &InterruptConfig{} + opt := WithLayerPayload("test payload") + opt(config) + assert.Equal(t, "test payload", config.LayerPayload) + }) + + // Test Case 2: Nil Payload + t.Run("NilPayload", func(t *testing.T) { + config := &InterruptConfig{LayerPayload: "existing"} + opt := WithLayerPayload(nil) + opt(config) + assert.Nil(t, config.LayerPayload) + }) + + // Test Case 3: Complex Payload + t.Run("ComplexPayload", func(t *testing.T) { + config := &InterruptConfig{} + payload := map[string]any{ + "key1": "value1", + "key2": 123, + } + opt := WithLayerPayload(payload) + opt(config) + assert.Equal(t, payload, config.LayerPayload) + }) +} + +func TestInterruptFunction(t *testing.T) { + // Test Case 1: Simple Interrupt without SubContexts + t.Run("SimpleInterrupt", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a mock address + expectedAddr := Address{{Type: AddressSegmentAgent, ID: "test-agent"}} + runCtx := &addrCtx{ + addr: expectedAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + info := "test info" + state := "test state" + + signal, err := Interrupt(ctx, info, state, nil) + assert.NoError(t, err) + assert.NotNil(t, signal) + assert.NotEmpty(t, signal.ID) + assert.Equal(t, info, signal.Info) + assert.Equal(t, state, signal.State) + assert.True(t, signal.IsRootCause) + assert.Equal(t, expectedAddr, signal.Address) + }) + + // Test Case 2: Interrupt with SubContexts + t.Run("InterruptWithSubContexts", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a mock address + expectedAddr := Address{{Type: AddressSegmentAgent, ID: "parent-agent"}} + runCtx := &addrCtx{ + addr: expectedAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + // Create sub contexts + subContexts := []*InterruptSignal{ + { + ID: "child1", + Address: Address{{Type: AddressSegmentAgent, ID: "child1"}}, + }, + { + ID: "child2", + Address: Address{{Type: AddressSegmentAgent, ID: "child2"}}, + }, + } + + info := "parent info" + state := "parent state" + + signal, err := Interrupt(ctx, info, state, subContexts) + assert.NoError(t, err) + assert.NotNil(t, signal) + assert.NotEmpty(t, signal.ID) + assert.Equal(t, info, signal.Info) + assert.Equal(t, state, signal.State) + assert.False(t, signal.IsRootCause) // Should be false when there are sub contexts + assert.Len(t, signal.Subs, 2) + assert.Equal(t, "child1", signal.Subs[0].ID) + assert.Equal(t, "child2", signal.Subs[1].ID) + }) + + // Test Case 3: Interrupt with Options + t.Run("InterruptWithOptions", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a mock address + expectedAddr := Address{{Type: AddressSegmentAgent, ID: "test-agent"}} + runCtx := &addrCtx{ + addr: expectedAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + info := "test info" + state := "test state" + layerPayload := "layer payload" + + signal, err := Interrupt(ctx, info, state, nil, WithLayerPayload(layerPayload)) + assert.NoError(t, err) + assert.NotNil(t, signal) + assert.Equal(t, layerPayload, signal.LayerSpecificPayload) + }) + + // Test Case 4: Empty SubContexts + t.Run("EmptySubContexts", func(t *testing.T) { + ctx := context.Background() + + // Create a context with a mock address + expectedAddr := Address{{Type: AddressSegmentAgent, ID: "test-agent"}} + runCtx := &addrCtx{ + addr: expectedAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + info := "test info" + state := "test state" + + signal, err := Interrupt(ctx, info, state, []*InterruptSignal{}) + assert.NoError(t, err) + assert.NotNil(t, signal) + assert.True(t, signal.IsRootCause) // Should be true when sub contexts is empty + assert.Empty(t, signal.Subs) + }) +} + +func TestAddressMethods(t *testing.T) { + // Test Case 1: Address.String() + t.Run("AddressString", func(t *testing.T) { + addr := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + {Type: AddressSegmentNode, ID: "node1", SubID: "sub1"}, + } + + result := addr.String() + expected := "agent:agent1;tool:tool1;node:node1:sub1" + assert.Equal(t, expected, result) + }) + + // Test Case 2: Address.String() with empty address + t.Run("EmptyAddressString", func(t *testing.T) { + var addr Address + result := addr.String() + assert.Equal(t, "", result) + }) + + // Test Case 3: Address.Equals() with equal addresses + t.Run("AddressEquals", func(t *testing.T) { + addr1 := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + } + addr2 := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + } + + assert.True(t, addr1.Equals(addr2)) + }) + + // Test Case 4: Address.Equals() with different addresses + t.Run("AddressNotEquals", func(t *testing.T) { + addr1 := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + } + addr2 := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool2"}, + } + + assert.False(t, addr1.Equals(addr2)) + }) + + // Test Case 5: Address.Equals() with different lengths + t.Run("AddressDifferentLengths", func(t *testing.T) { + addr1 := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + {Type: AddressSegmentTool, ID: "tool1"}, + } + addr2 := Address{ + {Type: AddressSegmentAgent, ID: "agent1"}, + } + + assert.False(t, addr1.Equals(addr2)) + }) + + // Test Case 6: Address.Equals() with SubID differences + t.Run("AddressSubIDDifference", func(t *testing.T) { + addr1 := Address{ + {Type: AddressSegmentAgent, ID: "agent1", SubID: "sub1"}, + } + addr2 := Address{ + {Type: AddressSegmentAgent, ID: "agent1", SubID: "sub2"}, + } + + assert.False(t, addr1.Equals(addr2)) + }) +} + +func TestAppendAddressSegment(t *testing.T) { + // Test Case 1: Append to empty address + t.Run("AppendToEmpty", func(t *testing.T) { + ctx := context.Background() + + newCtx := AppendAddressSegment(ctx, AddressSegmentAgent, "agent1", "") + + addr := GetCurrentAddress(newCtx) + assert.Len(t, addr, 1) + assert.Equal(t, AddressSegmentAgent, addr[0].Type) + assert.Equal(t, "agent1", addr[0].ID) + assert.Equal(t, "", addr[0].SubID) + }) + + // Test Case 2: Append to existing address + t.Run("AppendToExisting", func(t *testing.T) { + ctx := context.Background() + + // First append + ctx = AppendAddressSegment(ctx, AddressSegmentAgent, "agent1", "") + + // Second append + newCtx := AppendAddressSegment(ctx, AddressSegmentTool, "tool1", "call1") + + addr := GetCurrentAddress(newCtx) + assert.Len(t, addr, 2) + assert.Equal(t, AddressSegmentAgent, addr[0].Type) + assert.Equal(t, "agent1", addr[0].ID) + assert.Equal(t, AddressSegmentTool, addr[1].Type) + assert.Equal(t, "tool1", addr[1].ID) + assert.Equal(t, "call1", addr[1].SubID) + }) + + // Test Case 3: Append with SubID + t.Run("AppendWithSubID", func(t *testing.T) { + ctx := context.Background() + + newCtx := AppendAddressSegment(ctx, AddressSegmentTool, "tool1", "call123") + + addr := GetCurrentAddress(newCtx) + assert.Len(t, addr, 1) + assert.Equal(t, AddressSegmentTool, addr[0].Type) + assert.Equal(t, "tool1", addr[0].ID) + assert.Equal(t, "call123", addr[0].SubID) + }) +} + +func TestPopulateInterruptState(t *testing.T) { + // Test Case 1: Populate with matching address + t.Run("PopulateMatchingAddress", func(t *testing.T) { + ctx := context.Background() + + // Set up current address + currentAddr := Address{{Type: AddressSegmentAgent, ID: "agent1"}} + runCtx := &addrCtx{ + addr: currentAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + // Set up interrupt state data + id2Addr := map[string]Address{ + "interrupt1": currentAddr, + } + id2State := map[string]InterruptState{ + "interrupt1": {State: "test state"}, + } + + newCtx := PopulateInterruptState(ctx, id2Addr, id2State) + + // Verify the state was populated + wasInterrupted, hasState, state := GetInterruptState[string](newCtx) + assert.True(t, wasInterrupted) + assert.True(t, hasState) + assert.Equal(t, "test state", state) + }) + + // Test Case 2: Populate with non-matching address + t.Run("PopulateNonMatchingAddress", func(t *testing.T) { + ctx := context.Background() + + // Set up current address + currentAddr := Address{{Type: AddressSegmentAgent, ID: "agent1"}} + runCtx := &addrCtx{ + addr: currentAddr, + } + ctx = context.WithValue(ctx, addrCtxKey{}, runCtx) + + // Set up interrupt state data with different address + id2Addr := map[string]Address{ + "interrupt1": {{Type: AddressSegmentAgent, ID: "agent2"}}, + } + id2State := map[string]InterruptState{ + "interrupt1": {State: "test state"}, + } + + newCtx := PopulateInterruptState(ctx, id2Addr, id2State) + + // Verify the state was NOT populated (no matching address) + wasInterrupted, hasState, state := GetInterruptState[string](newCtx) + assert.False(t, wasInterrupted) + assert.False(t, hasState) + assert.Equal(t, "", state) + }) + + // Test Case 3: Populate with empty data + t.Run("PopulateEmptyData", func(t *testing.T) { + ctx := context.Background() + + newCtx := PopulateInterruptState(ctx, map[string]Address{}, map[string]InterruptState{}) + + // Verify no state was populated + wasInterrupted, hasState, state := GetInterruptState[string](newCtx) + assert.False(t, wasInterrupted) + assert.False(t, hasState) + assert.Equal(t, "", state) + }) +} + +func TestStringMethods(t *testing.T) { + // Test Case 1: InterruptSignal.Error() + t.Run("InterruptSignalError", func(t *testing.T) { + signal := &InterruptSignal{ + ID: "test-id", + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + InterruptInfo: InterruptInfo{ + Info: "test info", + }, + InterruptState: InterruptState{ + State: "test state", + LayerSpecificPayload: "test payload", + }, + Subs: []*InterruptSignal{ + {ID: "sub1"}, + }, + } + + errorStr := signal.Error() + expectedContains := []string{ + "interrupt signal:", + "ID=test-id", + "Addr=agent:agent1", + "Info=interrupt info: Info=test info, IsRootCause=false", + "State=interrupt state: State=test state, LayerSpecificPayload=test payload", + "SubsLen=1", + } + + for _, expected := range expectedContains { + assert.Contains(t, errorStr, expected) + } + }) + + // Test Case 2: InterruptState.String() + t.Run("InterruptStateString", func(t *testing.T) { + state := &InterruptState{ + State: "test state", + LayerSpecificPayload: "test payload", + } + + result := state.String() + expected := "interrupt state: State=test state, LayerSpecificPayload=test payload" + assert.Equal(t, expected, result) + }) + + // Test Case 3: InterruptState.String() with nil + t.Run("InterruptStateStringNil", func(t *testing.T) { + var state *InterruptState + result := state.String() + assert.Equal(t, "", result) + }) + + // Test Case 4: InterruptInfo.String() + t.Run("InterruptInfoString", func(t *testing.T) { + info := &InterruptInfo{ + Info: "test info", + IsRootCause: true, + } + + result := info.String() + expected := "interrupt info: Info=test info, IsRootCause=true" + assert.Equal(t, expected, result) + }) + + // Test Case 5: InterruptInfo.String() with nil + t.Run("InterruptInfoStringNil", func(t *testing.T) { + var info *InterruptInfo + result := info.String() + assert.Equal(t, "", result) + }) +} + +func TestInterruptCtxEqualsWithoutID(t *testing.T) { + // Test Case 1: Equal contexts + t.Run("EqualContexts", func(t *testing.T) { + ctx1 := &InterruptCtx{ + ID: "id1", + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Info: "info1", + IsRootCause: true, + } + ctx2 := &InterruptCtx{ + ID: "id2", // Different ID should be ignored + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Info: "info1", + IsRootCause: true, + } + + assert.True(t, ctx1.EqualsWithoutID(ctx2)) + }) + + // Test Case 2: Different addresses + t.Run("DifferentAddresses", func(t *testing.T) { + ctx1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + } + ctx2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent2"}}, + } + + assert.False(t, ctx1.EqualsWithoutID(ctx2)) + }) + + // Test Case 3: Different root cause flags + t.Run("DifferentRootCause", func(t *testing.T) { + ctx1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + IsRootCause: true, + } + ctx2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + IsRootCause: false, + } + + assert.False(t, ctx1.EqualsWithoutID(ctx2)) + }) + + // Test Case 4: Different info + t.Run("DifferentInfo", func(t *testing.T) { + ctx1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Info: "info1", + } + ctx2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Info: "info2", + } + + assert.False(t, ctx1.EqualsWithoutID(ctx2)) + }) + + // Test Case 5: Nil contexts + t.Run("NilContexts", func(t *testing.T) { + var ctx1 *InterruptCtx + var ctx2 *InterruptCtx + + assert.True(t, ctx1.EqualsWithoutID(ctx2)) + + ctx3 := &InterruptCtx{} + assert.False(t, ctx1.EqualsWithoutID(ctx3)) + assert.False(t, ctx3.EqualsWithoutID(ctx1)) + }) + + // Test Case 6: With parent contexts + t.Run("WithParentContexts", func(t *testing.T) { + parent1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "parent"}}, + } + parent2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "parent"}}, + } + + ctx1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Parent: parent1, + } + ctx2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Parent: parent2, + } + + assert.True(t, ctx1.EqualsWithoutID(ctx2)) + }) + + // Test Case 7: Different parent contexts + t.Run("DifferentParentContexts", func(t *testing.T) { + parent1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "parent1"}}, + } + parent2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "parent2"}}, + } + + ctx1 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Parent: parent1, + } + ctx2 := &InterruptCtx{ + Address: Address{{Type: AddressSegmentAgent, ID: "agent1"}}, + Parent: parent2, + } + + assert.False(t, ctx1.EqualsWithoutID(ctx2)) + }) +} diff --git a/internal/core/resume.go b/internal/core/resume.go new file mode 100644 index 0000000..c638bf2 --- /dev/null +++ b/internal/core/resume.go @@ -0,0 +1,109 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package core + +import "context" + +// GetInterruptState provides a type-safe way to check for and retrieve the persisted state from a previous interruption. +// It is the primary function a component should use to understand its past state. +// +// It returns three values: +// - wasInterrupted (bool): True if the node was part of a previous interruption, regardless of whether state was provided. +// - state (T): The typed state object, if it was provided and matches type `T`. +// - hasState (bool): True if state was provided during the original interrupt and successfully cast to type `T`. +func GetInterruptState[T any](ctx context.Context) (wasInterrupted bool, hasState bool, state T) { + rCtx, ok := getRunCtx(ctx) + if !ok || rCtx.interruptState == nil { + return + } + + wasInterrupted = true + if rCtx.interruptState.State == nil { + return + } + + state, hasState = rCtx.interruptState.State.(T) + return +} + +// GetResumeContext checks if the current component is the target of a resume operation +// and retrieves any data provided by the user for that resumption. +// +// This function is typically called *after* a component has already determined it is in a +// resumed state by calling GetInterruptState. +// +// It returns three values: +// - isResumeTarget: A boolean that is true if the current component's address OR any of its +// descendant addresses was explicitly targeted by a call to Resume() or ResumeWithData(). +// This allows composite components (like tools containing nested graphs) to know they should +// execute their children to reach the actual resume target. +// - hasData: A boolean that is true if data was provided for this specific component (i.e., not nil). +// - data: The typed data provided by the user. +// +// ### How to Use This Function: A Decision Framework +// +// The correct usage pattern depends on the application's desired resume strategy. +// +// #### Strategy 1: Implicit "Resume All" +// In some use cases, any resume operation implies that *all* interrupted points should proceed. +// For example, if an application's UI only provides a single "Continue" button for a set of +// interruptions. In this model, a component can often just use `GetInterruptState` to see if +// `wasInterrupted` is true and then proceed with its logic, as it can assume it is an intended target. +// It may still call `GetResumeContext` to check for optional data, but the `isResumeFlow` flag is less critical. +// +// #### Strategy 2: Explicit "Targeted Resume" (Most Common) +// For applications with multiple, distinct interrupt points that must be resumed independently, it is +// crucial to differentiate which point is being resumed. This is the primary use case for the `isResumeTarget` flag. +// - If `isResumeTarget` is `true`: Your component (or one of its descendants) is the target. +// If `hasData` is true, you are the direct target and should consume the data. +// If `hasData` is false, a descendant is the target—execute your children to reach it. +// - If `isResumeTarget` is `false`: Neither you nor your descendants are the target. You MUST +// re-interrupt (e.g., by returning `StatefulInterrupt(...)`) to preserve your state. +// +// ### Guidance for Composite Components +// +// Composite components (like `Graph` or other `Runnable`s that contain sub-processes) have a dual role: +// 1. Check for Self-Targeting: A composite component can itself be the target of a resume +// operation, for instance, to modify its internal state. It may call `GetResumeContext` +// to check for data targeted at its own address. +// 2. Act as a Conduit: After checking for itself, its primary role is to re-execute its children, +// allowing the resume context to flow down to them. It must not consume a resume signal +// intended for one of its descendants. +func GetResumeContext[T any](ctx context.Context) (isResumeTarget bool, hasData bool, data T) { + rCtx, ok := getRunCtx(ctx) + if !ok { + return + } + + isResumeTarget = rCtx.isResumeTarget + if !isResumeTarget { + return + } + + // It is a resume flow, now check for data + if rCtx.resumeData == nil { + return // hasData is false + } + + data, hasData = rCtx.resumeData.(T) + return +} + +func getRunCtx(ctx context.Context) (*addrCtx, bool) { + rCtx, ok := ctx.Value(addrCtxKey{}).(*addrCtx) + return rCtx, ok +} diff --git a/internal/generic/generic.go b/internal/generic/generic.go new file mode 100644 index 0000000..2f3293c --- /dev/null +++ b/internal/generic/generic.go @@ -0,0 +1,90 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package generic + +import ( + "reflect" +) + +// NewInstance create an instance of the given type T. +// the main purpose of this function is to create an instance of a type, can handle the type of T is a pointer or not. +// eg. NewInstance[int] returns 0. +// eg. NewInstance[*int] returns *0 (will be ptr of 0, not nil!). +func NewInstance[T any]() T { + typ := TypeOf[T]() + + switch typ.Kind() { + case reflect.Map: + return reflect.MakeMap(typ).Interface().(T) + case reflect.Slice, reflect.Array: + return reflect.MakeSlice(typ, 0, 0).Interface().(T) + case reflect.Ptr: + typ = typ.Elem() + origin := reflect.New(typ) + inst := origin + + for typ.Kind() == reflect.Ptr { + typ = typ.Elem() + inst = inst.Elem() + inst.Set(reflect.New(typ)) + } + + return origin.Interface().(T) + default: + var t T + return t + } +} + +// TypeOf returns the type of T. +// eg. TypeOf[int] returns reflect.TypeOf(int). +// eg. TypeOf[*int] returns reflect.TypeOf(*int). +func TypeOf[T any]() reflect.Type { + return reflect.TypeOf((*T)(nil)).Elem() +} + +// PtrOf returns a pointer of T. +// useful when you want to get a pointer of a value, in some config, for example. +// eg. PtrOf[int] returns *int. +// eg. PtrOf[*int] returns **int. +func PtrOf[T any](v T) *T { + return &v +} + +type Pair[F, S any] struct { + First F + Second S +} + +// Reverse returns a new slice with elements in reversed order. +func Reverse[S ~[]E, E any](s S) S { + d := make(S, len(s)) + for i := 0; i < len(s); i++ { + d[i] = s[len(s)-i-1] + } + + return d +} + +// CopyMap copies a map to a new map. +func CopyMap[K comparable, V any](src map[K]V) map[K]V { + dst := make(map[K]V, len(src)) + for k, v := range src { + dst[k] = v + } + return dst +} diff --git a/internal/generic/generic_test.go b/internal/generic/generic_test.go new file mode 100644 index 0000000..a44ee53 --- /dev/null +++ b/internal/generic/generic_test.go @@ -0,0 +1,104 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package generic + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNewInstance(t *testing.T) { + t.Run("struct", func(t *testing.T) { + type Test struct{} + + inst := NewInstance[Test]() + + assert.IsType(t, Test{}, inst) + }) + + t.Run("pointer", func(t *testing.T) { + type Test struct{} + + inst := NewInstance[*Test]() + + assert.IsType(t, &Test{}, inst) + }) + + t.Run("interface", func(t *testing.T) { + type Test interface{} + + inst := NewInstance[Test]() + assert.IsType(t, Test(nil), inst) + }) + + t.Run("pointer of pointer of pointer", func(t *testing.T) { + type Test struct { + Value int + } + + inst := NewInstance[***Test]() + + ptr := &Test{} + ptrOfPtr := &ptr + assert.NotNil(t, inst) + assert.NotNil(t, *inst) + assert.IsType(t, ptrOfPtr, *inst) + assert.NotNil(t, **inst) + assert.Equal(t, Test{Value: 0}, ***inst) + }) + + t.Run("primitive_map", func(t *testing.T) { + inst := NewInstance[map[string]any]() + assert.NotNil(t, inst) + inst["a"] = 1 + assert.Equal(t, map[string]any{"a": 1}, inst) + }) + + t.Run("primitive_slice", func(t *testing.T) { + inst := NewInstance[[]int]() + assert.NotNil(t, inst) + inst = append(inst, 1) + assert.Equal(t, []int{1}, inst) + }) + + t.Run("primitive_string", func(t *testing.T) { + inst := NewInstance[string]() + assert.Equal(t, "", inst) + }) + + t.Run("primitive_int64", func(t *testing.T) { + inst := NewInstance[int64]() + assert.Equal(t, int64(0), inst) + }) +} + +func TestReverse(t *testing.T) { + t.Run("reverse int slice", func(t *testing.T) { + input := []int{1, 2, 3, 4, 5} + expected := []int{5, 4, 3, 2, 1} + result := Reverse(input) + assert.Equal(t, expected, result) + }) + + t.Run("reverse string slice", func(t *testing.T) { + input := []string{"a", "b", "c"} + expected := []string{"c", "b", "a"} + result := Reverse(input) + assert.Equal(t, expected, result) + }) +} diff --git a/internal/generic/type_name.go b/internal/generic/type_name.go new file mode 100644 index 0000000..4829e74 --- /dev/null +++ b/internal/generic/type_name.go @@ -0,0 +1,71 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package generic + +import ( + "reflect" + "regexp" + "runtime" + "strings" +) + +var ( + regOfAnonymousFunc = regexp.MustCompile(`^func[0-9]+`) + regOfNumber = regexp.MustCompile(`^\d+$`) +) + +// ParseTypeName returns the name of the type of the given value. +// It takes a reflect.Value as input and processes it to determine the underlying type. If the type is a pointer, it dereferences it to get the actual type. (the optimization of this function) +// eg: ParseTypeName(reflect.ValueOf(&&myStruct{})) returns "myStruct" (not "**myStruct") +// +// If the type is a function, it retrieves the function's name, handling both named and anonymous functions. +// examples of function paths: [package_path].[receiver_type].[func_name] +// named function: xxx/utils.ParseTypeName +// method: xxx/utils.(*MyStruct).Method +// anonymous function: xxx/utils.TestParseTypeName.func6.1 +func ParseTypeName(val reflect.Value) string { + typ := val.Type() + + for typ.Kind() == reflect.Pointer { + typ = typ.Elem() + } + + if typ.Kind() == reflect.Func { + funcName := runtime.FuncForPC(val.Pointer()).Name() + idx := strings.LastIndex(funcName, ".") + if idx < 0 { + if funcName != "" { + return funcName + } + return "" + } + + name := funcName[idx+1:] + + if regOfAnonymousFunc.MatchString(name) { + return "" + } + + if regOfNumber.MatchString(name) { + return "" + } + + return name + } + + return typ.Name() +} diff --git a/internal/generic/type_name_test.go b/internal/generic/type_name_test.go new file mode 100644 index 0000000..4230174 --- /dev/null +++ b/internal/generic/type_name_test.go @@ -0,0 +1,86 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package generic + +import ( + "reflect" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestParseTypeName(t *testing.T) { + t.Run("named_struct", func(t *testing.T) { + type OpenAI struct{} + model := &OpenAI{} + name := ParseTypeName(reflect.Indirect(reflect.ValueOf(model))) + assert.Equal(t, "OpenAI", name) + }) + + t.Run("anonymous_struct", func(t *testing.T) { + model := &struct{}{} + name := ParseTypeName(reflect.ValueOf(model)) + assert.Equal(t, "", name) + }) + + t.Run("anonymous_struct_from_func", func(t *testing.T) { + model := genStruct() + name := ParseTypeName(reflect.ValueOf(model)) + assert.Equal(t, "", name) + }) + + t.Run("named_interface", func(t *testing.T) { + type OpenAI interface{} + model := OpenAI(&struct{}{}) + name := ParseTypeName(reflect.ValueOf(model)) + assert.Equal(t, "", name) + + name = ParseTypeName(reflect.ValueOf((*OpenAI)(nil))) + assert.Equal(t, "OpenAI", name) + }) + + t.Run("named_function", func(t *testing.T) { + f := genOpenAI + name := ParseTypeName(reflect.ValueOf(f)) + assert.Equal(t, "genOpenAI", name) + }) + + t.Run("anonymous_function", func(t *testing.T) { + f := genAnonymousFunc() + name := ParseTypeName(reflect.ValueOf(f)) + assert.Equal(t, "", name) + + ff := func(n string) { + _ = n + } + + name = ParseTypeName(reflect.ValueOf(ff)) + assert.Equal(t, "", name) + }) +} + +func genStruct() *struct{} { + return &struct{}{} +} + +func genOpenAI() {} + +func genAnonymousFunc() func(n string) { + return func(n string) { + _ = n + } +} diff --git a/internal/gmap/gmap.go b/internal/gmap/gmap.go new file mode 100644 index 0000000..7eae8b5 --- /dev/null +++ b/internal/gmap/gmap.go @@ -0,0 +1,121 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gmap + +// Concat returns the unions of maps as a new map. +// +// 💡 NOTE: +// +// - Once the key conflicts, the newer value always replace the older one ([DiscardOld]), +// - If the result is an empty set, always return an empty map instead of nil +// +// 🚀 EXAMPLE: +// +// m := map[int]int{1: 1, 2: 2} +// Concat(m, nil) ⏩ map[int]int{1: 1, 2: 2} +// Concat(m, map[int]{3: 3}) ⏩ map[int]int{1: 1, 2: 2, 3: 3} +// Concat(m, map[int]{2: -1}) ⏩ map[int]int{1: 1, 2: -1} // "2:2" is replaced by the newer "2:-1" +// +// 💡 AKA: Merge, Union, Combine +func Concat[K comparable, V any](ms ...map[K]V) map[K]V { + // FastPath: no map or only one map given. + if len(ms) == 0 { + return make(map[K]V) + } + if len(ms) == 1 { + return cloneWithoutNilCheck(ms[0]) + } + + var maxLen int + for _, m := range ms { + if len(m) > maxLen { + maxLen = len(m) + } + } + ret := make(map[K]V, maxLen) + // FastPath: all maps are empty. + if maxLen == 0 { + return ret + } + + // Concat all maps. + for _, m := range ms { + for k, v := range m { + ret[k] = v + } + } + return ret +} + +// Map applies function f to each key and value of map m. +// Results of f are returned as a new map. +// +// 🚀 EXAMPLE: +// +// f := func(k, v int) (string, string) { return strconv.Itoa(k), strconv.Itoa(v) } +// Map(map[int]int{1: 1}, f) ⏩ map[string]string{"1": "1"} +// Map(map[int]int{}, f) ⏩ map[string]string{} +func Map[K1, K2 comparable, V1, V2 any](m map[K1]V1, f func(K1, V1) (K2, V2)) map[K2]V2 { + r := make(map[K2]V2, len(m)) + for k, v := range m { + k2, v2 := f(k, v) + r[k2] = v2 + } + return r +} + +// Values returns the values of the map m. +// +// 🚀 EXAMPLE: +// +// m := map[int]string{1: "1", 2: "2", 3: "3", 4: "4"} +// Values(m) ⏩ []string{"1", "4", "2", "3"} //⚠️INDETERMINATE ORDER⚠️ +// +// ⚠️ WARNING: The keys values be in an indeterminate order, +func Values[K comparable, V any](m map[K]V) []V { + r := make([]V, 0, len(m)) + for _, v := range m { + r = append(r, v) + } + return r +} + +// Clone returns a shallow copy of map. +// If the given map is nil, nil is returned. +// +// 🚀 EXAMPLE: +// +// Clone(map[int]int{1: 1, 2: 2}) ⏩ map[int]int{1: 1, 2: 2} +// Clone(map[int]int{}) ⏩ map[int]int{} +// Clone[int, int](nil) ⏩ nil +// +// 💡 HINT: Both keys and values are copied using assignment (=), so this is a shallow clone. +// 💡 AKA: Copy +func Clone[K comparable, V any, M ~map[K]V](m M) M { + if m == nil { + return nil + } + return cloneWithoutNilCheck(m) +} + +func cloneWithoutNilCheck[K comparable, V any, M ~map[K]V](m M) M { + r := make(M, len(m)) + for k, v := range m { + r[k] = v + } + return r +} diff --git a/internal/gmap/gmap_test.go b/internal/gmap/gmap_test.go new file mode 100644 index 0000000..1c0f1ef --- /dev/null +++ b/internal/gmap/gmap_test.go @@ -0,0 +1,89 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gmap + +import ( + "fmt" + "sort" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestMerge(t *testing.T) { + assert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, + Concat(map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, nil)) + assert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, + Concat[int, int](nil, map[int]int{1: 1, 2: 2, 3: 3, 4: 4})) + assert.Equal(t, map[int]int{}, Concat[int, int](nil, nil)) + assert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, + Concat(map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, map[int]int{1: 1, 2: 2, 3: 3, 4: 4})) + assert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, + Concat(map[int]int{1: 0, 2: 0}, map[int]int{1: 1, 2: 2, 3: 3, 4: 4})) + assert.Equal(t, map[int]int{1: 1, 2: 2, 3: 3, 4: 4}, + Concat(map[int]int{1: 1, 2: 1}, map[int]int{2: 2, 3: 3, 4: 4})) +} + +func TestMap(t *testing.T) { + assert.Equal(t, + map[string]string{"1": "1", "2": "2"}, + Map(map[int]int{1: 1, 2: 2}, func(k, v int) (string, string) { + return strconv.Itoa(k), strconv.Itoa(v) + })) + assert.Equal(t, + map[string]string{}, + Map(map[int]int{}, func(k, v int) (string, string) { + return strconv.Itoa(k), strconv.Itoa(v) + })) +} + +func TestValues(t *testing.T) { + { + keys := Values(map[int]string{1: "1", 2: "2", 3: "3", 4: "4"}) + sort.Strings(keys) + assert.Equal(t, []string{"1", "2", "3", "4"}, keys) + } + assert.Equal(t, []string{}, Values(map[int]string{})) + assert.Equal(t, []string{}, Values[int, string](nil)) +} + +func TestClone(t *testing.T) { + assert.Equal(t, map[int]int{1: 1, 2: 2}, Clone(map[int]int{1: 1, 2: 2})) + var nilMap map[int]int + assert.Equal(t, map[int]int{}, Clone(map[int]int{})) + assert.NotEqual(t, (map[int]int)(nil), Clone(map[int]int{})) + assert.Equal(t, (map[int]int)(nil), Clone(nilMap)) + assert.NotEqual(t, map[int]int{}, Clone(nilMap)) + + // Test new type. + type I2I map[int]int + assert.Equal(t, I2I{1: 1, 2: 2}, Clone(I2I{1: 1, 2: 2})) + assert.Equal(t, "gmap.I2I", fmt.Sprintf("%T", Clone(I2I{}))) + + // Test shallow clone. + src := map[int]*int{1: ptr(1), 2: ptr(2)} + dst := Clone(src) + assert.Equal(t, src, dst) + assert.True(t, src[1] == dst[1]) + assert.True(t, src[2] == dst[2]) +} + +// Ptr returns a pointer to the given value. +func ptr[T any](v T) *T { + return &v +} diff --git a/internal/gslice/gslice.go b/internal/gslice/gslice.go new file mode 100644 index 0000000..a9c05fe --- /dev/null +++ b/internal/gslice/gslice.go @@ -0,0 +1,39 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gslice + +// ToMap collects elements of slice to map, both map keys and values are produced +// by mapping function f. +// +// 🚀 EXAMPLE: +// +// type Foo struct { +// ID int +// Name string +// } +// mapper := func(f Foo) (int, string) { return f.ID, f.Name } +// ToMap([]Foo{}, mapper) ⏩ map[int]string{} +// s := []Foo{{1, "one"}, {2, "two"}, {3, "three"}} +// ToMap(s, mapper) ⏩ map[int]string{1: "one", 2: "two", 3: "three"} +func ToMap[T, V any, K comparable](s []T, f func(T) (K, V)) map[K]V { + m := make(map[K]V, len(s)) + for _, e := range s { + k, v := f(e) + m[k] = v + } + return m +} diff --git a/internal/gslice/gslice_test.go b/internal/gslice/gslice_test.go new file mode 100644 index 0000000..b55c854 --- /dev/null +++ b/internal/gslice/gslice_test.go @@ -0,0 +1,36 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gslice + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestToMap(t *testing.T) { + type Foo struct { + ID int + Name string + } + mapper := func(f Foo) (int, string) { return f.ID, f.Name } + assert.Equal(t, map[int]string{}, ToMap([]Foo{}, mapper)) + assert.Equal(t, map[int]string{}, ToMap(nil, mapper)) + assert.Equal(t, + map[int]string{1: "one", 2: "two", 3: "three"}, + ToMap([]Foo{{1, "one"}, {2, "two"}, {3, "three"}}, mapper)) +} diff --git a/internal/merge.go b/internal/merge.go new file mode 100644 index 0000000..6d0f94d --- /dev/null +++ b/internal/merge.go @@ -0,0 +1,81 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package internal + +import ( + "fmt" + "reflect" + + "github.com/cloudwego/eino/internal/generic" +) + +var mergeFuncs = map[reflect.Type]any{} + +func RegisterValuesMergeFunc[T any](fn func([]T) (T, error)) { + mergeFuncs[generic.TypeOf[T]()] = fn +} + +func GetMergeFunc(typ reflect.Type) func([]any) (any, error) { + if fn, ok := mergeFuncs[typ]; ok { + return func(vs []any) (any, error) { + rvs := reflect.MakeSlice(reflect.SliceOf(typ), 0, len(vs)) + for _, v := range vs { + if t := reflect.TypeOf(v); t != typ { + return nil, fmt.Errorf( + "(values merge) field type mismatch. expected: '%v', got: '%v'", typ, t) + } + rvs = reflect.Append(rvs, reflect.ValueOf(v)) + } + + rets := reflect.ValueOf(fn).Call([]reflect.Value{rvs}) + var err error + if !rets[1].IsNil() { + err = rets[1].Interface().(error) + } + return rets[0].Interface(), err + } + } + + if typ.Kind() == reflect.Map { + return func(vs []any) (any, error) { + return mergeMap(typ, vs) + } + } + + return nil +} + +func mergeMap(typ reflect.Type, vs []any) (any, error) { + merged := reflect.MakeMap(typ) + for _, v := range vs { + if t := reflect.TypeOf(v); t != typ { + return nil, fmt.Errorf( + "(values merge map) field type mismatch. expected: '%v', got: '%v'", typ, t) + } + + iter := reflect.ValueOf(v).MapRange() + for iter.Next() { + key, val := iter.Key(), iter.Value() + if merged.MapIndex(key).IsValid() { + return nil, fmt.Errorf("(values merge map) duplicated key ('%v') found", key.Interface()) + } + merged.SetMapIndex(key, val) + } + } + + return merged.Interface(), nil +} diff --git a/internal/mock/adk/Agent_mock.go b/internal/mock/adk/Agent_mock.go new file mode 100644 index 0000000..9b98c2d --- /dev/null +++ b/internal/mock/adk/Agent_mock.go @@ -0,0 +1,171 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go +// +// Generated by this command: +// +// mockgen -destination ../internal/mock/adk/Agent_mock.go --package adk -source interface.go +// + +// Package adk is a generated GoMock package. +package adk + +import ( + context "context" + reflect "reflect" + + adk "github.com/cloudwego/eino/adk" + gomock "go.uber.org/mock/gomock" +) + +// MockAgent is a mock of Agent interface. +type MockAgent struct { + ctrl *gomock.Controller + recorder *MockAgentMockRecorder + isgomock struct{} +} + +// MockAgentMockRecorder is the mock recorder for MockAgent. +type MockAgentMockRecorder struct { + mock *MockAgent +} + +// NewMockAgent creates a new mock instance. +func NewMockAgent(ctrl *gomock.Controller) *MockAgent { + mock := &MockAgent{ctrl: ctrl} + mock.recorder = &MockAgentMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockAgent) EXPECT() *MockAgentMockRecorder { + return m.recorder +} + +// Description mocks base method. +func (m *MockAgent) Description(ctx context.Context) string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Description", ctx) + ret0, _ := ret[0].(string) + return ret0 +} + +// Description indicates an expected call of Description. +func (mr *MockAgentMockRecorder) Description(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Description", reflect.TypeOf((*MockAgent)(nil).Description), ctx) +} + +// Name mocks base method. +func (m *MockAgent) Name(ctx context.Context) string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Name", ctx) + ret0, _ := ret[0].(string) + return ret0 +} + +// Name indicates an expected call of Name. +func (mr *MockAgentMockRecorder) Name(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Name", reflect.TypeOf((*MockAgent)(nil).Name), ctx) +} + +// Run mocks base method. +func (m *MockAgent) Run(ctx context.Context, input *adk.AgentInput, options ...adk.AgentRunOption) *adk.AsyncIterator[*adk.AgentEvent] { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range options { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Run", varargs...) + ret0, _ := ret[0].(*adk.AsyncIterator[*adk.AgentEvent]) + return ret0 +} + +// Run indicates an expected call of Run. +func (mr *MockAgentMockRecorder) Run(ctx, input any, options ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, options...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Run", reflect.TypeOf((*MockAgent)(nil).Run), varargs...) +} + +// MockOnSubAgents is a mock of OnSubAgents interface. +type MockOnSubAgents struct { + ctrl *gomock.Controller + recorder *MockOnSubAgentsMockRecorder + isgomock struct{} +} + +// MockOnSubAgentsMockRecorder is the mock recorder for MockOnSubAgents. +type MockOnSubAgentsMockRecorder struct { + mock *MockOnSubAgents +} + +// NewMockOnSubAgents creates a new mock instance. +func NewMockOnSubAgents(ctrl *gomock.Controller) *MockOnSubAgents { + mock := &MockOnSubAgents{ctrl: ctrl} + mock.recorder = &MockOnSubAgentsMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockOnSubAgents) EXPECT() *MockOnSubAgentsMockRecorder { + return m.recorder +} + +// OnDisallowTransferToParent mocks base method. +func (m *MockOnSubAgents) OnDisallowTransferToParent(ctx context.Context) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OnDisallowTransferToParent", ctx) + ret0, _ := ret[0].(error) + return ret0 +} + +// OnDisallowTransferToParent indicates an expected call of OnDisallowTransferToParent. +func (mr *MockOnSubAgentsMockRecorder) OnDisallowTransferToParent(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnDisallowTransferToParent", reflect.TypeOf((*MockOnSubAgents)(nil).OnDisallowTransferToParent), ctx) +} + +// OnSetAsSubAgent mocks base method. +func (m *MockOnSubAgents) OnSetAsSubAgent(ctx context.Context, parent adk.Agent) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OnSetAsSubAgent", ctx, parent) + ret0, _ := ret[0].(error) + return ret0 +} + +// OnSetAsSubAgent indicates an expected call of OnSetAsSubAgent. +func (mr *MockOnSubAgentsMockRecorder) OnSetAsSubAgent(ctx, parent any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnSetAsSubAgent", reflect.TypeOf((*MockOnSubAgents)(nil).OnSetAsSubAgent), ctx, parent) +} + +// OnSetSubAgents mocks base method. +func (m *MockOnSubAgents) OnSetSubAgents(ctx context.Context, subAgents []adk.Agent) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OnSetSubAgents", ctx, subAgents) + ret0, _ := ret[0].(error) + return ret0 +} + +// OnSetSubAgents indicates an expected call of OnSetSubAgents. +func (mr *MockOnSubAgentsMockRecorder) OnSetSubAgents(ctx, subAgents any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnSetSubAgents", reflect.TypeOf((*MockOnSubAgents)(nil).OnSetSubAgents), ctx, subAgents) +} diff --git a/internal/mock/components/document/document_mock.go b/internal/mock/components/document/document_mock.go new file mode 100644 index 0000000..e941bae --- /dev/null +++ b/internal/mock/components/document/document_mock.go @@ -0,0 +1,116 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go + +// Package document is a generated GoMock package. +package document + +import ( + context "context" + reflect "reflect" + + document "github.com/cloudwego/eino/components/document" + schema "github.com/cloudwego/eino/schema" + gomock "go.uber.org/mock/gomock" +) + +// MockLoader is a mock of Loader interface. +type MockLoader struct { + ctrl *gomock.Controller + recorder *MockLoaderMockRecorder +} + +// MockLoaderMockRecorder is the mock recorder for MockLoader. +type MockLoaderMockRecorder struct { + mock *MockLoader +} + +// NewMockLoader creates a new mock instance. +func NewMockLoader(ctrl *gomock.Controller) *MockLoader { + mock := &MockLoader{ctrl: ctrl} + mock.recorder = &MockLoaderMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockLoader) EXPECT() *MockLoaderMockRecorder { + return m.recorder +} + +// Load mocks base method. +func (m *MockLoader) Load(ctx context.Context, src document.Source, opts ...document.LoaderOption) ([]*schema.Document, error) { + m.ctrl.T.Helper() + varargs := []interface{}{ctx, src} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Load", varargs...) + ret0, _ := ret[0].([]*schema.Document) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Load indicates an expected call of Load. +func (mr *MockLoaderMockRecorder) Load(ctx, src interface{}, opts ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{ctx, src}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Load", reflect.TypeOf((*MockLoader)(nil).Load), varargs...) +} + +// MockTransformer is a mock of Transformer interface. +type MockTransformer struct { + ctrl *gomock.Controller + recorder *MockTransformerMockRecorder +} + +// MockTransformerMockRecorder is the mock recorder for MockTransformer. +type MockTransformerMockRecorder struct { + mock *MockTransformer +} + +// NewMockTransformer creates a new mock instance. +func NewMockTransformer(ctrl *gomock.Controller) *MockTransformer { + mock := &MockTransformer{ctrl: ctrl} + mock.recorder = &MockTransformerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockTransformer) EXPECT() *MockTransformerMockRecorder { + return m.recorder +} + +// Transform mocks base method. +func (m *MockTransformer) Transform(ctx context.Context, src []*schema.Document, opts ...document.TransformerOption) ([]*schema.Document, error) { + m.ctrl.T.Helper() + varargs := []interface{}{ctx, src} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Transform", varargs...) + ret0, _ := ret[0].([]*schema.Document) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Transform indicates an expected call of Transform. +func (mr *MockTransformerMockRecorder) Transform(ctx, src interface{}, opts ...interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]interface{}{ctx, src}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Transform", reflect.TypeOf((*MockTransformer)(nil).Transform), varargs...) +} diff --git a/internal/mock/components/embedding/Embedding_mock.go b/internal/mock/components/embedding/Embedding_mock.go new file mode 100644 index 0000000..f0c1063 --- /dev/null +++ b/internal/mock/components/embedding/Embedding_mock.go @@ -0,0 +1,77 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go +// +// Generated by this command: +// +// mockgen -destination ../../internal/mock/components/embedding/Embedding_mock.go --package embedding -source interface.go +// + +// Package embedding is a generated GoMock package. +package embedding + +import ( + context "context" + reflect "reflect" + + embedding "github.com/cloudwego/eino/components/embedding" + gomock "go.uber.org/mock/gomock" +) + +// MockEmbedder is a mock of Embedder interface. +type MockEmbedder struct { + ctrl *gomock.Controller + recorder *MockEmbedderMockRecorder +} + +// MockEmbedderMockRecorder is the mock recorder for MockEmbedder. +type MockEmbedderMockRecorder struct { + mock *MockEmbedder +} + +// NewMockEmbedder creates a new mock instance. +func NewMockEmbedder(ctrl *gomock.Controller) *MockEmbedder { + mock := &MockEmbedder{ctrl: ctrl} + mock.recorder = &MockEmbedderMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockEmbedder) EXPECT() *MockEmbedderMockRecorder { + return m.recorder +} + +// EmbedStrings mocks base method. +func (m *MockEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, texts} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "EmbedStrings", varargs...) + ret0, _ := ret[0].([][]float64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// EmbedStrings indicates an expected call of EmbedStrings. +func (mr *MockEmbedderMockRecorder) EmbedStrings(ctx, texts any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, texts}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EmbedStrings", reflect.TypeOf((*MockEmbedder)(nil).EmbedStrings), varargs...) +} diff --git a/internal/mock/components/indexer/indexer_mock.go b/internal/mock/components/indexer/indexer_mock.go new file mode 100644 index 0000000..829af22 --- /dev/null +++ b/internal/mock/components/indexer/indexer_mock.go @@ -0,0 +1,78 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go +// +// Generated by this command: +// +// mockgen -destination ../../internal/mock/components/indexer/indexer_mock.go --package indexer -source interface.go +// + +// Package indexer is a generated GoMock package. +package indexer + +import ( + context "context" + reflect "reflect" + + indexer "github.com/cloudwego/eino/components/indexer" + schema "github.com/cloudwego/eino/schema" + gomock "go.uber.org/mock/gomock" +) + +// MockIndexer is a mock of Indexer interface. +type MockIndexer struct { + ctrl *gomock.Controller + recorder *MockIndexerMockRecorder +} + +// MockIndexerMockRecorder is the mock recorder for MockIndexer. +type MockIndexerMockRecorder struct { + mock *MockIndexer +} + +// NewMockIndexer creates a new mock instance. +func NewMockIndexer(ctrl *gomock.Controller) *MockIndexer { + mock := &MockIndexer{ctrl: ctrl} + mock.recorder = &MockIndexerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockIndexer) EXPECT() *MockIndexerMockRecorder { + return m.recorder +} + +// Store mocks base method. +func (m *MockIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...indexer.Option) ([]string, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, docs} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Store", varargs...) + ret0, _ := ret[0].([]string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Store indicates an expected call of Store. +func (mr *MockIndexerMockRecorder) Store(ctx, docs any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, docs}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Store", reflect.TypeOf((*MockIndexer)(nil).Store), varargs...) +} diff --git a/internal/mock/components/model/ChatModel_mock.go b/internal/mock/components/model/ChatModel_mock.go new file mode 100644 index 0000000..c597d7b --- /dev/null +++ b/internal/mock/components/model/ChatModel_mock.go @@ -0,0 +1,256 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go +// +// Generated by this command: +// +// mockgen -destination ../../internal/mock/components/model/ChatModel_mock.go --package model -source interface.go +// + +// Package model is a generated GoMock package. +package model + +import ( + context "context" + reflect "reflect" + + model "github.com/cloudwego/eino/components/model" + schema "github.com/cloudwego/eino/schema" + gomock "go.uber.org/mock/gomock" +) + +// MockBaseChatModel is a mock of BaseChatModel interface. +type MockBaseChatModel struct { + ctrl *gomock.Controller + recorder *MockBaseChatModelMockRecorder + isgomock struct{} +} + +// MockBaseChatModelMockRecorder is the mock recorder for MockBaseChatModel. +type MockBaseChatModelMockRecorder struct { + mock *MockBaseChatModel +} + +// NewMockBaseChatModel creates a new mock instance. +func NewMockBaseChatModel(ctrl *gomock.Controller) *MockBaseChatModel { + mock := &MockBaseChatModel{ctrl: ctrl} + mock.recorder = &MockBaseChatModelMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockBaseChatModel) EXPECT() *MockBaseChatModelMockRecorder { + return m.recorder +} + +// Generate mocks base method. +func (m *MockBaseChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Generate", varargs...) + ret0, _ := ret[0].(*schema.Message) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Generate indicates an expected call of Generate. +func (mr *MockBaseChatModelMockRecorder) Generate(ctx, input any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Generate", reflect.TypeOf((*MockBaseChatModel)(nil).Generate), varargs...) +} + +// Stream mocks base method. +func (m *MockBaseChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Stream", varargs...) + ret0, _ := ret[0].(*schema.StreamReader[*schema.Message]) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Stream indicates an expected call of Stream. +func (mr *MockBaseChatModelMockRecorder) Stream(ctx, input any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stream", reflect.TypeOf((*MockBaseChatModel)(nil).Stream), varargs...) +} + +// MockChatModel is a mock of ChatModel interface. +type MockChatModel struct { + ctrl *gomock.Controller + recorder *MockChatModelMockRecorder + isgomock struct{} +} + +// MockChatModelMockRecorder is the mock recorder for MockChatModel. +type MockChatModelMockRecorder struct { + mock *MockChatModel +} + +// NewMockChatModel creates a new mock instance. +func NewMockChatModel(ctrl *gomock.Controller) *MockChatModel { + mock := &MockChatModel{ctrl: ctrl} + mock.recorder = &MockChatModelMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockChatModel) EXPECT() *MockChatModelMockRecorder { + return m.recorder +} + +// BindTools mocks base method. +func (m *MockChatModel) BindTools(tools []*schema.ToolInfo) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "BindTools", tools) + ret0, _ := ret[0].(error) + return ret0 +} + +// BindTools indicates an expected call of BindTools. +func (mr *MockChatModelMockRecorder) BindTools(tools any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BindTools", reflect.TypeOf((*MockChatModel)(nil).BindTools), tools) +} + +// Generate mocks base method. +func (m *MockChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Generate", varargs...) + ret0, _ := ret[0].(*schema.Message) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Generate indicates an expected call of Generate. +func (mr *MockChatModelMockRecorder) Generate(ctx, input any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Generate", reflect.TypeOf((*MockChatModel)(nil).Generate), varargs...) +} + +// Stream mocks base method. +func (m *MockChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Stream", varargs...) + ret0, _ := ret[0].(*schema.StreamReader[*schema.Message]) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Stream indicates an expected call of Stream. +func (mr *MockChatModelMockRecorder) Stream(ctx, input any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stream", reflect.TypeOf((*MockChatModel)(nil).Stream), varargs...) +} + +// MockToolCallingChatModel is a mock of ToolCallingChatModel interface. +type MockToolCallingChatModel struct { + ctrl *gomock.Controller + recorder *MockToolCallingChatModelMockRecorder + isgomock struct{} +} + +// MockToolCallingChatModelMockRecorder is the mock recorder for MockToolCallingChatModel. +type MockToolCallingChatModelMockRecorder struct { + mock *MockToolCallingChatModel +} + +// NewMockToolCallingChatModel creates a new mock instance. +func NewMockToolCallingChatModel(ctrl *gomock.Controller) *MockToolCallingChatModel { + mock := &MockToolCallingChatModel{ctrl: ctrl} + mock.recorder = &MockToolCallingChatModelMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockToolCallingChatModel) EXPECT() *MockToolCallingChatModelMockRecorder { + return m.recorder +} + +// Generate mocks base method. +func (m *MockToolCallingChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Generate", varargs...) + ret0, _ := ret[0].(*schema.Message) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Generate indicates an expected call of Generate. +func (mr *MockToolCallingChatModelMockRecorder) Generate(ctx, input any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Generate", reflect.TypeOf((*MockToolCallingChatModel)(nil).Generate), varargs...) +} + +// Stream mocks base method. +func (m *MockToolCallingChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) { + m.ctrl.T.Helper() + varargs := []any{ctx, input} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Stream", varargs...) + ret0, _ := ret[0].(*schema.StreamReader[*schema.Message]) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Stream indicates an expected call of Stream. +func (mr *MockToolCallingChatModelMockRecorder) Stream(ctx, input any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, input}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stream", reflect.TypeOf((*MockToolCallingChatModel)(nil).Stream), varargs...) +} + +// WithTools mocks base method. +func (m *MockToolCallingChatModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WithTools", tools) + ret0, _ := ret[0].(model.ToolCallingChatModel) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// WithTools indicates an expected call of WithTools. +func (mr *MockToolCallingChatModelMockRecorder) WithTools(tools any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WithTools", reflect.TypeOf((*MockToolCallingChatModel)(nil).WithTools), tools) +} diff --git a/internal/mock/components/retriever/retriever_mock.go b/internal/mock/components/retriever/retriever_mock.go new file mode 100644 index 0000000..386371a --- /dev/null +++ b/internal/mock/components/retriever/retriever_mock.go @@ -0,0 +1,78 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Code generated by MockGen. DO NOT EDIT. +// Source: interface.go +// +// Generated by this command: +// +// mockgen -destination ../../internal/mock/components/retriever/retriever_mock.go --package retriever -source interface.go +// + +// Package retriever is a generated GoMock package. +package retriever + +import ( + context "context" + reflect "reflect" + + retriever "github.com/cloudwego/eino/components/retriever" + schema "github.com/cloudwego/eino/schema" + gomock "go.uber.org/mock/gomock" +) + +// MockRetriever is a mock of Retriever interface. +type MockRetriever struct { + ctrl *gomock.Controller + recorder *MockRetrieverMockRecorder +} + +// MockRetrieverMockRecorder is the mock recorder for MockRetriever. +type MockRetrieverMockRecorder struct { + mock *MockRetriever +} + +// NewMockRetriever creates a new mock instance. +func NewMockRetriever(ctrl *gomock.Controller) *MockRetriever { + mock := &MockRetriever{ctrl: ctrl} + mock.recorder = &MockRetrieverMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRetriever) EXPECT() *MockRetrieverMockRecorder { + return m.recorder +} + +// Retrieve mocks base method. +func (m *MockRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, query} + for _, a := range opts { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "Retrieve", varargs...) + ret0, _ := ret[0].([]*schema.Document) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Retrieve indicates an expected call of Retrieve. +func (mr *MockRetrieverMockRecorder) Retrieve(ctx, query any, opts ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, query}, opts...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Retrieve", reflect.TypeOf((*MockRetriever)(nil).Retrieve), varargs...) +} diff --git a/internal/mock/doc.go b/internal/mock/doc.go new file mode 100644 index 0000000..020235d --- /dev/null +++ b/internal/mock/doc.go @@ -0,0 +1,51 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package mock provides mock implementations for testing purposes. +// +// This package aims to provide mock implementations for interfaces in the components package, +// making it easier to use in testing environments. It includes mock implementations for +// various core components such as retrievers, tools, message handlers, and graph runners. +// +// Directory Structure: +// - components/: Contains mock implementations for various components +// - retriever/: Provides mock implementation for the Retriever interface +// - retriever_mock.go: Mock implementation for document retrieval +// - tool/: Mock implementations for tool-related interfaces +// - message/: Mock implementations for message handling components +// - graph/: Mock implementations for graph execution components +// - stream/: Mock implementations for streaming components +// +// Usage: +// These mock implementations are primarily used in unit tests and integration tests, +// allowing developers to conduct tests without depending on actual external services. +// Each mock component strictly follows the contract of its corresponding interface +// while providing controllable behaviors and results. +// +// Examples: +// +// - Using mock retriever: +// retriever := mock.NewMockRetriever() +// // Configure retriever behavior +// +// - Using mock tool: +// tool := mock.NewMockTool() +// // Configure tool behavior +// +// - Using mock graph runner: +// runner := mock.NewMockGraphRunner() +// // Configure runner behavior +package mock diff --git a/internal/safe/panic.go b/internal/safe/panic.go new file mode 100644 index 0000000..16bc1dc --- /dev/null +++ b/internal/safe/panic.go @@ -0,0 +1,40 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package safe + +import ( + "fmt" +) + +type panicErr struct { + info any + stack []byte +} + +func (p *panicErr) Error() string { + return fmt.Sprintf("panic error: %v, \nstack: %s", p.info, string(p.stack)) +} + +// NewPanicErr creates a new panic error. +// panicErr is a wrapper of panic info and stack trace. +// it implements the error interface, can print error message of info and stack trace. +func NewPanicErr(info any, stack []byte) error { + return &panicErr{ + info: info, + stack: stack, + } +} diff --git a/internal/safe/panic_test.go b/internal/safe/panic_test.go new file mode 100644 index 0000000..bbe0728 --- /dev/null +++ b/internal/safe/panic_test.go @@ -0,0 +1,28 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package safe + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestPanicErr(t *testing.T) { + err := NewPanicErr("info", []byte("stack")) + assert.Equal(t, "panic error: info, \nstack: stack", err.Error()) +} diff --git a/internal/serialization/serialization.go b/internal/serialization/serialization.go new file mode 100644 index 0000000..e59ed90 --- /dev/null +++ b/internal/serialization/serialization.go @@ -0,0 +1,697 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package serialization + +import ( + "encoding/json" + "fmt" + "reflect" + + "github.com/bytedance/sonic" +) + +var m = map[string]reflect.Type{} +var rm = map[reflect.Type]string{} + +func init() { + _ = GenericRegister[int]("_eino_int") + _ = GenericRegister[int8]("_eino_int8") + _ = GenericRegister[int16]("_eino_int16") + _ = GenericRegister[int32]("_eino_int32") + _ = GenericRegister[int64]("_eino_int64") + _ = GenericRegister[uint]("_eino_uint") + _ = GenericRegister[uint8]("_eino_uint8") + _ = GenericRegister[uint16]("_eino_uint16") + _ = GenericRegister[uint32]("_eino_uint32") + _ = GenericRegister[uint64]("_eino_uint64") + _ = GenericRegister[float32]("_eino_float32") + _ = GenericRegister[float64]("_eino_float64") + _ = GenericRegister[complex64]("_eino_complex64") + _ = GenericRegister[complex128]("_eino_complex128") + _ = GenericRegister[uintptr]("_eino_uintptr") + _ = GenericRegister[bool]("_eino_bool") + _ = GenericRegister[string]("_eino_string") + _ = GenericRegister[any]("_eino_any") +} + +func GenericRegister[T any](key string) error { + t := reflect.TypeOf((*T)(nil)).Elem() + for t.Kind() == reflect.Ptr { + t = t.Elem() + } + if nt, ok := m[key]; ok { + return fmt.Errorf("key[%s] already registered to %s", key, nt.String()) + } + if nk, ok := rm[t]; ok { + return fmt.Errorf("type[%s] already registered to %s", t.String(), nk) + } + m[key] = t + rm[t] = key + return nil +} + +type InternalSerializer struct{} + +func (i *InternalSerializer) Marshal(v any) ([]byte, error) { + is, err := internalMarshal(v, nil) + if err != nil { + return nil, err + } + + return sonic.Marshal(is) +} + +func (i *InternalSerializer) Unmarshal(data []byte, v any) error { + val, err := unmarshal(data, reflect.TypeOf(v)) + if err != nil { + return fmt.Errorf("failed to unmarshal: %w", err) + } + + rv := reflect.ValueOf(v) + if rv.Kind() != reflect.Ptr || rv.IsNil() { + return fmt.Errorf("failed to unmarshal: value must be a non-nil pointer") + } + + target := rv.Elem() + if !target.CanSet() { + return fmt.Errorf("failed to unmarshal: output value must be settable") + } + + if val == nil { + target.Set(reflect.Zero(target.Type())) + return nil + } + + source := reflect.ValueOf(val) + + var set func(target, source reflect.Value) bool + set = func(target, source reflect.Value) bool { + if !source.IsValid() { + target.Set(reflect.Zero(target.Type())) + return true + } + if source.Type().AssignableTo(target.Type()) { + target.Set(source) + return true + } + + if target.Kind() == reflect.Ptr { + if target.IsNil() { + if !target.CanSet() { + return false + } + target.Set(reflect.New(target.Type().Elem())) + } + return set(target.Elem(), source) + } + + if source.Kind() == reflect.Ptr { + if source.IsNil() { + target.Set(reflect.Zero(target.Type())) + return true + } + return set(target, source.Elem()) + } + + if source.Type().ConvertibleTo(target.Type()) { + target.Set(source.Convert(target.Type())) + return true + } + + return false + } + + if set(target, source) { + return nil + } + + return fmt.Errorf("failed to unmarshal: cannot assign %s to %s", reflect.TypeOf(val), target.Type()) +} + +func unmarshal(data []byte, t reflect.Type) (any, error) { + is := &internalStruct{} + err := sonic.Unmarshal(data, is) + if err != nil { + return nil, err + } + return internalUnmarshal(is, t) +} + +type internalStruct struct { + Type *valueType `json:",omitempty"` + + JSONValue json.RawMessage `json:",omitempty"` + + // map or struct + // in map, the key is the serialized map key anyway todo: if key is string, don't serialize + // in struct, the key is the original field name + MapValues map[string]*internalStruct `json:",omitempty"` + + // slice + SliceValues []*internalStruct `json:",omitempty"` +} + +type valueType struct { + PointerNum uint32 `json:",omitempty"` + + SimpleType string `json:",omitempty"` + + StructType string `json:",omitempty"` + + MapKeyType *valueType `json:",omitempty"` + MapValueType *valueType `json:",omitempty"` + + SliceValueType *valueType `json:",omitempty"` +} + +func extractType(t reflect.Type) (*valueType, error) { + ret := &valueType{} + for t.Kind() == reflect.Ptr { + ret.PointerNum += 1 + t = t.Elem() + } + var err error + if t.Kind() == reflect.Map { + ret.MapKeyType, err = extractType(t.Key()) + if err != nil { + return nil, err + } + ret.MapValueType, err = extractType(t.Elem()) + if err != nil { + return nil, err + } + } else if t.Kind() == reflect.Slice || t.Kind() == reflect.Array { + ret.SliceValueType, err = extractType(t.Elem()) + if err != nil { + return nil, err + } + } else { + key, ok := rm[t] + if !ok { + return ret, fmt.Errorf("unknown type: %s", t.String()) + } + ret.SimpleType = key + } + return ret, nil +} + +func restoreType(vt *valueType) (reflect.Type, error) { + if vt.SimpleType != "" { + rt, ok := m[vt.SimpleType] + if !ok { + return nil, fmt.Errorf("unknown type: %s", vt.SimpleType) + } + return resolvePointerNum(vt.PointerNum, rt), nil + } + if vt.StructType != "" { + rt, ok := m[vt.StructType] + if !ok { + return nil, fmt.Errorf("unknown type: %s", vt.StructType) + } + return resolvePointerNum(vt.PointerNum, rt), nil + } + if vt.MapKeyType != nil { + rkt, err := restoreType(vt.MapKeyType) + if err != nil { + return nil, err + } + rvt, err := restoreType(vt.MapValueType) + if err != nil { + return nil, err + } + return resolvePointerNum(vt.PointerNum, reflect.MapOf(rkt, rvt)), nil + } + if vt.SliceValueType != nil { + rt, err := restoreType(vt.SliceValueType) + if err != nil { + return nil, err + } + return resolvePointerNum(vt.PointerNum, reflect.SliceOf(rt)), nil + } + return nil, fmt.Errorf("empty value") +} + +func internalMarshal(v any, fieldType reflect.Type) (*internalStruct, error) { + if v == nil || + (reflect.ValueOf(v).IsZero() && fieldType != nil && fieldType.Kind() != reflect.Interface) { + return nil, nil + } + + ret := &internalStruct{} + rv := reflect.ValueOf(v) + rt := rv.Type() + typeUnspecific := fieldType == nil || fieldType.Kind() == reflect.Interface + + var pointerNum uint32 + for rt.Kind() == reflect.Ptr { + pointerNum++ + if !rv.IsNil() { + rv = rv.Elem() + rt = rt.Elem() + continue + } + for rt.Kind() == reflect.Ptr { + rt = rt.Elem() + } + if typeUnspecific { + // need type registered + key, ok := rm[rt] + if !ok { + return nil, fmt.Errorf("unknown type: %v", rt) + } + ret.Type = &valueType{ + PointerNum: pointerNum, + SimpleType: key, + } + } + ret.JSONValue = json.RawMessage("null") + return ret, nil + } + + switch rt.Kind() { + case reflect.Struct: + if typeUnspecific { + // need type registered + key, ok := rm[rt] + if !ok { + return nil, fmt.Errorf("unknown type: %v", rt) + } + + if checkMarshaler(rt) { + ret.Type = &valueType{ + PointerNum: pointerNum, + SimpleType: key, + } + } else { + ret.Type = &valueType{ + PointerNum: pointerNum, + StructType: key, + } + } + } + + if checkMarshaler(rt) { + // Use rv.Addr() when possible so that pointer-receiver MarshalJSON methods + // are callable. rv is addressable when obtained from pointer dereference. + // When not addressable, copy into an addressable temporary. + var marshalTarget any + if rv.CanAddr() { + marshalTarget = rv.Addr().Interface() + } else { + tmp := reflect.New(rt) + tmp.Elem().Set(rv) + marshalTarget = tmp.Interface() + } + jsonBytes, err := json.Marshal(marshalTarget) + if err != nil { + return nil, err + } + ret.JSONValue = jsonBytes + return ret, nil + } + + ret.MapValues = make(map[string]*internalStruct) + + for i := 0; i < rt.NumField(); i++ { + field := rt.Field(i) + // only handle exported fields + if field.PkgPath == "" { + k := field.Name + v := rv.Field(i) + + internalValue, err := internalMarshal(v.Interface(), field.Type) + if err != nil { + return nil, err + } + + ret.MapValues[k] = internalValue + } + } + + return ret, nil + case reflect.Map: + if typeUnspecific { + var err error + ret.Type = &valueType{ + PointerNum: pointerNum, + } + // map key type + ret.Type.MapKeyType, err = extractType(rt.Key()) + if err != nil { + return nil, err + } + + // map value type + ret.Type.MapValueType, err = extractType(rt.Elem()) + if err != nil { + return nil, err + } + } + + ret.MapValues = make(map[string]*internalStruct) + + iter := rv.MapRange() + for iter.Next() { + k := iter.Key() + v := iter.Value() + + internalValue, err := internalMarshal(v.Interface(), rt.Elem()) + if err != nil { + return nil, err + } + + keyStr, err := sonic.MarshalString(k.Interface()) + if err != nil { + return nil, fmt.Errorf("marshaling map key[%v] fail: %v", k.Interface(), err) + } + ret.MapValues[keyStr] = internalValue + } + + return ret, nil + case reflect.Slice, reflect.Array: + if typeUnspecific { + var err error + ret.Type = &valueType{PointerNum: pointerNum} + ret.Type.SliceValueType, err = extractType(rt.Elem()) + if err != nil { + return nil, err + } + } + + length := rv.Len() + ret.SliceValues = make([]*internalStruct, length) + + for i := 0; i < length; i++ { + internalValue, err := internalMarshal(rv.Index(i).Interface(), rt.Elem()) + if err != nil { + return nil, err + } + ret.SliceValues[i] = internalValue + } + + return ret, nil + + default: + if typeUnspecific { + key, ok := rm[rv.Type()] + if !ok { + return nil, fmt.Errorf("unknown type: %v", rt) + } + ret.Type = &valueType{ + PointerNum: pointerNum, + SimpleType: key, + } + } + + jsonBytes, err := json.Marshal(rv.Interface()) + if err != nil { + return nil, err + } + ret.JSONValue = jsonBytes + return ret, nil + } +} + +func internalUnmarshal(v *internalStruct, typ reflect.Type) (any, error) { + if v == nil { + return nil, nil + } + + if v.Type == nil { + // specific type + if checkMarshaler(typ) { + pv := reflect.New(typ) + err := json.Unmarshal(v.JSONValue, pv.Interface()) + if err != nil { + return nil, err + } + return pv.Elem().Interface(), nil + } + return internalSpecificTypeUnmarshal(v, typ) + } + + if len(v.Type.SimpleType) != 0 { + // based type + t, ok := m[v.Type.SimpleType] + if !ok { + return nil, fmt.Errorf("unknown type key: %v", v.Type) + } + pResult := reflect.New(resolvePointerNum(v.Type.PointerNum, t)) + err := sonic.Unmarshal(v.JSONValue, pResult.Interface()) + if err != nil { + return nil, fmt.Errorf("unmarshal type[%s] fail: %v, data: %s", t.String(), err, string(v.JSONValue)) + } + return pResult.Elem().Interface(), nil + } + + if len(v.Type.StructType) > 0 { + // struct + rt, ok := m[v.Type.StructType] + if !ok { + return nil, fmt.Errorf("unknown type key: %v", v.Type.StructType) + } + result, dResult := createValueFromType(resolvePointerNum(v.Type.PointerNum, rt)) + + err := setStructFields(dResult, v.MapValues) + if err != nil { + return nil, err + } + + return result.Interface(), nil + } + + if v.Type.MapKeyType != nil { + // map + rkt, err := restoreType(v.Type.MapKeyType) + if err != nil { + return nil, err + } + rvt, err := restoreType(v.Type.MapValueType) + if err != nil { + return nil, err + } + + result, dResult := createValueFromType(reflect.MapOf(rkt, rvt)) + err = setMapKVs(dResult, v.MapValues) + if err != nil { + return nil, err + } + return result.Interface(), nil + } + + // slice + rvt, err := restoreType(v.Type.SliceValueType) + if err != nil { + return nil, err + } + + result, dResult := createValueFromType(reflect.SliceOf(rvt)) + err = setSliceElems(dResult, v.SliceValues) + if err != nil { + return nil, err + } + return result.Interface(), nil +} + +func internalSpecificTypeUnmarshal(is *internalStruct, typ reflect.Type) (any, error) { + _, dtyp := derefPointerNum(typ) + result, dResult := createValueFromType(typ) + + if dtyp.Kind() == reflect.Struct { + err := setStructFields(dResult, is.MapValues) + if err != nil { + return nil, err + } + return result.Interface(), nil + } else if dtyp.Kind() == reflect.Map { + err := setMapKVs(dResult, is.MapValues) + if err != nil { + return nil, err + } + return result.Interface(), nil + } else if dtyp.Kind() == reflect.Array || dtyp.Kind() == reflect.Slice { + err := setSliceElems(dResult, is.SliceValues) + if err != nil { + return nil, err + } + return result.Interface(), nil + } + // simple type + v := reflect.New(typ) + err := sonic.Unmarshal(is.JSONValue, v.Interface()) + if err != nil { + return nil, fmt.Errorf("unmarshal type[%s] fail: %v", typ.String(), err) + } + return v.Elem().Interface(), nil +} + +func setSliceElems(dResult reflect.Value, values []*internalStruct) error { + t := dResult.Type() + + // Handle arrays differently from slices + // Arrays have fixed size and cannot use reflect.Append + if dResult.Kind() == reflect.Array { + for i, internalValue := range values { + if i >= dResult.Len() { + return fmt.Errorf("array index out of bounds: trying to set index %d in array of length %d", i, dResult.Len()) + } + value, err := internalUnmarshal(internalValue, t.Elem()) + if err != nil { + return fmt.Errorf("unmarshal array[%s] element %d fail: %v", t.Elem(), i, err) + } + if value == nil { + dResult.Index(i).Set(reflect.Zero(t.Elem())) + } else { + dResult.Index(i).Set(reflect.ValueOf(value)) + } + } + return nil + } + + // For slices, use Append as before + for _, internalValue := range values { + value, err := internalUnmarshal(internalValue, t.Elem()) + if err != nil { + return fmt.Errorf("unmarshal slice[%s] fail: %v", t.Elem(), err) + } + if value == nil { + // empty value + dResult.Set(reflect.Append(dResult, reflect.New(t.Elem()).Elem())) + } else { + dResult.Set(reflect.Append(dResult, reflect.ValueOf(value))) + } + } + return nil +} + +func setMapKVs(dResult reflect.Value, values map[string]*internalStruct) error { + t := dResult.Type() + for marshaledMapKey, internalValue := range values { + prkv := reflect.New(t.Key()) + err := sonic.UnmarshalString(marshaledMapKey, prkv.Interface()) + if err != nil { + return fmt.Errorf("unmarshal map key[%v] to type[%s] fail: %v", marshaledMapKey, t.Key(), err) + } + + value, err := internalUnmarshal(internalValue, t.Elem()) + if err != nil { + return fmt.Errorf("unmarshal map value fail: %v", err) + } + if value == nil { + dResult.SetMapIndex(prkv.Elem(), reflect.New(t.Elem()).Elem()) + } else { + dResult.SetMapIndex(prkv.Elem(), reflect.ValueOf(value)) + } + } + return nil +} + +func setStructFields(dResult reflect.Value, values map[string]*internalStruct) error { + t := dResult.Type() + for k, internalValue := range values { + sf, ok := t.FieldByName(k) + if !ok { + continue + } + value, err := internalUnmarshal(internalValue, sf.Type) + if err != nil { + return fmt.Errorf("unmarshal map field[%v] fail: %v", k, err) + } + err = setStructField(t, dResult, k, value) + if err != nil { + return err + } + } + return nil +} + +func setStructField(t reflect.Type, s reflect.Value, fieldName string, val any) error { + field := s.FieldByName(fieldName) + if !field.CanSet() { + return fmt.Errorf("unmarshal map fail, can not set field %v", fieldName) + } + if val == nil { + rft, ok := t.FieldByName(fieldName) + if !ok { + return fmt.Errorf("unmarshal map fail, cannot find field: %v", fieldName) + } + field.Set(reflect.New(rft.Type).Elem()) + } else { + field.Set(reflect.ValueOf(val)) + } + return nil +} + +func resolvePointerNum(pointerNum uint32, t reflect.Type) reflect.Type { + for i := uint32(0); i < pointerNum; i++ { + t = reflect.PointerTo(t) + } + return t +} + +func derefPointerNum(t reflect.Type) (uint32, reflect.Type) { + var ptrCount uint32 = 0 + + for t != nil && t.Kind() == reflect.Ptr { + t = t.Elem() + ptrCount++ + } + + return ptrCount, t +} + +func createValueFromType(t reflect.Type) (value reflect.Value, derefValue reflect.Value) { + value = reflect.New(t).Elem() + + derefValue = value + for derefValue.Kind() == reflect.Ptr { + if derefValue.IsNil() { + derefValue.Set(reflect.New(derefValue.Type().Elem())) + } + derefValue = derefValue.Elem() + } + + if derefValue.Kind() == reflect.Map && derefValue.IsNil() { + derefValue.Set(reflect.MakeMap(derefValue.Type())) + } + + // Use Len() == 0 instead of IsNil() for slices to avoid panic + // IsNil() can panic on uninitialized slice values created via reflect.New().Elem() + if derefValue.Kind() == reflect.Slice { + if derefValue.Len() == 0 && derefValue.Cap() == 0 { + derefValue.Set(reflect.MakeSlice(derefValue.Type(), 0, 0)) + } + } + // Arrays cannot be nil and don't need initialization + + return value, derefValue +} + +var marshalerType = reflect.TypeOf((*json.Marshaler)(nil)).Elem() +var unmarshalerType = reflect.TypeOf((*json.Unmarshaler)(nil)).Elem() + +func checkMarshaler(t reflect.Type) bool { + for t.Kind() == reflect.Ptr { + t = t.Elem() + } + + if (t.Implements(marshalerType) || reflect.PointerTo(t).Implements(marshalerType)) && + (t.Implements(unmarshalerType) || reflect.PointerTo(t).Implements(unmarshalerType)) { + return true + } + return false +} diff --git a/internal/serialization/serialization_test.go b/internal/serialization/serialization_test.go new file mode 100644 index 0000000..014c7fd --- /dev/null +++ b/internal/serialization/serialization_test.go @@ -0,0 +1,351 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package serialization + +import ( + "reflect" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type myInterface interface { + Method() +} +type myStruct struct { + A string +} + +func (m *myStruct) Method() {} + +type myStruct2 struct { + A any + B myInterface + C map[string]**myStruct + D map[myStruct]any + E []any + f string + G myStruct3 + H *myStruct4 + I []*myStruct3 + J map[string]myStruct3 + K myStruct4 + L []*myStruct4 + M map[string]myStruct4 +} + +type myStruct3 struct { + FieldA string +} + +type myStruct4 struct { + FieldA string +} + +func (m *myStruct4) UnmarshalJSON(bytes []byte) error { + m.FieldA = string(bytes) + return nil +} + +func (m myStruct4) MarshalJSON() ([]byte, error) { + return []byte(m.FieldA), nil +} + +func TestSerialization(t *testing.T) { + _ = GenericRegister[myStruct]("myStruct") + _ = GenericRegister[myStruct2]("myStruct2") + _ = GenericRegister[myInterface]("myInterface") + ms := myStruct{A: "test"} + pms := &ms + pointerOfPointerOfMyStruct := &pms + + ms1 := myStruct{A: "1"} + ms2 := myStruct{A: "2"} + ms3 := myStruct{A: "3"} + ms4 := myStruct{A: "4"} + values := []any{ + 10, + "test", + ms, + pms, + pointerOfPointerOfMyStruct, + myInterface(pms), + []int{1, 2, 3}, + []any{1, "test"}, + []myInterface{nil, &myStruct{A: "1"}, &myStruct{A: "2"}}, + map[string]string{"123": "123", "abc": "abc"}, + map[string]myInterface{"1": nil, "2": pms}, + map[string]any{"123": 1, "abc": &myStruct{A: "1"}, "bcd": nil}, + map[myStruct]any{ + ms1: 1, + ms2: &myStruct{ + A: "2", + }, + ms3: nil, + ms4: []any{ + 1, + pointerOfPointerOfMyStruct, + "123", &myStruct{ + A: "1", + }, + nil, + map[myStruct]any{ + ms1: 1, + ms2: nil, + }, + }, + }, + myStruct2{ + A: "123", + B: &myStruct{ + A: "test", + }, + C: map[string]**myStruct{ + "a": pointerOfPointerOfMyStruct, + }, + D: map[myStruct]any{{"a"}: 1}, + E: []any{1, "2", 3}, + f: "", + G: myStruct3{ + FieldA: "1", + }, + H: nil, + I: []*myStruct3{ + {FieldA: "2"}, {FieldA: "3"}, + }, + J: map[string]myStruct3{ + "1": {FieldA: "4"}, + "2": {FieldA: "5"}, + }, + K: myStruct4{ + FieldA: "1", + }, + L: []*myStruct4{ + {FieldA: "2"}, {FieldA: "3"}, + }, + M: map[string]myStruct4{ + "1": {FieldA: "4"}, + "2": {FieldA: "5"}, + }, + }, + map[string]map[string][]map[string][][]string{ + "1": { + "a": []map[string][][]string{ + {"b": { + {"c"}, + {"d"}, + }}, + }, + }, + }, + []*myStruct{}, + &myStruct{}, + } + + for _, value := range values { + data, err := (&InternalSerializer{}).Marshal(value) + assert.NoError(t, err) + v := reflect.New(reflect.TypeOf(value)).Interface() + err = (&InternalSerializer{}).Unmarshal(data, v) + assert.NoError(t, err) + assert.Equal(t, value, reflect.ValueOf(v).Elem().Interface()) + } +} + +type myStruct5 struct { + FieldA string +} + +func (m *myStruct5) UnmarshalJSON(bytes []byte) error { + m.FieldA = "FieldA" + return nil +} + +func (m myStruct5) MarshalJSON() ([]byte, error) { + return []byte("1"), nil +} + +func TestMarshalStruct(t *testing.T) { + assert.NoError(t, GenericRegister[myStruct5]("myStruct5")) + s := myStruct5{FieldA: "1"} + data, err := (&InternalSerializer{}).Marshal(s) + assert.NoError(t, err) + result := &myStruct5{} + err = (&InternalSerializer{}).Unmarshal(data, result) + assert.NoError(t, err) + assert.Equal(t, myStruct5{FieldA: "FieldA"}, *result) + + ma := map[string]any{ + "1": s, + } + data, err = (&InternalSerializer{}).Marshal(ma) + assert.NoError(t, err) + result2 := map[string]any{} + err = (&InternalSerializer{}).Unmarshal(data, &result2) + assert.NoError(t, err) + assert.Equal(t, map[string]any{ + "1": myStruct5{FieldA: "FieldA"}, + }, result2) +} + +type unmarshalTestStruct struct { + Foo string + Bar int +} + +func init() { + // Register types for the serializer to work. + // This is necessary for the serializer to know how to handle custom struct types. + err := GenericRegister[unmarshalTestStruct]("unmarshalTestStruct") + if err != nil { + panic(err) + } +} + +func TestInternalSerializer_Unmarshal(t *testing.T) { + s := InternalSerializer{} + + t.Run("success cases", func(t *testing.T) { + // Helper to create a pointer to a value, needed for the expected value in one test case. + ptr := func(i int) *int { return &i } + + testCases := []struct { + name string + inputValue any + outputPtr any + expectedVal any + }{ + { + name: "simple type", + inputValue: 123, + outputPtr: new(int), + expectedVal: 123, + }, + { + name: "struct type", + inputValue: unmarshalTestStruct{Foo: "hello", Bar: 42}, + outputPtr: new(unmarshalTestStruct), + expectedVal: unmarshalTestStruct{Foo: "hello", Bar: 42}, + }, + { + name: "pointer to struct", + inputValue: &unmarshalTestStruct{Foo: "world", Bar: 99}, + outputPtr: new(*unmarshalTestStruct), + expectedVal: &unmarshalTestStruct{Foo: "world", Bar: 99}, + }, + { + name: "unmarshal pointer to value", + inputValue: &unmarshalTestStruct{Foo: "p2v", Bar: 1}, + outputPtr: new(unmarshalTestStruct), + expectedVal: unmarshalTestStruct{Foo: "p2v", Bar: 1}, + }, + { + name: "unmarshal value to pointer", + inputValue: unmarshalTestStruct{Foo: "v2p", Bar: 2}, + outputPtr: new(*unmarshalTestStruct), + expectedVal: &unmarshalTestStruct{Foo: "v2p", Bar: 2}, + }, + { + name: "unmarshal nil pointer", + inputValue: (*unmarshalTestStruct)(nil), + outputPtr: &struct{ v *unmarshalTestStruct }{v: &unmarshalTestStruct{}}, // placeholder to be replaced + expectedVal: (*unmarshalTestStruct)(nil), + }, + { + name: "convertible types", + inputValue: int32(42), + outputPtr: new(int64), + expectedVal: int64(42), + }, + { + name: "pointer to pointer destination", + inputValue: 12345, + outputPtr: new(*int), + expectedVal: ptr(12345), + }, + { + name: "unmarshal to any", + inputValue: unmarshalTestStruct{Foo: "any", Bar: 101}, + outputPtr: new(any), + expectedVal: unmarshalTestStruct{Foo: "any", Bar: 101}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + data, err := s.Marshal(tc.inputValue) + require.NoError(t, err) + + // Special handling for the nil test case to correctly pass the pointer. + if tc.name == "unmarshal nil pointer" { + target := tc.outputPtr.(*struct{ v *unmarshalTestStruct }) + err = s.Unmarshal(data, &target.v) + require.NoError(t, err) + assert.Nil(t, target.v) + return + } + + err = s.Unmarshal(data, tc.outputPtr) + require.NoError(t, err) + + // Dereference the pointer to get the actual value for comparison. + actualVal := reflect.ValueOf(tc.outputPtr).Elem().Interface() + assert.Equal(t, tc.expectedVal, actualVal) + }) + } + }) + + t.Run("error cases", func(t *testing.T) { + data, err := s.Marshal(123) + require.NoError(t, err) + + t.Run("destination not a pointer", func(t *testing.T) { + var output int + err := s.Unmarshal(data, output) + require.Error(t, err) + assert.Contains(t, err.Error(), "value must be a non-nil pointer") + }) + + t.Run("destination is a nil pointer", func(t *testing.T) { + var output *int // nil + err := s.Unmarshal(data, output) + require.Error(t, err) + assert.Contains(t, err.Error(), "value must be a non-nil pointer") + }) + + t.Run("type mismatch", func(t *testing.T) { + strData, mErr := s.Marshal("i am a string") + require.NoError(t, mErr) + + var output int + err := s.Unmarshal(strData, &output) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot assign") + }) + + t.Run("unconvertible types", func(t *testing.T) { + intData, mErr := s.Marshal(123) + require.NoError(t, mErr) + + var output bool + err := s.Unmarshal(intData, &output) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot assign") + }) + }) +} diff --git a/schema/agentic_message.go b/schema/agentic_message.go new file mode 100644 index 0000000..2474e5d --- /dev/null +++ b/schema/agentic_message.go @@ -0,0 +1,2284 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "bytes" + "context" + "encoding/gob" + "encoding/json" + "fmt" + "reflect" + "sort" + "strings" + + "github.com/bytedance/sonic" + "github.com/eino-contrib/jsonschema" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/schema/claude" + "github.com/cloudwego/eino/schema/gemini" + "github.com/cloudwego/eino/schema/openai" +) + +type ContentBlockType string + +const ( + ContentBlockTypeReasoning ContentBlockType = "reasoning" + ContentBlockTypeUserInputText ContentBlockType = "user_input_text" + ContentBlockTypeUserInputImage ContentBlockType = "user_input_image" + ContentBlockTypeUserInputAudio ContentBlockType = "user_input_audio" + ContentBlockTypeUserInputVideo ContentBlockType = "user_input_video" + ContentBlockTypeUserInputFile ContentBlockType = "user_input_file" + ContentBlockTypeToolSearchResult ContentBlockType = "tool_search_result" + ContentBlockTypeAssistantGenText ContentBlockType = "assistant_gen_text" + ContentBlockTypeAssistantGenImage ContentBlockType = "assistant_gen_image" + ContentBlockTypeAssistantGenAudio ContentBlockType = "assistant_gen_audio" + ContentBlockTypeAssistantGenVideo ContentBlockType = "assistant_gen_video" + ContentBlockTypeFunctionToolCall ContentBlockType = "function_tool_call" + ContentBlockTypeFunctionToolResult ContentBlockType = "function_tool_result" + ContentBlockTypeServerToolCall ContentBlockType = "server_tool_call" + ContentBlockTypeServerToolResult ContentBlockType = "server_tool_result" + ContentBlockTypeMCPToolCall ContentBlockType = "mcp_tool_call" + ContentBlockTypeMCPToolResult ContentBlockType = "mcp_tool_result" + ContentBlockTypeMCPListToolsResult ContentBlockType = "mcp_list_tools_result" + ContentBlockTypeMCPToolApprovalRequest ContentBlockType = "mcp_tool_approval_request" + ContentBlockTypeMCPToolApprovalResponse ContentBlockType = "mcp_tool_approval_response" +) + +type AgenticRoleType string + +const ( + AgenticRoleTypeSystem AgenticRoleType = "system" + AgenticRoleTypeUser AgenticRoleType = "user" + AgenticRoleTypeAssistant AgenticRoleType = "assistant" +) + +type AgenticMessage struct { + // Role is the message role. + Role AgenticRoleType `json:"role"` + + // ContentBlocks is the list of content blocks. + ContentBlocks []*ContentBlock `json:"content_blocks,omitempty"` + + // ResponseMeta is the response metadata. + ResponseMeta *AgenticResponseMeta `json:"response_meta,omitempty"` + + // Extra is the additional information. + Extra map[string]any `json:"extra,omitempty"` +} + +type AgenticResponseMeta struct { + // TokenUsage is the token usage. + TokenUsage *TokenUsage `json:"token_usage,omitempty"` + + // OpenAIExtension is the extension for OpenAI. + OpenAIExtension *openai.ResponseMetaExtension `json:"openai_extension,omitempty"` + + // GeminiExtension is the extension for Gemini. + GeminiExtension *gemini.ResponseMetaExtension `json:"gemini_extension,omitempty"` + + // ClaudeExtension is the extension for Claude. + ClaudeExtension *claude.ResponseMetaExtension `json:"claude_extension,omitempty"` + + // Extension is the extension for other models, supplied by the component implementer. + Extension any `json:"extension,omitempty"` +} + +type ContentBlock struct { + Type ContentBlockType `json:"type"` + + // Reasoning contains the reasoning content generated by the model. + Reasoning *Reasoning `json:"reasoning,omitempty"` + + // UserInputText contains the text content provided by the user. + UserInputText *UserInputText `json:"user_input_text,omitempty"` + + // UserInputImage contains the image content provided by the user. + UserInputImage *UserInputImage `json:"user_input_image,omitempty"` + + // UserInputAudio contains the audio content provided by the user. + UserInputAudio *UserInputAudio `json:"user_input_audio,omitempty"` + + // UserInputVideo contains the video content provided by the user. + UserInputVideo *UserInputVideo `json:"user_input_video,omitempty"` + + // UserInputFile contains the file content provided by the user. + UserInputFile *UserInputFile `json:"user_input_file,omitempty"` + + // AssistantGenText contains the text content generated by the model. + AssistantGenText *AssistantGenText `json:"assistant_gen_text,omitempty"` + + // AssistantGenImage contains the image content generated by the model. + AssistantGenImage *AssistantGenImage `json:"assistant_gen_image,omitempty"` + + // AssistantGenAudio contains the audio content generated by the model. + AssistantGenAudio *AssistantGenAudio `json:"assistant_gen_audio,omitempty"` + + // AssistantGenVideo contains the video content generated by the model. + AssistantGenVideo *AssistantGenVideo `json:"assistant_gen_video,omitempty"` + + // FunctionToolCall contains the invocation details for a user-defined tool. + FunctionToolCall *FunctionToolCall `json:"function_tool_call,omitempty"` + + // FunctionToolResult contains the result returned from a user-defined tool call. + FunctionToolResult *FunctionToolResult `json:"function_tool_result,omitempty"` + + // ToolSearchFunctionToolResult contains the result of a client-side custom tool search tool call. + // It carries the full definitions of newly discovered tools so that the model can + // recognize which tools have been added and are now available for invocation. + ToolSearchFunctionToolResult *ToolSearchFunctionToolResult `json:"tool_search_function_tool_result,omitempty"` + + // ServerToolCall contains the invocation details for a provider built-in tool executed on the model server. + ServerToolCall *ServerToolCall `json:"server_tool_call,omitempty"` + + // ServerToolResult contains the result returned from a provider built-in tool executed on the model server. + ServerToolResult *ServerToolResult `json:"server_tool_result,omitempty"` + + // MCPToolCall contains the invocation details for an MCP tool managed by the model server. + MCPToolCall *MCPToolCall `json:"mcp_tool_call,omitempty"` + + // MCPToolResult contains the result returned from an MCP tool managed by the model server. + MCPToolResult *MCPToolResult `json:"mcp_tool_result,omitempty"` + + // MCPListToolsResult contains the list of available MCP tools reported by the model server. + MCPListToolsResult *MCPListToolsResult `json:"mcp_list_tools_result,omitempty"` + + // MCPToolApprovalRequest contains the user approval request for an MCP tool call when required. + MCPToolApprovalRequest *MCPToolApprovalRequest `json:"mcp_tool_approval_request,omitempty"` + + // MCPToolApprovalResponse contains the user's approval decision for an MCP tool call. + MCPToolApprovalResponse *MCPToolApprovalResponse `json:"mcp_tool_approval_response,omitempty"` + + // StreamingMeta contains metadata for streaming responses. + // Only set for streaming responses. + StreamingMeta *StreamingMeta `json:"streaming_meta,omitempty"` + + // Extra contains additional information for the content block. + Extra map[string]any `json:"extra,omitempty"` +} + +type StreamingMeta struct { + // Index specifies the index position of this block in the final response. + Index int `json:"index"` +} + +type UserInputText struct { + // Text is the text content. + Text string `json:"text,omitempty"` +} + +type UserInputImage struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "image/png". + MIMEType string `json:"mime_type,omitempty"` + + // Detail is the quality of the image url. + Detail ImageURLDetail `json:"detail,omitempty"` +} + +type UserInputAudio struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "audio/wav". + MIMEType string `json:"mime_type,omitempty"` +} + +type UserInputVideo struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "video/mp4". + MIMEType string `json:"mime_type,omitempty"` +} + +type UserInputFile struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Name is the filename. + Name string `json:"name,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "application/pdf". + MIMEType string `json:"mime_type,omitempty"` +} + +type AssistantGenText struct { + // Text is the generated text. + Text string `json:"text,omitempty"` + + // OpenAIExtension is the extension for OpenAI. + OpenAIExtension *openai.AssistantGenTextExtension `json:"openai_extension,omitempty"` + + // ClaudeExtension is the extension for Claude. + ClaudeExtension *claude.AssistantGenTextExtension `json:"claude_extension,omitempty"` + + // Extension is the extension for other models, supplied by the component implementer. + Extension any `json:"extension,omitempty"` +} + +type AssistantGenImage struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "image/png". + MIMEType string `json:"mime_type,omitempty"` +} + +type AssistantGenAudio struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "audio/wav". + MIMEType string `json:"mime_type,omitempty"` +} + +type AssistantGenVideo struct { + // URL is the HTTP/HTTPS link. + URL string `json:"url,omitempty"` + + // Base64Data is the binary data in Base64 encoded string format. + Base64Data string `json:"base64_data,omitempty"` + + // MIMEType is the mime type, e.g. "video/mp4". + MIMEType string `json:"mime_type,omitempty"` +} + +type Reasoning struct { + // Text is either the thought summary or the raw reasoning text itself. + Text string `json:"text,omitempty"` + + // Signature contains encrypted reasoning tokens. + // Required by some models when passing reasoning text back. + Signature string `json:"signature,omitempty"` + + // OpenAIExtension is the extension for OpenAI. + OpenAIExtension *openai.ReasoningExtension `json:"openai_extension,omitempty"` +} + +type FunctionToolCall struct { + // CallID is the unique identifier for the tool call. + CallID string `json:"call_id,omitempty"` + + // Name specifies the function tool invoked. + Name string `json:"name"` + + // Arguments is the JSON string arguments for the function tool call. + Arguments string `json:"arguments,omitempty"` +} + +// FunctionToolResultContentBlockType identifies which media field of a +// FunctionToolResultContentBlock is populated. +type FunctionToolResultContentBlockType string + +const ( + FunctionToolResultContentBlockTypeText FunctionToolResultContentBlockType = "text" + FunctionToolResultContentBlockTypeImage FunctionToolResultContentBlockType = "image" + FunctionToolResultContentBlockTypeAudio FunctionToolResultContentBlockType = "audio" + FunctionToolResultContentBlockTypeVideo FunctionToolResultContentBlockType = "video" + FunctionToolResultContentBlockTypeFile FunctionToolResultContentBlockType = "file" +) + +// FunctionToolResultContentBlock represents a single content block within a multimodal +// function tool result. Type identifies which of the media fields is populated; +// exactly one of the media fields should be set to match Type. +type FunctionToolResultContentBlock struct { + // Type identifies which media field below is populated. + Type FunctionToolResultContentBlockType `json:"type"` + // Text contains the text content of the block. + Text *UserInputText `json:"text,omitempty"` + // Image contains the image content of the block. + Image *UserInputImage `json:"image,omitempty"` + // Audio contains the audio content of the block. + Audio *UserInputAudio `json:"audio,omitempty"` + // Video contains the video content of the block. + Video *UserInputVideo `json:"video,omitempty"` + // File contains the file content of the block. + File *UserInputFile `json:"file,omitempty"` + + // Extra holds additional metadata for model-specific or custom extensions. + Extra map[string]any `json:"extra,omitempty"` +} + +func (b *FunctionToolResultContentBlock) String() string { + switch b.Type { + case FunctionToolResultContentBlockTypeText: + if b.Text != nil { + return b.Text.String() + } + return "empty text block\n" + case FunctionToolResultContentBlockTypeImage: + if b.Image != nil { + return b.Image.String() + } + return "empty image block\n" + case FunctionToolResultContentBlockTypeAudio: + if b.Audio != nil { + return b.Audio.String() + } + return "empty audio block\n" + case FunctionToolResultContentBlockTypeVideo: + if b.Video != nil { + return b.Video.String() + } + return "empty video block\n" + case FunctionToolResultContentBlockTypeFile: + if b.File != nil { + return b.File.String() + } + return "empty file block\n" + case "": + return "unknown block type: \n" + default: + return fmt.Sprintf("unknown block type: %s\n", b.Type) + } +} + +type FunctionToolResult struct { + // CallID is the unique identifier for the tool call. + CallID string `json:"call_id,omitempty"` + + // Name specifies the function tool invoked. + Name string `json:"name"` + + // Content holds the tool execution output as an ordered list of content blocks. + // Each block carries its own type (text, image, audio, video, file), allowing + // text-only and multimodal results to share a uniform representation. + Content []*FunctionToolResultContentBlock `json:"content,omitempty"` +} + +// ToolSearchFunctionToolResult represents the result of a client-side custom tool search +// function tool call. Unlike a regular FunctionToolResult, this carries a ToolSearchResult +// containing the full definitions of newly discovered tools, so the model can recognize +// which tools have been added and are now available for invocation. +type ToolSearchFunctionToolResult struct { + // CallID is the unique identifier for the tool call. + CallID string `json:"call_id,omitempty"` + + // Name specifies the function tool invoked. + Name string `json:"name"` + + // Result is the function tool result returned by the user + Result *ToolSearchResult `json:"result,omitempty"` +} + +func (t *ToolSearchFunctionToolResult) String() string { + if t.Result != nil { + return t.Result.String() + } + return "" +} + +type ServerToolCall struct { + // Name specifies the server-side tool invoked. + // Supplied by the model server (e.g., `web_search` for OpenAI, `googleSearch` for Gemini). + Name string `json:"name"` + + // CallID is the unique identifier for the tool call. + // Empty if not provided by the model server. + CallID string `json:"call_id,omitempty"` + + // Arguments are the raw inputs to the server-side tool, + // supplied by the component implementer. + Arguments any `json:"arguments,omitempty"` +} + +type ServerToolResult struct { + // Name specifies the server-side tool invoked. + // Supplied by the model server (e.g., `web_search` for OpenAI, `googleSearch` for Gemini). + Name string `json:"name"` + + // CallID is the unique identifier for the tool call. + // Empty if not provided by the model server. + CallID string `json:"call_id,omitempty"` + + // Content refers to the raw output generated by the server-side tool, + // supplied by the component implementer. + Content any `json:"content,omitempty"` +} + +type MCPToolCall struct { + // ServerLabel is the MCP server label used to identify it in tool calls + ServerLabel string `json:"server_label,omitempty"` + + // ApprovalRequestID is the approval request ID. + ApprovalRequestID string `json:"approval_request_id,omitempty"` + + // CallID is the unique ID of the tool call. + CallID string `json:"call_id,omitempty"` + + // Name is the name of the tool to run. + Name string `json:"name"` + + // Arguments is the JSON string arguments for the tool call. + Arguments string `json:"arguments,omitempty"` +} + +type MCPToolResult struct { + // ServerLabel is the MCP server label used to identify it in tool calls + ServerLabel string `json:"server_label,omitempty"` + + // CallID is the unique ID of the tool call. + CallID string `json:"call_id,omitempty"` + + // Name is the name of the tool to run. + Name string `json:"name"` + + // Content is the JSON string with the tool result. + Content string `json:"content,omitempty"` + + // Error returned when the server fails to run the tool. + Error *MCPToolCallError `json:"error,omitempty"` +} + +type MCPToolCallError struct { + // Code is the error code. + Code *int64 `json:"code,omitempty"` + + // Message is the error message. + Message string `json:"message,omitempty"` +} + +type MCPListToolsResult struct { + // ServerLabel is the MCP server label used to identify it in tool calls. + ServerLabel string `json:"server_label,omitempty"` + + // Tools is the list of tools available on the server. + Tools []*MCPListToolsItem `json:"tools,omitempty"` + + // Error returned when the server fails to list tools. + Error string `json:"error,omitempty"` +} + +type MCPListToolsItem struct { + // Name is the name of the tool. + Name string `json:"name"` + + // Description is the description of the tool. + Description string `json:"description"` + + // InputSchema is the JSON schema that describes the tool input parameters. + InputSchema *jsonschema.Schema `json:"input_schema,omitempty"` +} + +type mcpListToolsItemGob struct { + Name string + Description string + InputSchemaJSON []byte +} + +func (m *MCPListToolsItem) GobEncode() ([]byte, error) { + g := mcpListToolsItemGob{ + Name: m.Name, + Description: m.Description, + } + if m.InputSchema != nil { + b, err := json.Marshal(m.InputSchema) + if err != nil { + return nil, fmt.Errorf("failed to marshal MCPListToolsItem.InputSchema: %w", err) + } + g.InputSchemaJSON = b + } + var buf bytes.Buffer + if err := gob.NewEncoder(&buf).Encode(&g); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func (m *MCPListToolsItem) GobDecode(data []byte) error { + var g mcpListToolsItemGob + if err := gob.NewDecoder(bytes.NewReader(data)).Decode(&g); err != nil { + return err + } + m.Name = g.Name + m.Description = g.Description + if len(g.InputSchemaJSON) > 0 { + m.InputSchema = &jsonschema.Schema{} + if err := sonic.Unmarshal(g.InputSchemaJSON, m.InputSchema); err != nil { + return fmt.Errorf("failed to unmarshal MCPListToolsItem.InputSchema: %w", err) + } + } + return nil +} + +type MCPToolApprovalRequest struct { + // ID is the approval request ID. + ID string `json:"id,omitempty"` + + // Name is the name of the tool to run. + Name string `json:"name"` + + // Arguments is the JSON string arguments for the tool call. + Arguments string `json:"arguments,omitempty"` + + // ServerLabel is the MCP server label used to identify it in tool calls. + ServerLabel string `json:"server_label,omitempty"` +} + +type MCPToolApprovalResponse struct { + // ApprovalRequestID is the approval request ID being responded to. + ApprovalRequestID string `json:"approval_request_id,omitempty"` + + // Approve indicates whether the request is approved. + Approve bool `json:"approve"` + + // Reason is the rationale for the decision. + // Optional. + Reason string `json:"reason,omitempty"` +} + +// SystemAgenticMessage represents a message with AgenticRoleType "system". +func SystemAgenticMessage(text string) *AgenticMessage { + return &AgenticMessage{ + Role: AgenticRoleTypeSystem, + ContentBlocks: []*ContentBlock{NewContentBlock(&UserInputText{Text: text})}, + } +} + +// UserAgenticMessage represents a message with AgenticRoleType "user". +func UserAgenticMessage(text string) *AgenticMessage { + return &AgenticMessage{ + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{NewContentBlock(&UserInputText{Text: text})}, + } +} + +type contentBlockVariant interface { + Reasoning | userInputVariant | assistantGenVariant | functionToolCallVariant | serverToolCallVariant | mcpToolCallVariant +} + +type userInputVariant interface { + UserInputText | UserInputImage | UserInputAudio | UserInputVideo | UserInputFile +} + +type assistantGenVariant interface { + AssistantGenText | AssistantGenImage | AssistantGenAudio | AssistantGenVideo +} + +type functionToolCallVariant interface { + FunctionToolCall | FunctionToolResult | ToolSearchFunctionToolResult +} + +type serverToolCallVariant interface { + ServerToolCall | ServerToolResult +} + +type mcpToolCallVariant interface { + MCPToolCall | MCPToolResult | MCPListToolsResult | MCPToolApprovalRequest | MCPToolApprovalResponse +} + +// NewContentBlock creates a new ContentBlock with the given content. +func NewContentBlock[T contentBlockVariant](content *T) *ContentBlock { + switch b := any(content).(type) { + case *Reasoning: + return &ContentBlock{Type: ContentBlockTypeReasoning, Reasoning: b} + case *UserInputText: + return &ContentBlock{Type: ContentBlockTypeUserInputText, UserInputText: b} + case *UserInputImage: + return &ContentBlock{Type: ContentBlockTypeUserInputImage, UserInputImage: b} + case *UserInputAudio: + return &ContentBlock{Type: ContentBlockTypeUserInputAudio, UserInputAudio: b} + case *UserInputVideo: + return &ContentBlock{Type: ContentBlockTypeUserInputVideo, UserInputVideo: b} + case *UserInputFile: + return &ContentBlock{Type: ContentBlockTypeUserInputFile, UserInputFile: b} + case *ToolSearchFunctionToolResult: + return &ContentBlock{Type: ContentBlockTypeToolSearchResult, ToolSearchFunctionToolResult: b} + case *AssistantGenText: + return &ContentBlock{Type: ContentBlockTypeAssistantGenText, AssistantGenText: b} + case *AssistantGenImage: + return &ContentBlock{Type: ContentBlockTypeAssistantGenImage, AssistantGenImage: b} + case *AssistantGenAudio: + return &ContentBlock{Type: ContentBlockTypeAssistantGenAudio, AssistantGenAudio: b} + case *AssistantGenVideo: + return &ContentBlock{Type: ContentBlockTypeAssistantGenVideo, AssistantGenVideo: b} + case *FunctionToolCall: + return &ContentBlock{Type: ContentBlockTypeFunctionToolCall, FunctionToolCall: b} + case *FunctionToolResult: + return &ContentBlock{Type: ContentBlockTypeFunctionToolResult, FunctionToolResult: b} + case *ServerToolCall: + return &ContentBlock{Type: ContentBlockTypeServerToolCall, ServerToolCall: b} + case *ServerToolResult: + return &ContentBlock{Type: ContentBlockTypeServerToolResult, ServerToolResult: b} + case *MCPToolCall: + return &ContentBlock{Type: ContentBlockTypeMCPToolCall, MCPToolCall: b} + case *MCPToolResult: + return &ContentBlock{Type: ContentBlockTypeMCPToolResult, MCPToolResult: b} + case *MCPListToolsResult: + return &ContentBlock{Type: ContentBlockTypeMCPListToolsResult, MCPListToolsResult: b} + case *MCPToolApprovalRequest: + return &ContentBlock{Type: ContentBlockTypeMCPToolApprovalRequest, MCPToolApprovalRequest: b} + case *MCPToolApprovalResponse: + return &ContentBlock{Type: ContentBlockTypeMCPToolApprovalResponse, MCPToolApprovalResponse: b} + default: + return nil + } +} + +// NewContentBlockChunk creates a new ContentBlock with the given content and streaming metadata. +func NewContentBlockChunk[T contentBlockVariant](content *T, meta *StreamingMeta) *ContentBlock { + block := NewContentBlock(content) + block.StreamingMeta = meta + return block +} + +// AgenticMessagesTemplate is the interface for agentic messages template. +// It's used to render a template to a list of agentic messages. +// e.g. +// +// chatTemplate := prompt.FromAgenticMessages( +// &schema.AgenticMessage{ +// Role: schema.AgenticRoleTypeSystem, +// ContentBlocks: []*schema.ContentBlock{ +// {Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "you are an eino helper"}}, +// }, +// }, +// schema.AgenticMessagesPlaceholder("history", false), // <= this will use the value of "history" in params +// ) +// msgs, err := chatTemplate.Format(ctx, params) +type AgenticMessagesTemplate interface { + Format(ctx context.Context, vs map[string]any, formatType FormatType) ([]*AgenticMessage, error) +} + +var _ AgenticMessagesTemplate = &AgenticMessage{} +var _ AgenticMessagesTemplate = AgenticMessagesPlaceholder("", false) + +type agenticMessagesPlaceholder struct { + key string + optional bool +} + +// AgenticMessagesPlaceholder can render a placeholder to a list of agentic messages in params. +// e.g. +// +// placeholder := AgenticMessagesPlaceholder("history", false) +// params := map[string]any{ +// "history": []*schema.AgenticMessage{ +// &schema.AgenticMessage{ +// Role: schema.AgenticRoleTypeSystem, +// ContentBlocks: []*schema.ContentBlock{ +// {Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "you are an eino helper"}}, +// }, +// }, +// }, +// } +// chatTemplate := chatTpl := prompt.FromMessages( +// schema.AgenticMessagesPlaceholder("history", false), // <= this will use the value of "history" in params +// ) +// msgs, err := chatTemplate.Format(ctx, params) +func AgenticMessagesPlaceholder(key string, optional bool) AgenticMessagesTemplate { + return &agenticMessagesPlaceholder{ + key: key, + optional: optional, + } +} + +func (p *agenticMessagesPlaceholder) Format(_ context.Context, vs map[string]any, _ FormatType) ([]*AgenticMessage, error) { + v, ok := vs[p.key] + if !ok { + if p.optional { + return []*AgenticMessage{}, nil + } + + return nil, fmt.Errorf("message placeholder format: %s not found", p.key) + } + + msgs, ok := v.([]*AgenticMessage) + if !ok { + return nil, fmt.Errorf("only agentic messages can be used to format message placeholder, key: %v, actual type: %v", p.key, reflect.TypeOf(v)) + } + + return msgs, nil +} + +// Format returns the agentic messages after rendering by the given formatType. +// It formats only the user input fields (UserInputText, UserInputImage, UserInputAudio, UserInputVideo, UserInputFile). +// e.g. +// +// msg := &schema.AgenticMessage{ +// Role: schema.AgenticRoleTypeUser, +// ContentBlocks: []*schema.ContentBlock{ +// {Type: schema.ContentBlockTypeUserInputText, UserInputText: &schema.UserInputText{Text: "hello {name}"}}, +// }, +// } +// msgs, err := msg.Format(ctx, map[string]any{"name": "eino"}, schema.FString) +// // msgs[0].ContentBlocks[0].UserInputText.Text will be "hello eino" +func (m *AgenticMessage) Format(_ context.Context, vs map[string]any, formatType FormatType) ([]*AgenticMessage, error) { + copied := *m + + if len(m.ContentBlocks) > 0 { + copiedBlocks := make([]*ContentBlock, len(m.ContentBlocks)) + for i, block := range m.ContentBlocks { + if block == nil { + copiedBlocks[i] = nil + continue + } + + copiedBlock := *block + var err error + + switch block.Type { + case ContentBlockTypeUserInputText: + if block.UserInputText != nil { + copiedBlock.UserInputText, err = formatUserInputText(block.UserInputText, vs, formatType) + if err != nil { + return nil, err + } + } + case ContentBlockTypeUserInputImage: + if block.UserInputImage != nil { + copiedBlock.UserInputImage, err = formatUserInputImage(block.UserInputImage, vs, formatType) + if err != nil { + return nil, err + } + } + case ContentBlockTypeUserInputAudio: + if block.UserInputAudio != nil { + copiedBlock.UserInputAudio, err = formatUserInputAudio(block.UserInputAudio, vs, formatType) + if err != nil { + return nil, err + } + } + case ContentBlockTypeUserInputVideo: + if block.UserInputVideo != nil { + copiedBlock.UserInputVideo, err = formatUserInputVideo(block.UserInputVideo, vs, formatType) + if err != nil { + return nil, err + } + } + case ContentBlockTypeUserInputFile: + if block.UserInputFile != nil { + copiedBlock.UserInputFile, err = formatUserInputFile(block.UserInputFile, vs, formatType) + if err != nil { + return nil, err + } + } + } + + copiedBlocks[i] = &copiedBlock + } + copied.ContentBlocks = copiedBlocks + } + + return []*AgenticMessage{&copied}, nil +} + +func formatUserInputText(uit *UserInputText, vs map[string]any, formatType FormatType) (*UserInputText, error) { + text, err := formatContent(uit.Text, vs, formatType) + if err != nil { + return nil, err + } + copied := *uit + copied.Text = text + return &copied, nil +} + +func formatUserInputImage(uii *UserInputImage, vs map[string]any, formatType FormatType) (*UserInputImage, error) { + copied := *uii + if uii.URL != "" { + url, err := formatContent(uii.URL, vs, formatType) + if err != nil { + return nil, err + } + copied.URL = url + } + if uii.Base64Data != "" { + base64data, err := formatContent(uii.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copied.Base64Data = base64data + } + return &copied, nil +} + +func formatUserInputAudio(uia *UserInputAudio, vs map[string]any, formatType FormatType) (*UserInputAudio, error) { + copied := *uia + if uia.URL != "" { + url, err := formatContent(uia.URL, vs, formatType) + if err != nil { + return nil, err + } + copied.URL = url + } + if uia.Base64Data != "" { + base64data, err := formatContent(uia.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copied.Base64Data = base64data + } + return &copied, nil +} + +func formatUserInputVideo(uiv *UserInputVideo, vs map[string]any, formatType FormatType) (*UserInputVideo, error) { + copied := *uiv + if uiv.URL != "" { + url, err := formatContent(uiv.URL, vs, formatType) + if err != nil { + return nil, err + } + copied.URL = url + } + if uiv.Base64Data != "" { + base64data, err := formatContent(uiv.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copied.Base64Data = base64data + } + return &copied, nil +} + +func formatUserInputFile(uif *UserInputFile, vs map[string]any, formatType FormatType) (*UserInputFile, error) { + copied := *uif + if uif.URL != "" { + url, err := formatContent(uif.URL, vs, formatType) + if err != nil { + return nil, err + } + copied.URL = url + } + if uif.Name != "" { + name, err := formatContent(uif.Name, vs, formatType) + if err != nil { + return nil, err + } + copied.Name = name + } + if uif.Base64Data != "" { + base64data, err := formatContent(uif.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copied.Base64Data = base64data + } + return &copied, nil +} + +// ConcatAgenticMessagesArray concatenates multiple streams of AgenticMessage into a single slice of AgenticMessage. +func ConcatAgenticMessagesArray(mas [][]*AgenticMessage) ([]*AgenticMessage, error) { + return buildConcatGenericArray[AgenticMessage](ConcatAgenticMessages)(mas) +} + +// ConcatAgenticMessages concatenates a list of AgenticMessage chunks into a single AgenticMessage. +func ConcatAgenticMessages(msgs []*AgenticMessage) (*AgenticMessage, error) { + var ( + role AgenticRoleType + blocks []*ContentBlock + metas []*AgenticResponseMeta + extra map[string]any + blockIndices []int + indexToBlocks = map[int][]*ContentBlock{} + extraList = make([]map[string]any, 0, len(msgs)) + ) + + if len(msgs) == 1 { + return msgs[0], nil + } + + for idx, msg := range msgs { + if msg == nil { + return nil, fmt.Errorf("message at index %d is nil", idx) + } + + if msg.Role != "" { + if role == "" { + role = msg.Role + } else if role != msg.Role { + return nil, fmt.Errorf("cannot concat messages with different roles: got '%s' and '%s'", role, msg.Role) + } + } + + for _, block := range msg.ContentBlocks { + if block == nil { + continue + } + if block.StreamingMeta == nil { + // Non-streaming block + if len(blockIndices) > 0 { + // Cannot mix streaming and non-streaming blocks + return nil, fmt.Errorf("found non-streaming block after streaming blocks") + } + // Collect non-streaming block + blocks = append(blocks, block) + } else { + // Streaming block + if len(blocks) > 0 { + // Cannot mix non-streaming and streaming blocks + return nil, fmt.Errorf("found streaming block after non-streaming blocks") + } + // Collect streaming block by index + if blocks_, ok := indexToBlocks[block.StreamingMeta.Index]; ok { + indexToBlocks[block.StreamingMeta.Index] = append(blocks_, block) + } else { + blockIndices = append(blockIndices, block.StreamingMeta.Index) + indexToBlocks[block.StreamingMeta.Index] = []*ContentBlock{block} + } + } + } + + if msg.ResponseMeta != nil { + metas = append(metas, msg.ResponseMeta) + } + + if msg.Extra != nil { + extraList = append(extraList, msg.Extra) + } + } + + meta, err := concatAgenticResponseMeta(metas) + if err != nil { + return nil, fmt.Errorf("failed to concat agentic response meta: %w", err) + } + + if len(blockIndices) > 0 { + // All blocks are streaming, concat each group by index + indexToBlock := map[int]*ContentBlock{} + for idx, bs := range indexToBlocks { + var b *ContentBlock + b, err = concatChunksOfSameContentBlock(bs) + if err != nil { + return nil, err + } + indexToBlock[idx] = b + } + blocks = make([]*ContentBlock, 0, len(blockIndices)) + sort.Slice(blockIndices, func(i, j int) bool { + return blockIndices[i] < blockIndices[j] + }) + for _, idx := range blockIndices { + blocks = append(blocks, indexToBlock[idx]) + } + } + + if len(extraList) > 0 { + extra, err = concatExtra(extraList) + if err != nil { + return nil, err + } + } + + return &AgenticMessage{ + Role: role, + ResponseMeta: meta, + ContentBlocks: blocks, + Extra: extra, + }, nil +} + +func concatAgenticResponseMeta(metas []*AgenticResponseMeta) (ret *AgenticResponseMeta, err error) { + if len(metas) == 0 { + return nil, nil + } + + openaiExtensions := make([]*openai.ResponseMetaExtension, 0, len(metas)) + claudeExtensions := make([]*claude.ResponseMetaExtension, 0, len(metas)) + geminiExtensions := make([]*gemini.ResponseMetaExtension, 0, len(metas)) + tokenUsages := make([]*TokenUsage, 0, len(metas)) + + var ( + extType reflect.Type + extensions reflect.Value + ) + + for _, meta := range metas { + if meta.TokenUsage != nil { + tokenUsages = append(tokenUsages, meta.TokenUsage) + } + + var isConsistent bool + + if meta.Extension != nil { + extType, isConsistent = validateExtensionType(extType, meta.Extension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in response meta chunks: '%s' vs '%s'", + extType, reflect.TypeOf(meta.Extension)) + } + if !extensions.IsValid() { + extensions = reflect.MakeSlice(reflect.SliceOf(extType), 0, len(metas)) + } + extensions = reflect.Append(extensions, reflect.ValueOf(meta.Extension)) + } + + if meta.OpenAIExtension != nil { + extType, isConsistent = validateExtensionType(extType, meta.OpenAIExtension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in response meta chunks: '%s' vs '%s'", + extType, reflect.TypeOf(meta.OpenAIExtension)) + } + openaiExtensions = append(openaiExtensions, meta.OpenAIExtension) + } + + if meta.ClaudeExtension != nil { + extType, isConsistent = validateExtensionType(extType, meta.ClaudeExtension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in response meta chunks: '%s' vs '%s'", + extType, reflect.TypeOf(meta.ClaudeExtension)) + } + claudeExtensions = append(claudeExtensions, meta.ClaudeExtension) + } + + if meta.GeminiExtension != nil { + extType, isConsistent = validateExtensionType(extType, meta.GeminiExtension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in response meta chunks: '%s' vs '%s'", + extType, reflect.TypeOf(meta.GeminiExtension)) + } + geminiExtensions = append(geminiExtensions, meta.GeminiExtension) + } + } + + ret = &AgenticResponseMeta{ + TokenUsage: concatTokenUsage(tokenUsages), + } + + if extensions.IsValid() && !extensions.IsZero() { + var extension reflect.Value + extension, err = internal.ConcatSliceValue(extensions) + if err != nil { + return nil, fmt.Errorf("failed to concat extensions: %w", err) + } + ret.Extension = extension.Interface() + } + + if len(openaiExtensions) > 0 { + ret.OpenAIExtension, err = openai.ConcatResponseMetaExtensions(openaiExtensions) + if err != nil { + return nil, fmt.Errorf("failed to concat openai extensions: %w", err) + } + } + + if len(claudeExtensions) > 0 { + ret.ClaudeExtension, err = claude.ConcatResponseMetaExtensions(claudeExtensions) + if err != nil { + return nil, fmt.Errorf("failed to concat claude extensions: %w", err) + } + } + + if len(geminiExtensions) > 0 { + ret.GeminiExtension, err = gemini.ConcatResponseMetaExtensions(geminiExtensions) + if err != nil { + return nil, fmt.Errorf("failed to concat gemini extensions: %w", err) + } + } + + return ret, nil +} + +func concatTokenUsage(usages []*TokenUsage) *TokenUsage { + if len(usages) == 0 { + return nil + } + + ret := &TokenUsage{} + + for _, usage := range usages { + if usage == nil { + continue + } + + if usage.PromptTokens > ret.PromptTokens { + ret.PromptTokens = usage.PromptTokens + } + if usage.CompletionTokens > ret.CompletionTokens { + ret.CompletionTokens = usage.CompletionTokens + } + if usage.TotalTokens > ret.TotalTokens { + ret.TotalTokens = usage.TotalTokens + } + if usage.PromptTokenDetails.CachedTokens > ret.PromptTokenDetails.CachedTokens { + ret.PromptTokenDetails.CachedTokens = usage.PromptTokenDetails.CachedTokens + } + if usage.CompletionTokensDetails.ReasoningTokens > ret.CompletionTokensDetails.ReasoningTokens { + ret.CompletionTokensDetails.ReasoningTokens = usage.CompletionTokensDetails.ReasoningTokens + } + } + + return ret +} + +func concatChunksOfSameContentBlock(blocks []*ContentBlock) (*ContentBlock, error) { + if len(blocks) == 0 { + return nil, fmt.Errorf("no content blocks to concat") + } + + blockType := blocks[0].Type + + switch blockType { + case ContentBlockTypeReasoning: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *Reasoning { return b.Reasoning }, + concatReasoning) + + case ContentBlockTypeUserInputText: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *UserInputText { return b.UserInputText }, + concatUserInputTexts) + + case ContentBlockTypeUserInputImage: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *UserInputImage { return b.UserInputImage }, + concatUserInputImages) + + case ContentBlockTypeUserInputAudio: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *UserInputAudio { return b.UserInputAudio }, + concatUserInputAudios) + + case ContentBlockTypeUserInputVideo: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *UserInputVideo { return b.UserInputVideo }, + concatUserInputVideos) + + case ContentBlockTypeUserInputFile: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *UserInputFile { return b.UserInputFile }, + concatUserInputFiles) + + case ContentBlockTypeToolSearchResult: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *ToolSearchFunctionToolResult { return b.ToolSearchFunctionToolResult }, + concatToolSearchFunctionToolResult) + + case ContentBlockTypeAssistantGenText: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *AssistantGenText { return b.AssistantGenText }, + concatAssistantGenTexts) + + case ContentBlockTypeAssistantGenImage: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *AssistantGenImage { return b.AssistantGenImage }, + concatAssistantGenImages) + + case ContentBlockTypeAssistantGenAudio: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *AssistantGenAudio { return b.AssistantGenAudio }, + concatAssistantGenAudios) + + case ContentBlockTypeAssistantGenVideo: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *AssistantGenVideo { return b.AssistantGenVideo }, + concatAssistantGenVideos) + + case ContentBlockTypeFunctionToolCall: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *FunctionToolCall { return b.FunctionToolCall }, + concatFunctionToolCalls) + + case ContentBlockTypeFunctionToolResult: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *FunctionToolResult { return b.FunctionToolResult }, + concatFunctionToolResults) + + case ContentBlockTypeServerToolCall: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *ServerToolCall { return b.ServerToolCall }, + concatServerToolCalls) + + case ContentBlockTypeServerToolResult: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *ServerToolResult { return b.ServerToolResult }, + concatServerToolResults) + + case ContentBlockTypeMCPToolCall: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *MCPToolCall { return b.MCPToolCall }, + concatMCPToolCalls) + + case ContentBlockTypeMCPToolResult: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *MCPToolResult { return b.MCPToolResult }, + concatMCPToolResults) + + case ContentBlockTypeMCPListToolsResult: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *MCPListToolsResult { return b.MCPListToolsResult }, + concatMCPListToolsResults) + + case ContentBlockTypeMCPToolApprovalRequest: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *MCPToolApprovalRequest { return b.MCPToolApprovalRequest }, + concatMCPToolApprovalRequests) + + case ContentBlockTypeMCPToolApprovalResponse: + return concatContentBlockHelper(blocks, blockType, + func(b *ContentBlock) *MCPToolApprovalResponse { return b.MCPToolApprovalResponse }, + concatMCPToolApprovalResponses) + + default: + return nil, fmt.Errorf("unknown content block type: %s", blockType) + } +} + +// concatContentBlockHelper is a generic helper function that reduces code duplication +// for concatenating content blocks of a specific type. +func concatContentBlockHelper[T contentBlockVariant]( + blocks []*ContentBlock, + expectedType ContentBlockType, + getter func(*ContentBlock) *T, + concatFunc func([]*T) (*T, error), +) (*ContentBlock, error) { + items, err := genericGetTFromContentBlocks(blocks, func(block *ContentBlock) (*T, error) { + if block.Type != expectedType { + return nil, fmt.Errorf("content block type mismatch: expected '%s', but got '%s'", expectedType, block.Type) + } + item := getter(block) + if item == nil { + return nil, fmt.Errorf("'%s' content is nil", expectedType) + } + return item, nil + }) + if err != nil { + return nil, err + } + + concatenated, err := concatFunc(items) + if err != nil { + return nil, fmt.Errorf("failed to concat '%s' content blocks: %w", expectedType, err) + } + + extras := make([]map[string]any, 0, len(blocks)) + for _, block := range blocks { + if len(block.Extra) > 0 { + extras = append(extras, block.Extra) + } + } + + var extra map[string]any + if len(extras) > 0 { + extra, err = internal.ConcatItems(extras) + if err != nil { + return nil, fmt.Errorf("failed to concat content block extras: %w", err) + } + } + + block := NewContentBlock(concatenated) + block.Extra = extra + + return block, nil +} + +func genericGetTFromContentBlocks[T any](blocks []*ContentBlock, checkAndGetter func(block *ContentBlock) (T, error)) ([]T, error) { + ret := make([]T, 0, len(blocks)) + for _, block := range blocks { + t, err := checkAndGetter(block) + if err != nil { + return nil, err + } + ret = append(ret, t) + } + return ret, nil +} + +func concatReasoning(reasons []*Reasoning) (ret *Reasoning, err error) { + if len(reasons) == 0 { + return nil, fmt.Errorf("no reasoning found") + } + + ret = &Reasoning{} + + openaiExtensions := make([]*openai.ReasoningExtension, 0, len(reasons)) + + for _, r := range reasons { + if r == nil { + continue + } + if r.Text != "" { + ret.Text += r.Text + } + if r.Signature != "" { + ret.Signature += r.Signature + } + if r.OpenAIExtension != nil { + openaiExtensions = append(openaiExtensions, r.OpenAIExtension) + } + } + + if len(openaiExtensions) > 0 { + ret.OpenAIExtension, err = openai.ConcatReasoningExtensions(openaiExtensions) + if err != nil { + return nil, fmt.Errorf("failed to concat openai reasoning extensions: %w", err) + } + } + + return ret, nil +} + +func concatUserInputTexts(texts []*UserInputText) (*UserInputText, error) { + if len(texts) == 0 { + return nil, fmt.Errorf("no user input text found") + } + if len(texts) == 1 { + return texts[0], nil + } + return nil, fmt.Errorf("cannot concat multiple user input texts") +} + +func concatUserInputImages(images []*UserInputImage) (*UserInputImage, error) { + if len(images) == 0 { + return nil, fmt.Errorf("no user input image found") + } + if len(images) == 1 { + return images[0], nil + } + return nil, fmt.Errorf("cannot concat multiple user input images") +} + +func concatUserInputAudios(audios []*UserInputAudio) (*UserInputAudio, error) { + if len(audios) == 0 { + return nil, fmt.Errorf("no user input audio found") + } + if len(audios) == 1 { + return audios[0], nil + } + return nil, fmt.Errorf("cannot concat multiple user input audios") +} + +func concatUserInputVideos(videos []*UserInputVideo) (*UserInputVideo, error) { + if len(videos) == 0 { + return nil, fmt.Errorf("no user input video found") + } + if len(videos) == 1 { + return videos[0], nil + } + return nil, fmt.Errorf("cannot concat multiple user input videos") +} + +func concatUserInputFiles(files []*UserInputFile) (*UserInputFile, error) { + if len(files) == 0 { + return nil, fmt.Errorf("no user input file found") + } + if len(files) == 1 { + return files[0], nil + } + return nil, fmt.Errorf("cannot concat multiple user input files") +} + +func concatToolSearchFunctionToolResult(results []*ToolSearchFunctionToolResult) (*ToolSearchFunctionToolResult, error) { + if len(results) == 0 { + return nil, fmt.Errorf("no tool search results found") + } + if len(results) == 1 { + return results[0], nil + } + return nil, fmt.Errorf("cannot concat multiple tool search results") +} + +func concatAssistantGenTexts(texts []*AssistantGenText) (ret *AssistantGenText, err error) { + if len(texts) == 0 { + return nil, fmt.Errorf("no assistant generated text found") + } + if len(texts) == 1 { + return texts[0], nil + } + + ret = &AssistantGenText{} + + openaiExtensions := make([]*openai.AssistantGenTextExtension, 0, len(texts)) + claudeExtensions := make([]*claude.AssistantGenTextExtension, 0, len(texts)) + + var ( + extType reflect.Type + extensions reflect.Value + ) + + for _, t := range texts { + if t == nil { + continue + } + + ret.Text += t.Text + + var isConsistent bool + + if t.Extension != nil { + extType, isConsistent = validateExtensionType(extType, t.Extension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in assistant generated text chunks: '%s' vs '%s'", + extType, reflect.TypeOf(t.Extension)) + } + if !extensions.IsValid() { + extensions = reflect.MakeSlice(reflect.SliceOf(extType), 0, len(texts)) + } + extensions = reflect.Append(extensions, reflect.ValueOf(t.Extension)) + } + + if t.OpenAIExtension != nil { + extType, isConsistent = validateExtensionType(extType, t.OpenAIExtension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in assistant generated text chunks: '%s' vs '%s'", + extType, reflect.TypeOf(t.OpenAIExtension)) + } + openaiExtensions = append(openaiExtensions, t.OpenAIExtension) + } + + if t.ClaudeExtension != nil { + extType, isConsistent = validateExtensionType(extType, t.ClaudeExtension) + if !isConsistent { + return nil, fmt.Errorf("inconsistent extension types in assistant generated text chunks: '%s' vs '%s'", + extType, reflect.TypeOf(t.ClaudeExtension)) + } + claudeExtensions = append(claudeExtensions, t.ClaudeExtension) + } + } + + if extensions.IsValid() && !extensions.IsZero() { + ret.Extension, err = internal.ConcatSliceValue(extensions) + if err != nil { + return nil, err + } + } + + if len(openaiExtensions) > 0 { + ret.OpenAIExtension, err = openai.ConcatAssistantGenTextExtensions(openaiExtensions) + if err != nil { + return nil, err + } + } + + if len(claudeExtensions) > 0 { + ret.ClaudeExtension, err = claude.ConcatAssistantGenTextExtensions(claudeExtensions) + if err != nil { + return nil, err + } + } + + return ret, nil +} + +func concatAssistantGenImages(images []*AssistantGenImage) (*AssistantGenImage, error) { + if len(images) == 0 { + return nil, fmt.Errorf("no assistant gen image found") + } + if len(images) == 1 { + return images[0], nil + } + + ret := &AssistantGenImage{} + + for _, img := range images { + if img == nil { + continue + } + + ret.Base64Data += img.Base64Data + + if ret.URL == "" { + ret.URL = img.URL + } else if img.URL != "" && ret.URL != img.URL { + return nil, fmt.Errorf("inconsistent URLs in assistant generated image chunks: '%s' vs '%s'", ret.URL, img.URL) + } + + if ret.MIMEType == "" { + ret.MIMEType = img.MIMEType + } else if img.MIMEType != "" && ret.MIMEType != img.MIMEType { + return nil, fmt.Errorf("inconsistent MIME types in assistant generated image chunks: '%s' vs '%s'", ret.MIMEType, img.MIMEType) + } + } + + return ret, nil +} + +func concatAssistantGenAudios(audios []*AssistantGenAudio) (*AssistantGenAudio, error) { + if len(audios) == 0 { + return nil, fmt.Errorf("no assistant gen audio found") + } + if len(audios) == 1 { + return audios[0], nil + } + + ret := &AssistantGenAudio{} + + for _, audio := range audios { + if audio == nil { + continue + } + + ret.Base64Data += audio.Base64Data + + if ret.URL == "" { + ret.URL = audio.URL + } else if audio.URL != "" && ret.URL != audio.URL { + return nil, fmt.Errorf("inconsistent URLs in assistant generated audio chunks: '%s' vs '%s'", ret.URL, audio.URL) + } + + if ret.MIMEType == "" { + ret.MIMEType = audio.MIMEType + } else if audio.MIMEType != "" && ret.MIMEType != audio.MIMEType { + return nil, fmt.Errorf("inconsistent MIME types in assistant generated audio chunks: '%s' vs '%s'", ret.MIMEType, audio.MIMEType) + } + } + + return ret, nil +} + +func concatAssistantGenVideos(videos []*AssistantGenVideo) (*AssistantGenVideo, error) { + if len(videos) == 0 { + return nil, fmt.Errorf("no assistant gen video found") + } + if len(videos) == 1 { + return videos[0], nil + } + + ret := &AssistantGenVideo{} + + for _, video := range videos { + if video == nil { + continue + } + + ret.Base64Data += video.Base64Data + + if ret.URL == "" { + ret.URL = video.URL + } else if video.URL != "" && ret.URL != video.URL { + return nil, fmt.Errorf("inconsistent URLs in assistant generated video chunks: '%s' vs '%s'", ret.URL, video.URL) + } + + if ret.MIMEType == "" { + ret.MIMEType = video.MIMEType + } else if video.MIMEType != "" && ret.MIMEType != video.MIMEType { + return nil, fmt.Errorf("inconsistent MIME types in assistant generated video chunks: '%s' vs '%s'", ret.MIMEType, video.MIMEType) + } + } + + return ret, nil +} + +func concatFunctionToolCalls(calls []*FunctionToolCall) (*FunctionToolCall, error) { + if len(calls) == 0 { + return nil, fmt.Errorf("no function tool call found") + } + if len(calls) == 1 { + return calls[0], nil + } + + ret := &FunctionToolCall{} + + for _, c := range calls { + if c == nil { + continue + } + + if ret.CallID == "" { + ret.CallID = c.CallID + } else if c.CallID != "" && c.CallID != ret.CallID { + return nil, fmt.Errorf("expected call ID '%s' for function tool call, but got '%s'", ret.CallID, c.CallID) + } + + if ret.Name == "" { + ret.Name = c.Name + } else if c.Name != "" && c.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for function tool call, but got '%s'", ret.Name, c.Name) + } + + ret.Arguments += c.Arguments + } + + return ret, nil +} + +func concatFunctionToolResults(results []*FunctionToolResult) (*FunctionToolResult, error) { + if len(results) == 0 { + return nil, fmt.Errorf("no function tool result found") + } + if len(results) == 1 { + return results[0], nil + } + + ret := &FunctionToolResult{} + + for _, r := range results { + if r == nil { + continue + } + + if ret.CallID == "" { + ret.CallID = r.CallID + } else if r.CallID != "" && r.CallID != ret.CallID { + return nil, fmt.Errorf("expected call ID '%s' for function tool result, but got '%s'", ret.CallID, r.CallID) + } + + if ret.Name == "" { + ret.Name = r.Name + } else if r.Name != "" && r.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for function tool result, but got '%s'", ret.Name, r.Name) + } + + for _, b := range r.Content { + if b == nil { + continue + } + ret.Content = append(ret.Content, b) + } + } + + return ret, nil +} + +func concatServerToolCalls(calls []*ServerToolCall) (ret *ServerToolCall, err error) { + if len(calls) == 0 { + return nil, fmt.Errorf("no server tool call found") + } + if len(calls) == 1 { + return calls[0], nil + } + + ret = &ServerToolCall{} + + var ( + argsType reflect.Type + argsChunks reflect.Value + ) + + for _, c := range calls { + if c == nil { + continue + } + + if ret.CallID == "" { + ret.CallID = c.CallID + } else if c.CallID != "" && c.CallID != ret.CallID { + return nil, fmt.Errorf("expected call ID '%s' for server tool call, but got '%s'", ret.CallID, c.CallID) + } + + if ret.Name == "" { + ret.Name = c.Name + } else if c.Name != "" && c.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for server tool call, but got '%s'", ret.Name, c.Name) + } + + if c.Arguments != nil { + argsType_ := reflect.TypeOf(c.Arguments) + if argsType == nil { + argsType = argsType_ + argsChunks = reflect.MakeSlice(reflect.SliceOf(argsType), 0, len(calls)) + } else if argsType != argsType_ { + return nil, fmt.Errorf("expected type '%s' for server tool call arguments, but got '%s'", argsType, argsType_) + } + argsChunks = reflect.Append(argsChunks, reflect.ValueOf(c.Arguments)) + } + } + + if argsChunks.IsValid() && !argsChunks.IsZero() { + arguments, err := internal.ConcatSliceValue(argsChunks) + if err != nil { + return nil, err + } + ret.Arguments = arguments.Interface() + } + + return ret, nil +} + +func concatServerToolResults(results []*ServerToolResult) (ret *ServerToolResult, err error) { + if len(results) == 0 { + return nil, fmt.Errorf("no server tool result found") + } + if len(results) == 1 { + return results[0], nil + } + + ret = &ServerToolResult{} + + var ( + resType reflect.Type + resChunks reflect.Value + ) + + for _, r := range results { + if r == nil { + continue + } + + if ret.CallID == "" { + ret.CallID = r.CallID + } else if r.CallID != "" && r.CallID != ret.CallID { + return nil, fmt.Errorf("expected call ID '%s' for server tool result, but got '%s'", ret.CallID, r.CallID) + } + + if ret.Name == "" { + ret.Name = r.Name + } else if r.Name != "" && r.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for server tool result, but got '%s'", ret.Name, r.Name) + } + + if r.Content != nil { + resType_ := reflect.TypeOf(r.Content) + if resType == nil { + resType = resType_ + resChunks = reflect.MakeSlice(reflect.SliceOf(resType), 0, len(results)) + } else if resType != resType_ { + return nil, fmt.Errorf("expected type '%s' for server tool result, but got '%s'", resType, resType_) + } + resChunks = reflect.Append(resChunks, reflect.ValueOf(r.Content)) + } + } + + if resChunks.IsValid() && !resChunks.IsZero() { + result, err := internal.ConcatSliceValue(resChunks) + if err != nil { + return nil, fmt.Errorf("failed to concat server tool result: %v", err) + } + ret.Content = result.Interface() + } + + return ret, nil +} + +func concatMCPToolCalls(calls []*MCPToolCall) (*MCPToolCall, error) { + if len(calls) == 0 { + return nil, fmt.Errorf("no mcp tool call found") + } + if len(calls) == 1 { + return calls[0], nil + } + + ret := &MCPToolCall{} + + for _, c := range calls { + if c == nil { + continue + } + + ret.Arguments += c.Arguments + + if ret.ServerLabel == "" { + ret.ServerLabel = c.ServerLabel + } else if c.ServerLabel != "" && c.ServerLabel != ret.ServerLabel { + return nil, fmt.Errorf("expected server label '%s' for mcp tool call, but got '%s'", ret.ServerLabel, c.ServerLabel) + } + + if ret.CallID == "" { + ret.CallID = c.CallID + } else if c.CallID != "" && c.CallID != ret.CallID { + return nil, fmt.Errorf("expected call ID '%s' for mcp tool call, but got '%s'", ret.CallID, c.CallID) + } + + if ret.Name == "" { + ret.Name = c.Name + } else if c.Name != "" && c.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for mcp tool call, but got '%s'", ret.Name, c.Name) + } + } + + return ret, nil +} + +func concatMCPToolResults(results []*MCPToolResult) (*MCPToolResult, error) { + if len(results) == 0 { + return nil, fmt.Errorf("no mcp tool result found") + } + if len(results) == 1 { + return results[0], nil + } + + ret := &MCPToolResult{} + + for _, r := range results { + if r == nil { + continue + } + + if r.Content != "" { + ret.Content = r.Content + } + + if ret.ServerLabel == "" { + ret.ServerLabel = r.ServerLabel + } else if r.ServerLabel != "" && r.ServerLabel != ret.ServerLabel { + return nil, fmt.Errorf("expected server label '%s' for mcp tool result, but got '%s'", ret.ServerLabel, r.ServerLabel) + } + + if ret.CallID == "" { + ret.CallID = r.CallID + } else if r.CallID != "" && r.CallID != ret.CallID { + return nil, fmt.Errorf("expected call ID '%s' for mcp tool result, but got '%s'", ret.CallID, r.CallID) + } + + if ret.Name == "" { + ret.Name = r.Name + } else if r.Name != "" && r.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for mcp tool result, but got '%s'", ret.Name, r.Name) + } + + if r.Error != nil { + ret.Error = r.Error + } + } + + return ret, nil +} + +func concatMCPListToolsResults(results []*MCPListToolsResult) (*MCPListToolsResult, error) { + if len(results) == 0 { + return nil, fmt.Errorf("no mcp list tools result found") + } + if len(results) == 1 { + return results[0], nil + } + + ret := &MCPListToolsResult{} + + for _, r := range results { + if r == nil { + continue + } + + ret.Tools = append(ret.Tools, r.Tools...) + + if r.Error != "" { + ret.Error = r.Error + } + + if ret.ServerLabel == "" { + ret.ServerLabel = r.ServerLabel + } else if r.ServerLabel != "" && r.ServerLabel != ret.ServerLabel { + return nil, fmt.Errorf("expected server label '%s' for mcp list tools result, but got '%s'", ret.ServerLabel, r.ServerLabel) + } + } + + return ret, nil +} + +func concatMCPToolApprovalRequests(requests []*MCPToolApprovalRequest) (*MCPToolApprovalRequest, error) { + if len(requests) == 0 { + return nil, fmt.Errorf("no mcp tool approval request found") + } + if len(requests) == 1 { + return requests[0], nil + } + + ret := &MCPToolApprovalRequest{} + + for _, r := range requests { + if r == nil { + continue + } + + ret.Arguments += r.Arguments + + if ret.ID == "" { + ret.ID = r.ID + } else if r.ID != "" && r.ID != ret.ID { + return nil, fmt.Errorf("expected request ID '%s' for mcp tool approval request, but got '%s'", ret.ID, r.ID) + } + + if ret.Name == "" { + ret.Name = r.Name + } else if r.Name != "" && r.Name != ret.Name { + return nil, fmt.Errorf("expected tool name '%s' for mcp tool approval request, but got '%s'", ret.Name, r.Name) + } + + if ret.ServerLabel == "" { + ret.ServerLabel = r.ServerLabel + } else if r.ServerLabel != "" && r.ServerLabel != ret.ServerLabel { + return nil, fmt.Errorf("expected server label '%s' for mcp tool approval request, but got '%s'", ret.ServerLabel, r.ServerLabel) + } + } + + return ret, nil +} + +func concatMCPToolApprovalResponses(responses []*MCPToolApprovalResponse) (*MCPToolApprovalResponse, error) { + if len(responses) == 0 { + return nil, fmt.Errorf("no mcp tool approval response found") + } + if len(responses) == 1 { + return responses[0], nil + } + return nil, fmt.Errorf("cannot concat multiple mcp tool approval responses") +} + +// String returns the string representation of AgenticMessage. +func (m *AgenticMessage) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf("role: %s\n", m.Role)) + + if len(m.ContentBlocks) > 0 { + sb.WriteString("content_blocks:\n") + for i, block := range m.ContentBlocks { + if block == nil { + continue + } + sb.WriteString(fmt.Sprintf(" [%d] %s", i, block.String())) + } + } + + if m.ResponseMeta != nil { + sb.WriteString(m.ResponseMeta.String()) + } + + return sb.String() +} + +// String returns the string representation of ContentBlock. +// nolint +func (b *ContentBlock) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf("type: %s\n", b.Type)) + + switch b.Type { + case ContentBlockTypeReasoning: + if b.Reasoning != nil { + sb.WriteString(b.Reasoning.String()) + } + case ContentBlockTypeUserInputText: + if b.UserInputText != nil { + sb.WriteString(b.UserInputText.String()) + } + case ContentBlockTypeUserInputImage: + if b.UserInputImage != nil { + sb.WriteString(b.UserInputImage.String()) + } + case ContentBlockTypeUserInputAudio: + if b.UserInputAudio != nil { + sb.WriteString(b.UserInputAudio.String()) + } + case ContentBlockTypeUserInputVideo: + if b.UserInputVideo != nil { + sb.WriteString(b.UserInputVideo.String()) + } + case ContentBlockTypeUserInputFile: + if b.UserInputFile != nil { + sb.WriteString(b.UserInputFile.String()) + } + case ContentBlockTypeToolSearchResult: + if b.ToolSearchFunctionToolResult != nil { + sb.WriteString(b.ToolSearchFunctionToolResult.String()) + } + case ContentBlockTypeAssistantGenText: + if b.AssistantGenText != nil { + sb.WriteString(b.AssistantGenText.String()) + } + case ContentBlockTypeAssistantGenImage: + if b.AssistantGenImage != nil { + sb.WriteString(b.AssistantGenImage.String()) + } + case ContentBlockTypeAssistantGenAudio: + if b.AssistantGenAudio != nil { + sb.WriteString(b.AssistantGenAudio.String()) + } + case ContentBlockTypeAssistantGenVideo: + if b.AssistantGenVideo != nil { + sb.WriteString(b.AssistantGenVideo.String()) + } + case ContentBlockTypeFunctionToolCall: + if b.FunctionToolCall != nil { + sb.WriteString(b.FunctionToolCall.String()) + } + case ContentBlockTypeFunctionToolResult: + if b.FunctionToolResult != nil { + sb.WriteString(b.FunctionToolResult.String()) + } + case ContentBlockTypeServerToolCall: + if b.ServerToolCall != nil { + sb.WriteString(b.ServerToolCall.String()) + } + case ContentBlockTypeServerToolResult: + if b.ServerToolResult != nil { + sb.WriteString(b.ServerToolResult.String()) + } + case ContentBlockTypeMCPToolCall: + if b.MCPToolCall != nil { + sb.WriteString(b.MCPToolCall.String()) + } + case ContentBlockTypeMCPToolResult: + if b.MCPToolResult != nil { + sb.WriteString(b.MCPToolResult.String()) + } + case ContentBlockTypeMCPListToolsResult: + if b.MCPListToolsResult != nil { + sb.WriteString(b.MCPListToolsResult.String()) + } + case ContentBlockTypeMCPToolApprovalRequest: + if b.MCPToolApprovalRequest != nil { + sb.WriteString(b.MCPToolApprovalRequest.String()) + } + case ContentBlockTypeMCPToolApprovalResponse: + if b.MCPToolApprovalResponse != nil { + sb.WriteString(b.MCPToolApprovalResponse.String()) + } + } + + if b.StreamingMeta != nil { + sb.WriteString(fmt.Sprintf(" stream_index: %d\n", b.StreamingMeta.Index)) + } + + return sb.String() +} + +// String returns the string representation of Reasoning. +func (r *Reasoning) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" text: %s\n", r.Text)) + if r.Signature != "" { + sb.WriteString(fmt.Sprintf(" signature: %s\n", truncateString(r.Signature, 50))) + } + return sb.String() +} + +// String returns the string representation of UserInputText. +func (u *UserInputText) String() string { + return fmt.Sprintf(" text: %s\n", u.Text) +} + +// String returns the string representation of UserInputImage. +func (u *UserInputImage) String() string { + return formatMediaString(u.URL, u.Base64Data, u.MIMEType, string(u.Detail)) +} + +// String returns the string representation of UserInputAudio. +func (u *UserInputAudio) String() string { + return formatMediaString(u.URL, u.Base64Data, u.MIMEType, "") +} + +// String returns the string representation of UserInputVideo. +func (u *UserInputVideo) String() string { + return formatMediaString(u.URL, u.Base64Data, u.MIMEType, "") +} + +// String returns the string representation of UserInputFile. +func (u *UserInputFile) String() string { + sb := &strings.Builder{} + if u.Name != "" { + sb.WriteString(fmt.Sprintf(" name: %s\n", u.Name)) + } + sb.WriteString(formatMediaString(u.URL, u.Base64Data, u.MIMEType, "")) + return sb.String() +} + +// String returns the string representation of AssistantGenText. +func (a *AssistantGenText) String() string { + return fmt.Sprintf(" text: %s\n", a.Text) +} + +// String returns the string representation of AssistantGenImage. +func (a *AssistantGenImage) String() string { + return formatMediaString(a.URL, a.Base64Data, a.MIMEType, "") +} + +// String returns the string representation of AssistantGenAudio. +func (a *AssistantGenAudio) String() string { + return formatMediaString(a.URL, a.Base64Data, a.MIMEType, "") +} + +// String returns the string representation of AssistantGenVideo. +func (a *AssistantGenVideo) String() string { + return formatMediaString(a.URL, a.Base64Data, a.MIMEType, "") +} + +// String returns the string representation of FunctionToolCall. +func (f *FunctionToolCall) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" call_id: %s\n", f.CallID)) + sb.WriteString(fmt.Sprintf(" name: %s\n", f.Name)) + sb.WriteString(fmt.Sprintf(" arguments: %s\n", f.Arguments)) + return sb.String() +} + +// String returns the string representation of FunctionToolResult. +func (f *FunctionToolResult) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" call_id: %s\n", f.CallID)) + sb.WriteString(fmt.Sprintf(" name: %s\n", f.Name)) + if len(f.Content) > 0 { + sb.WriteString(fmt.Sprintf(" content: (%d blocks)\n", len(f.Content))) + for i, block := range f.Content { + if block == nil { + continue + } + sb.WriteString(fmt.Sprintf(" [%d] %s", i, block.String())) + } + } + return sb.String() +} + +// String returns the string representation of ServerToolCall. +func (s *ServerToolCall) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" name: %s\n", s.Name)) + if s.CallID != "" { + sb.WriteString(fmt.Sprintf(" call_id: %s\n", s.CallID)) + } + sb.WriteString(fmt.Sprintf(" arguments: %s\n", printAny(s.Arguments))) + return sb.String() +} + +// String returns the string representation of ServerToolResult. +func (s *ServerToolResult) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" name: %s\n", s.Name)) + if s.CallID != "" { + sb.WriteString(fmt.Sprintf(" call_id: %s\n", s.CallID)) + } + sb.WriteString(fmt.Sprintf(" content: %s\n", printAny(s.Content))) + return sb.String() +} + +// String returns the string representation of MCPToolCall. +func (m *MCPToolCall) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" server_label: %s\n", m.ServerLabel)) + sb.WriteString(fmt.Sprintf(" call_id: %s\n", m.CallID)) + sb.WriteString(fmt.Sprintf(" name: %s\n", m.Name)) + sb.WriteString(fmt.Sprintf(" arguments: %s\n", m.Arguments)) + return sb.String() +} + +// String returns the string representation of MCPToolResult. +func (m *MCPToolResult) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" call_id: %s\n", m.CallID)) + sb.WriteString(fmt.Sprintf(" name: %s\n", m.Name)) + sb.WriteString(fmt.Sprintf(" content: %s\n", m.Content)) + if m.Error != nil { + if m.Error.Code != nil { + sb.WriteString(fmt.Sprintf(" error: [%d] %s\n", *m.Error.Code, m.Error.Message)) + } else { + sb.WriteString(fmt.Sprintf(" error: %s\n", m.Error.Message)) + } + } + return sb.String() +} + +// String returns the string representation of MCPListToolsResult. +func (m *MCPListToolsResult) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" server_label: %s\n", m.ServerLabel)) + sb.WriteString(fmt.Sprintf(" tools: %d items\n", len(m.Tools))) + for _, tool := range m.Tools { + sb.WriteString(fmt.Sprintf(" - %s: %s\n", tool.Name, tool.Description)) + } + if m.Error != "" { + sb.WriteString(fmt.Sprintf(" error: %s\n", m.Error)) + } + return sb.String() +} + +// String returns the string representation of MCPToolApprovalRequest. +func (m *MCPToolApprovalRequest) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" server_label: %s\n", m.ServerLabel)) + sb.WriteString(fmt.Sprintf(" id: %s\n", m.ID)) + sb.WriteString(fmt.Sprintf(" name: %s\n", m.Name)) + sb.WriteString(fmt.Sprintf(" arguments: %s\n", m.Arguments)) + return sb.String() +} + +// String returns the string representation of MCPToolApprovalResponse. +func (m *MCPToolApprovalResponse) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf(" approval_request_id: %s\n", m.ApprovalRequestID)) + sb.WriteString(fmt.Sprintf(" approve: %v\n", m.Approve)) + if m.Reason != "" { + sb.WriteString(fmt.Sprintf(" reason: %s\n", m.Reason)) + } + return sb.String() +} + +// String returns the string representation of AgenticResponseMeta. +func (a *AgenticResponseMeta) String() string { + sb := &strings.Builder{} + sb.WriteString("response_meta:\n") + if a.TokenUsage != nil { + sb.WriteString(fmt.Sprintf(" token_usage: prompt=%d, completion=%d, total=%d\n", + a.TokenUsage.PromptTokens, + a.TokenUsage.CompletionTokens, + a.TokenUsage.TotalTokens)) + } + return sb.String() +} + +// truncateString truncates a string to maxLen characters, adding "..." if truncated +func truncateString(s string, maxLen int) string { + if len(s) <= maxLen { + return s + } + return s[:maxLen] + "..." +} + +// formatMediaString formats URL, Base64Data, MIMEType and Detail for media content +func formatMediaString(url, base64Data string, mimeType string, detail string) string { + sb := &strings.Builder{} + if url != "" { + sb.WriteString(fmt.Sprintf(" url: %s\n", truncateString(url, 100))) + } + if base64Data != "" { + // Only show first few characters of base64 data + sb.WriteString(fmt.Sprintf(" base64_data: %s... (%d bytes)\n", truncateString(base64Data, 20), len(base64Data))) + } + if mimeType != "" { + sb.WriteString(fmt.Sprintf(" mime_type: %s\n", mimeType)) + } + if detail != "" { + sb.WriteString(fmt.Sprintf(" detail: %s\n", detail)) + } + return sb.String() +} + +func validateExtensionType(expected reflect.Type, actual any) (reflect.Type, bool) { + if actual == nil { + return expected, true + } + actualType := reflect.TypeOf(actual) + if expected == nil { + return actualType, true + } + if expected != actualType { + return expected, false + } + return expected, true +} + +func printAny(a any) string { + switch v := a.(type) { + case string: + return v + case fmt.Stringer: + return v.String() + default: + b, err := json.MarshalIndent(a, "", " ") + if err != nil { + return fmt.Sprintf("%v", a) + } + return string(b) + } +} diff --git a/schema/agentic_message_test.go b/schema/agentic_message_test.go new file mode 100644 index 0000000..32dc96c --- /dev/null +++ b/schema/agentic_message_test.go @@ -0,0 +1,1729 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "context" + "reflect" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConcatAgenticMessages(t *testing.T) { + t.Run("single message", func(t *testing.T) { + msg := &AgenticMessage{ + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Hello", + }, + }, + }, + } + + result, err := ConcatAgenticMessages([]*AgenticMessage{msg}) + assert.NoError(t, err) + assert.Equal(t, msg, result) + }) + + t.Run("nil message in stream", func(t *testing.T) { + msgs := []*AgenticMessage{ + {Role: AgenticRoleTypeAssistant}, + nil, + {Role: AgenticRoleTypeAssistant}, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "message at index 1 is nil") + }) + + t.Run("different roles", func(t *testing.T) { + msgs := []*AgenticMessage{ + {Role: AgenticRoleTypeUser}, + {Role: AgenticRoleTypeAssistant}, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "cannot concat messages with different roles") + }) + + t.Run("concat text blocks", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Hello ", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "World!", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, AgenticRoleTypeAssistant, result.Role) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "Hello World!", result.ContentBlocks[0].AssistantGenText.Text) + }) + + t.Run("concat reasoning with nil index", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeReasoning, + Reasoning: &Reasoning{ + Text: "First ", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeReasoning, + Reasoning: &Reasoning{ + Text: "Second", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "First Second", result.ContentBlocks[0].Reasoning.Text) + }) + + t.Run("concat reasoning with index", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeReasoning, + Reasoning: &Reasoning{ + Text: "Part1-", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeReasoning, + Reasoning: &Reasoning{ + Text: "Part3", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "Part1-Part3", result.ContentBlocks[0].Reasoning.Text) + }) + + t.Run("concat user input text", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Hello ", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "World!", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "Hello World!", result.ContentBlocks[0].AssistantGenText.Text) + }) + + t.Run("concat assistant gen image", func(t *testing.T) { + base1 := "1" + base2 := "2" + + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenImage, + AssistantGenImage: &AssistantGenImage{ + Base64Data: base1, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenImage, + AssistantGenImage: &AssistantGenImage{ + Base64Data: base2, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "12", result.ContentBlocks[0].AssistantGenImage.Base64Data) + }) + + t.Run("concat user input audio - should error", func(t *testing.T) { + url1 := "https://example.com/audio1.mp3" + url2 := "https://example.com/audio2.mp3" + + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputAudio, + UserInputAudio: &UserInputAudio{ + URL: url1, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputAudio, + UserInputAudio: &UserInputAudio{ + URL: url2, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "cannot concat multiple user input audios") + }) + + t.Run("concat user input video - should error", func(t *testing.T) { + url1 := "https://example.com/video1.mp4" + url2 := "https://example.com/video2.mp4" + + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputVideo, + UserInputVideo: &UserInputVideo{ + URL: url1, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputVideo, + UserInputVideo: &UserInputVideo{ + URL: url2, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "cannot concat multiple user input videos") + }) + + t.Run("concat assistant gen text", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Generated ", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Text", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "Generated Text", result.ContentBlocks[0].AssistantGenText.Text) + }) + + t.Run("concat assistant gen image", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenImage, + AssistantGenImage: &AssistantGenImage{ + Base64Data: "part1", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenImage, + AssistantGenImage: &AssistantGenImage{ + Base64Data: "part2", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "part1part2", result.ContentBlocks[0].AssistantGenImage.Base64Data) + }) + + t.Run("concat assistant gen audio", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenAudio, + AssistantGenAudio: &AssistantGenAudio{ + Base64Data: "audio1", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenAudio, + AssistantGenAudio: &AssistantGenAudio{ + Base64Data: "audio2", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "audio1audio2", result.ContentBlocks[0].AssistantGenAudio.Base64Data) + }) + + t.Run("concat assistant gen video", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenVideo, + AssistantGenVideo: &AssistantGenVideo{ + Base64Data: "video1", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenVideo, + AssistantGenVideo: &AssistantGenVideo{ + Base64Data: "video2", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "video1video2", result.ContentBlocks[0].AssistantGenVideo.Base64Data) + }) + + t.Run("concat function tool call", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeFunctionToolCall, + FunctionToolCall: &FunctionToolCall{ + CallID: "call_123", + Name: "get_weather", + Arguments: `{"location`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeFunctionToolCall, + FunctionToolCall: &FunctionToolCall{ + Arguments: `":"NYC"}`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "call_123", result.ContentBlocks[0].FunctionToolCall.CallID) + assert.Equal(t, "get_weather", result.ContentBlocks[0].FunctionToolCall.Name) + assert.Equal(t, `{"location":"NYC"}`, result.ContentBlocks[0].FunctionToolCall.Arguments) + }) + + t.Run("concat function tool result", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeFunctionToolResult, + FunctionToolResult: &FunctionToolResult{ + CallID: "call_123", + Name: "get_weather", + Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: `{"temp`}}, + }, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeFunctionToolResult, + FunctionToolResult: &FunctionToolResult{ + Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: `":72}`}}, + }, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "call_123", result.ContentBlocks[0].FunctionToolResult.CallID) + assert.Equal(t, "get_weather", result.ContentBlocks[0].FunctionToolResult.Name) + assert.Equal(t, 2, len(result.ContentBlocks[0].FunctionToolResult.Content)) + assert.Equal(t, `{"temp`, result.ContentBlocks[0].FunctionToolResult.Content[0].Text.Text) + assert.Equal(t, `":72}`, result.ContentBlocks[0].FunctionToolResult.Content[1].Text.Text) + }) + + t.Run("concat server tool call", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeServerToolCall, + ServerToolCall: &ServerToolCall{ + CallID: "server_call_1", + Name: "server_func", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeServerToolCall, + ServerToolCall: &ServerToolCall{ + Arguments: map[string]any{"key": "value"}, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "server_call_1", result.ContentBlocks[0].ServerToolCall.CallID) + assert.Equal(t, "server_func", result.ContentBlocks[0].ServerToolCall.Name) + assert.NotNil(t, result.ContentBlocks[0].ServerToolCall.Arguments) + }) + + t.Run("concat server tool result", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeServerToolResult, + ServerToolResult: &ServerToolResult{ + CallID: "server_call_1", + Name: "server_func", + Content: "result1", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeServerToolResult, + ServerToolResult: &ServerToolResult{}, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "server_call_1", result.ContentBlocks[0].ServerToolResult.CallID) + assert.Equal(t, "server_func", result.ContentBlocks[0].ServerToolResult.Name) + assert.Equal(t, "result1", result.ContentBlocks[0].ServerToolResult.Content) + }) + + t.Run("concat mcp tool call", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolCall, + MCPToolCall: &MCPToolCall{ + ServerLabel: "mcp-server", + CallID: "mcp_call_1", + Name: "mcp_func", + Arguments: `{"arg`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolCall, + MCPToolCall: &MCPToolCall{ + Arguments: `":123}`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "mcp-server", result.ContentBlocks[0].MCPToolCall.ServerLabel) + assert.Equal(t, "mcp_call_1", result.ContentBlocks[0].MCPToolCall.CallID) + assert.Equal(t, "mcp_func", result.ContentBlocks[0].MCPToolCall.Name) + assert.Equal(t, `{"arg":123}`, result.ContentBlocks[0].MCPToolCall.Arguments) + }) + + t.Run("concat mcp tool result", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolResult, + MCPToolResult: &MCPToolResult{ + ServerLabel: "mcp-server", + CallID: "mcp_call_1", + Name: "mcp_func", + Content: `First`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolResult, + MCPToolResult: &MCPToolResult{ + Content: `Second`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "mcp-server", result.ContentBlocks[0].MCPToolResult.ServerLabel) + assert.Equal(t, "mcp_call_1", result.ContentBlocks[0].MCPToolResult.CallID) + assert.Equal(t, "mcp_func", result.ContentBlocks[0].MCPToolResult.Name) + assert.Equal(t, `Second`, result.ContentBlocks[0].MCPToolResult.Content) + }) + + t.Run("concat mcp list tools", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPListToolsResult, + MCPListToolsResult: &MCPListToolsResult{ + ServerLabel: "mcp-server", + Tools: []*MCPListToolsItem{ + {Name: "tool1"}, + }, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPListToolsResult, + MCPListToolsResult: &MCPListToolsResult{ + Tools: []*MCPListToolsItem{ + {Name: "tool2"}, + }, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "mcp-server", result.ContentBlocks[0].MCPListToolsResult.ServerLabel) + assert.Len(t, result.ContentBlocks[0].MCPListToolsResult.Tools, 2) + }) + + t.Run("concat mcp tool approval request", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolApprovalRequest, + MCPToolApprovalRequest: &MCPToolApprovalRequest{ + ID: "approval_1", + Name: "approval_func", + ServerLabel: "mcp-server", + Arguments: `{"request`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolApprovalRequest, + MCPToolApprovalRequest: &MCPToolApprovalRequest{ + Arguments: `":1}`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "approval_1", result.ContentBlocks[0].MCPToolApprovalRequest.ID) + assert.Equal(t, "approval_func", result.ContentBlocks[0].MCPToolApprovalRequest.Name) + assert.Equal(t, "mcp-server", result.ContentBlocks[0].MCPToolApprovalRequest.ServerLabel) + assert.Equal(t, `{"request":1}`, result.ContentBlocks[0].MCPToolApprovalRequest.Arguments) + }) + + t.Run("concat mcp tool approval response - should error", func(t *testing.T) { + response1 := &MCPToolApprovalResponse{ + ApprovalRequestID: "approval_1", + Approve: false, + } + response2 := &MCPToolApprovalResponse{ + ApprovalRequestID: "approval_1", + Approve: true, + } + + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolApprovalResponse, + MCPToolApprovalResponse: response1, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolApprovalResponse, + MCPToolApprovalResponse: response2, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "cannot concat multiple mcp tool approval responses") + }) + + t.Run("concat response meta", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ResponseMeta: &AgenticResponseMeta{ + TokenUsage: &TokenUsage{ + PromptTokens: 10, + CompletionTokens: 5, + PromptTokenDetails: PromptTokenDetails{ + CachedTokens: 3, + }, + CompletionTokensDetails: CompletionTokensDetails{ + ReasoningTokens: 2, + }, + TotalTokens: 15, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ResponseMeta: &AgenticResponseMeta{ + TokenUsage: &TokenUsage{ + PromptTokens: 8, + CompletionTokens: 15, + PromptTokenDetails: PromptTokenDetails{ + CachedTokens: 5, + }, + CompletionTokensDetails: CompletionTokensDetails{ + ReasoningTokens: 4, + }, + TotalTokens: 23, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.NotNil(t, result.ResponseMeta) + assert.Equal(t, 15, result.ResponseMeta.TokenUsage.CompletionTokens) + assert.Equal(t, 10, result.ResponseMeta.TokenUsage.PromptTokens) + assert.Equal(t, 23, result.ResponseMeta.TokenUsage.TotalTokens) + assert.Equal(t, 5, result.ResponseMeta.TokenUsage.PromptTokenDetails.CachedTokens) + assert.Equal(t, 4, result.ResponseMeta.TokenUsage.CompletionTokensDetails.ReasoningTokens) + }) + + t.Run("mixed streaming and non-streaming blocks error", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Hello", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "World", + }, + // No StreamingMeta - non-streaming + }, + }, + }, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "found non-streaming block after streaming blocks") + }) + + t.Run("concat MCP tool call", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolCall, + MCPToolCall: &MCPToolCall{ + ServerLabel: "mcp-server", + CallID: "call_456", + Name: "list_files", + Arguments: `{"path`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeMCPToolCall, + MCPToolCall: &MCPToolCall{ + Arguments: `":"/tmp"}`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 1) + assert.Equal(t, "mcp-server", result.ContentBlocks[0].MCPToolCall.ServerLabel) + assert.Equal(t, "call_456", result.ContentBlocks[0].MCPToolCall.CallID) + assert.Equal(t, `{"path":"/tmp"}`, result.ContentBlocks[0].MCPToolCall.Arguments) + }) + + t.Run("concat user input text - should error", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputText, + UserInputText: &UserInputText{ + Text: "What is ", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + { + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputText, + UserInputText: &UserInputText{ + Text: "the weather?", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + }, + }, + } + + _, err := ConcatAgenticMessages(msgs) + assert.Error(t, err) + assert.ErrorContains(t, err, "cannot concat multiple user input texts") + }) + + t.Run("multiple stream indexes - sparse indexes", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Index0-", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Index2-", + }, + StreamingMeta: &StreamingMeta{Index: 2}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Part2", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Part2", + }, + StreamingMeta: &StreamingMeta{Index: 2}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 2) + assert.Equal(t, "Index0-Part2", result.ContentBlocks[0].AssistantGenText.Text) + assert.Equal(t, "Index2-Part2", result.ContentBlocks[1].AssistantGenText.Text) + }) + + t.Run("multiple stream indexes - mixed content types", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Text ", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeFunctionToolCall, + FunctionToolCall: &FunctionToolCall{ + CallID: "call_1", + Name: "func1", + Arguments: `{"a`, + }, + StreamingMeta: &StreamingMeta{Index: 1}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "Content", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeFunctionToolCall, + FunctionToolCall: &FunctionToolCall{ + Arguments: `":1}`, + }, + StreamingMeta: &StreamingMeta{Index: 1}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 2) + assert.Equal(t, "Text Content", result.ContentBlocks[0].AssistantGenText.Text) + assert.Equal(t, "call_1", result.ContentBlocks[1].FunctionToolCall.CallID) + assert.Equal(t, "func1", result.ContentBlocks[1].FunctionToolCall.Name) + assert.Equal(t, `{"a":1}`, result.ContentBlocks[1].FunctionToolCall.Arguments) + }) + + t.Run("multiple stream indexes - three indexes", func(t *testing.T) { + msgs := []*AgenticMessage{ + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "A", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "B", + }, + StreamingMeta: &StreamingMeta{Index: 1}, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "C", + }, + StreamingMeta: &StreamingMeta{Index: 2}, + }, + }, + }, + { + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "1", + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "2", + }, + StreamingMeta: &StreamingMeta{Index: 1}, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "3", + }, + StreamingMeta: &StreamingMeta{Index: 2}, + }, + }, + }, + } + + result, err := ConcatAgenticMessages(msgs) + assert.NoError(t, err) + assert.Len(t, result.ContentBlocks, 3) + assert.Equal(t, "A1", result.ContentBlocks[0].AssistantGenText.Text) + assert.Equal(t, "B2", result.ContentBlocks[1].AssistantGenText.Text) + assert.Equal(t, "C3", result.ContentBlocks[2].AssistantGenText.Text) + }) +} + +func TestAgenticMessageFormat(t *testing.T) { + m := &AgenticMessage{ + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputText, + UserInputText: &UserInputText{Text: "{a}"}, + }, + { + Type: ContentBlockTypeUserInputImage, + UserInputImage: &UserInputImage{ + URL: "{b}", + Base64Data: "{c}", + }, + }, + { + Type: ContentBlockTypeUserInputAudio, + UserInputAudio: &UserInputAudio{ + URL: "{d}", + Base64Data: "{e}", + }, + }, + { + Type: ContentBlockTypeUserInputVideo, + UserInputVideo: &UserInputVideo{ + URL: "{f}", + Base64Data: "{g}", + }, + }, + { + Type: ContentBlockTypeUserInputFile, + UserInputFile: &UserInputFile{ + URL: "{h}", + Base64Data: "{i}", + }, + }, + }, + } + + result, err := m.Format(context.Background(), map[string]any{ + "a": "1", "b": "2", "c": "3", "d": "4", "e": "5", "f": "6", "g": "7", "h": "8", "i": "9", + }, FString) + assert.NoError(t, err) + assert.Equal(t, []*AgenticMessage{{ + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputText, + UserInputText: &UserInputText{Text: "1"}, + }, + { + Type: ContentBlockTypeUserInputImage, + UserInputImage: &UserInputImage{ + URL: "2", + Base64Data: "3", + }, + }, + { + Type: ContentBlockTypeUserInputAudio, + UserInputAudio: &UserInputAudio{ + URL: "4", + Base64Data: "5", + }, + }, + { + Type: ContentBlockTypeUserInputVideo, + UserInputVideo: &UserInputVideo{ + URL: "6", + Base64Data: "7", + }, + }, + { + Type: ContentBlockTypeUserInputFile, + UserInputFile: &UserInputFile{ + URL: "8", + Base64Data: "9", + }, + }, + }, + }}, result) +} + +func TestAgenticPlaceholderFormat(t *testing.T) { + ctx := context.Background() + ph := AgenticMessagesPlaceholder("a", false) + + result, err := ph.Format(ctx, map[string]any{ + "a": []*AgenticMessage{{Role: AgenticRoleTypeUser}, {Role: AgenticRoleTypeUser}}, + }, FString) + assert.NoError(t, err) + assert.Equal(t, 2, len(result)) + + ph = AgenticMessagesPlaceholder("a", true) + + result, err = ph.Format(ctx, map[string]any{}, FString) + assert.NoError(t, err) + assert.Equal(t, 0, len(result)) +} + +func ptrOf[T any](v T) *T { + return &v +} + +func TestAgenticMessageString(t *testing.T) { + longBase64 := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + + msg := &AgenticMessage{ + Role: AgenticRoleTypeAssistant, + ContentBlocks: []*ContentBlock{ + { + Type: ContentBlockTypeUserInputText, + UserInputText: &UserInputText{ + Text: "What's the weather like in New York City today?", + }, + }, + { + Type: ContentBlockTypeUserInputImage, + UserInputImage: &UserInputImage{ + URL: "https://example.com/weather-map.jpg", + Base64Data: longBase64, + MIMEType: "image/jpeg", + Detail: ImageURLDetailHigh, + }, + }, + { + Type: ContentBlockTypeUserInputAudio, + UserInputAudio: &UserInputAudio{ + URL: "http://audio.com", + Base64Data: "audio_data", + MIMEType: "audio/mp3", + }, + }, + { + Type: ContentBlockTypeUserInputVideo, + UserInputVideo: &UserInputVideo{ + URL: "http://video.com", + Base64Data: "video_data", + MIMEType: "video/mp4", + }, + }, + { + Type: ContentBlockTypeUserInputFile, + UserInputFile: &UserInputFile{ + URL: "http://file.com", + Name: "file.txt", + Base64Data: "file_data", + MIMEType: "text/plain", + }, + }, + { + Type: ContentBlockTypeAssistantGenText, + AssistantGenText: &AssistantGenText{ + Text: "I'll check the current weather in New York City for you.", + }, + }, + { + Type: ContentBlockTypeAssistantGenImage, + AssistantGenImage: &AssistantGenImage{ + URL: "http://gen_image.com", + Base64Data: "gen_image_data", + MIMEType: "image/png", + }, + }, + { + Type: ContentBlockTypeAssistantGenAudio, + AssistantGenAudio: &AssistantGenAudio{ + URL: "http://gen_audio.com", + Base64Data: "gen_audio_data", + MIMEType: "audio/wav", + }, + }, + { + Type: ContentBlockTypeAssistantGenVideo, + AssistantGenVideo: &AssistantGenVideo{ + URL: "http://gen_video.com", + Base64Data: "gen_video_data", + MIMEType: "video/mp4", + }, + }, + { + Type: ContentBlockTypeReasoning, + Reasoning: &Reasoning{ + Text: "First, I need to identify the location (New York City) from the user's query.\n" + + "Then, I should call the weather API to get current conditions.\n" + + "Finally, I'll format the response in a user-friendly way with temperature and conditions.", + Signature: "encrypted_reasoning_content_that_is_very_long_and_will_be_truncated_for_display", + }, + }, + { + Type: ContentBlockTypeFunctionToolCall, + FunctionToolCall: &FunctionToolCall{ + CallID: "call_weather_123", + Name: "get_current_weather", + Arguments: `{"location":"New York City","unit":"fahrenheit"}`, + }, + StreamingMeta: &StreamingMeta{Index: 0}, + }, + { + Type: ContentBlockTypeFunctionToolResult, + FunctionToolResult: &FunctionToolResult{ + CallID: "call_weather_123", + Name: "get_current_weather", + Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: `{"temperature":72,"condition":"sunny","humidity":45,"wind_speed":8}`}}, + }, + }, + }, + { + Type: ContentBlockTypeServerToolCall, + ServerToolCall: &ServerToolCall{ + Name: "server_tool", + CallID: "call_1", + Arguments: map[string]any{"a": 1}, + }, + }, + { + Type: ContentBlockTypeServerToolResult, + ServerToolResult: &ServerToolResult{ + Name: "server_tool", + CallID: "call_1", + Content: map[string]any{"success": true}, + }, + }, + { + Type: ContentBlockTypeMCPToolApprovalRequest, + MCPToolApprovalRequest: &MCPToolApprovalRequest{ + ID: "req_1", + Name: "mcp_tool", + ServerLabel: "mcp_server", + Arguments: "{}", + }, + }, + { + Type: ContentBlockTypeMCPToolApprovalResponse, + MCPToolApprovalResponse: &MCPToolApprovalResponse{ + ApprovalRequestID: "req_1", + Approve: true, + Reason: "looks good", + }, + }, + { + Type: ContentBlockTypeMCPToolCall, + MCPToolCall: &MCPToolCall{ + ServerLabel: "weather-mcp-server", + CallID: "mcp_forecast_456", + Name: "get_7day_forecast", + Arguments: `{"city":"New York","days":7}`, + }, + }, + { + Type: ContentBlockTypeMCPToolResult, + MCPToolResult: &MCPToolResult{ + CallID: "mcp_forecast_456", + Name: "get_7day_forecast", + Content: `{"status":"partial","days_available":3}`, + Error: &MCPToolCallError{ + Code: ptrOf[int64](503), + Message: "Service temporarily unavailable for full 7-day forecast", + }, + }, + }, + { + Type: ContentBlockTypeMCPListToolsResult, + MCPListToolsResult: &MCPListToolsResult{ + ServerLabel: "weather-mcp-server", + Tools: []*MCPListToolsItem{ + {Name: "get_current_weather", Description: "Get current weather conditions for a location"}, + {Name: "get_7day_forecast", Description: "Get 7-day weather forecast"}, + {Name: "get_weather_alerts", Description: "Get active weather alerts and warnings"}, + }, + }, + }, + }, + ResponseMeta: &AgenticResponseMeta{ + TokenUsage: &TokenUsage{ + PromptTokens: 250, + CompletionTokens: 180, + TotalTokens: 430, + }, + }, + } + + // Print the formatted output + output := msg.String() + + assert.Equal(t, `role: assistant +content_blocks: + [0] type: user_input_text + text: What's the weather like in New York City today? + [1] type: user_input_image + url: https://example.com/weather-map.jpg + base64_data: iVBORw0KGgoAAAANSUhE...... (96 bytes) + mime_type: image/jpeg + detail: high + [2] type: user_input_audio + url: http://audio.com + base64_data: audio_data... (10 bytes) + mime_type: audio/mp3 + [3] type: user_input_video + url: http://video.com + base64_data: video_data... (10 bytes) + mime_type: video/mp4 + [4] type: user_input_file + name: file.txt + url: http://file.com + base64_data: file_data... (9 bytes) + mime_type: text/plain + [5] type: assistant_gen_text + text: I'll check the current weather in New York City for you. + [6] type: assistant_gen_image + url: http://gen_image.com + base64_data: gen_image_data... (14 bytes) + mime_type: image/png + [7] type: assistant_gen_audio + url: http://gen_audio.com + base64_data: gen_audio_data... (14 bytes) + mime_type: audio/wav + [8] type: assistant_gen_video + url: http://gen_video.com + base64_data: gen_video_data... (14 bytes) + mime_type: video/mp4 + [9] type: reasoning + text: First, I need to identify the location (New York City) from the user's query. +Then, I should call the weather API to get current conditions. +Finally, I'll format the response in a user-friendly way with temperature and conditions. + signature: encrypted_reasoning_content_that_is_very_long_and_... + [10] type: function_tool_call + call_id: call_weather_123 + name: get_current_weather + arguments: {"location":"New York City","unit":"fahrenheit"} + stream_index: 0 + [11] type: function_tool_result + call_id: call_weather_123 + name: get_current_weather + content: (1 blocks) + [0] text: {"temperature":72,"condition":"sunny","humidity":45,"wind_speed":8} + [12] type: server_tool_call + name: server_tool + call_id: call_1 + arguments: { + "a": 1 +} + [13] type: server_tool_result + name: server_tool + call_id: call_1 + content: { + "success": true +} + [14] type: mcp_tool_approval_request + server_label: mcp_server + id: req_1 + name: mcp_tool + arguments: {} + [15] type: mcp_tool_approval_response + approval_request_id: req_1 + approve: true + reason: looks good + [16] type: mcp_tool_call + server_label: weather-mcp-server + call_id: mcp_forecast_456 + name: get_7day_forecast + arguments: {"city":"New York","days":7} + [17] type: mcp_tool_result + call_id: mcp_forecast_456 + name: get_7day_forecast + content: {"status":"partial","days_available":3} + error: [503] Service temporarily unavailable for full 7-day forecast + [18] type: mcp_list_tools_result + server_label: weather-mcp-server + tools: 3 items + - get_current_weather: Get current weather conditions for a location + - get_7day_forecast: Get 7-day weather forecast + - get_weather_alerts: Get active weather alerts and warnings +response_meta: + token_usage: prompt=250, completion=180, total=430 +`, output) + + t.Run("nil/empty fields", func(t *testing.T) { + msg := &AgenticMessage{ + Role: AgenticRoleTypeUser, + ContentBlocks: []*ContentBlock{ + {Type: ContentBlockTypeUserInputAudio, UserInputAudio: &UserInputAudio{}}, // empty + {Type: ContentBlockTypeUserInputVideo, UserInputVideo: &UserInputVideo{}}, + {Type: ContentBlockTypeUserInputFile, UserInputFile: &UserInputFile{}}, + {Type: ContentBlockTypeAssistantGenImage, AssistantGenImage: &AssistantGenImage{}}, + {Type: ContentBlockTypeAssistantGenAudio, AssistantGenAudio: &AssistantGenAudio{}}, + {Type: ContentBlockTypeAssistantGenVideo, AssistantGenVideo: &AssistantGenVideo{}}, + {Type: ContentBlockTypeServerToolCall, ServerToolCall: &ServerToolCall{Name: "t"}}, // No CallID + {Type: ContentBlockTypeServerToolResult, ServerToolResult: &ServerToolResult{Name: "t"}}, // No CallID + {Type: ContentBlockTypeMCPToolResult, MCPToolResult: &MCPToolResult{Name: "t"}}, // No Error + {Type: ContentBlockTypeMCPListToolsResult, MCPListToolsResult: &MCPListToolsResult{}}, // No Error + {Type: ContentBlockTypeMCPToolApprovalResponse, MCPToolApprovalResponse: &MCPToolApprovalResponse{Approve: false}}, // No Reason + nil, // Nil block in slice + }, + } + + s := msg.String() + assert.Contains(t, s, "type: user_input_audio") + assert.NotContains(t, s, "mime_type:") + assert.Contains(t, s, "type: server_tool_call") + }) + + t.Run("nil content struct in block", func(t *testing.T) { + // Test cases where the specific content struct is nil but type is set + // This shouldn't crash and should just print type + msg := &AgenticMessage{ + ContentBlocks: []*ContentBlock{ + {Type: ContentBlockTypeReasoning, Reasoning: nil}, + {Type: ContentBlockTypeUserInputText, UserInputText: nil}, + {Type: ContentBlockTypeUserInputImage, UserInputImage: nil}, + {Type: ContentBlockTypeUserInputAudio, UserInputAudio: nil}, + {Type: ContentBlockTypeUserInputVideo, UserInputVideo: nil}, + {Type: ContentBlockTypeUserInputFile, UserInputFile: nil}, + {Type: ContentBlockTypeAssistantGenText, AssistantGenText: nil}, + {Type: ContentBlockTypeAssistantGenImage, AssistantGenImage: nil}, + {Type: ContentBlockTypeAssistantGenAudio, AssistantGenAudio: nil}, + {Type: ContentBlockTypeAssistantGenVideo, AssistantGenVideo: nil}, + {Type: ContentBlockTypeFunctionToolCall, FunctionToolCall: nil}, + {Type: ContentBlockTypeFunctionToolResult, FunctionToolResult: nil}, + {Type: ContentBlockTypeServerToolCall, ServerToolCall: nil}, + {Type: ContentBlockTypeServerToolResult, ServerToolResult: nil}, + {Type: ContentBlockTypeMCPToolCall, MCPToolCall: nil}, + {Type: ContentBlockTypeMCPToolResult, MCPToolResult: nil}, + {Type: ContentBlockTypeMCPListToolsResult, MCPListToolsResult: nil}, + {Type: ContentBlockTypeMCPToolApprovalRequest, MCPToolApprovalRequest: nil}, + {Type: ContentBlockTypeMCPToolApprovalResponse, MCPToolApprovalResponse: nil}, + }, + } + s := msg.String() + assert.Contains(t, s, "type: reasoning") + // ensure no panic and basic output present + }) +} + +func TestSystemAgenticMessage(t *testing.T) { + t.Run("basic", func(t *testing.T) { + msg := SystemAgenticMessage("system") + assert.Equal(t, AgenticRoleTypeSystem, msg.Role) + assert.Len(t, msg.ContentBlocks, 1) + assert.Equal(t, "system", msg.ContentBlocks[0].UserInputText.Text) + }) +} + +func TestUserAgenticMessage(t *testing.T) { + t.Run("basic", func(t *testing.T) { + msg := UserAgenticMessage("user") + assert.Equal(t, AgenticRoleTypeUser, msg.Role) + assert.Len(t, msg.ContentBlocks, 1) + assert.Equal(t, "user", msg.ContentBlocks[0].UserInputText.Text) + }) +} + +func TestNewContentBlock(t *testing.T) { + cbType := reflect.TypeOf(ContentBlock{}) + for i := 0; i < cbType.NumField(); i++ { + field := cbType.Field(i) + + // Skip non-content fields + if field.Name == "Type" || field.Name == "Extra" || field.Name == "StreamingMeta" { + continue + } + + t.Run(field.Name, func(t *testing.T) { + // Ensure field is a pointer + assert.Equal(t, reflect.Ptr, field.Type.Kind(), "Field %s should be a pointer", field.Name) + + // Create a new instance of the field's type + // field.Type is *T, so Elem() is T. reflect.New(T) returns *T. + elemType := field.Type.Elem() + inputVal := reflect.New(elemType) + input := inputVal.Interface() + + // Call NewContentBlock (generic) via type switch + var block *ContentBlock + switch v := input.(type) { + case *Reasoning: + block = NewContentBlock(v) + case *UserInputText: + block = NewContentBlock(v) + case *UserInputImage: + block = NewContentBlock(v) + case *UserInputAudio: + block = NewContentBlock(v) + case *UserInputVideo: + block = NewContentBlock(v) + case *UserInputFile: + block = NewContentBlock(v) + case *ToolSearchFunctionToolResult: + block = NewContentBlock(v) + case *AssistantGenText: + block = NewContentBlock(v) + case *AssistantGenImage: + block = NewContentBlock(v) + case *AssistantGenAudio: + block = NewContentBlock(v) + case *AssistantGenVideo: + block = NewContentBlock(v) + case *FunctionToolCall: + block = NewContentBlock(v) + case *FunctionToolResult: + block = NewContentBlock(v) + case *ServerToolCall: + block = NewContentBlock(v) + case *ServerToolResult: + block = NewContentBlock(v) + case *MCPToolCall: + block = NewContentBlock(v) + case *MCPToolResult: + block = NewContentBlock(v) + case *MCPListToolsResult: + block = NewContentBlock(v) + case *MCPToolApprovalRequest: + block = NewContentBlock(v) + case *MCPToolApprovalResponse: + block = NewContentBlock(v) + default: + t.Fatalf("unsupported ContentBlock field type: %T", input) + } + + // Assertions + assert.NotNil(t, block, "NewContentBlock should return non-nil for type %T", input) + + // Check if the corresponding field in block is set equals to input + blockVal := reflect.ValueOf(block).Elem() + fieldVal := blockVal.FieldByName(field.Name) + assert.True(t, fieldVal.IsValid(), "Field %s not found in result", field.Name) + assert.Equal(t, input, fieldVal.Interface(), "Field %s should match input", field.Name) + + // Check Type is set + typeVal := blockVal.FieldByName("Type") + assert.NotEmpty(t, typeVal.String(), "Type should be set for %s", field.Name) + }) + } +} + +func TestNewContentBlockChunk_NilMeta(t *testing.T) { + require.NotPanics(t, func() { + block := NewContentBlockChunk(&AssistantGenText{Text: "test"}, nil) + require.NotNil(t, block) + assert.Nil(t, block.StreamingMeta) + }, "NewContentBlockChunk should handle nil meta without panic") +} + +func TestConcatAssistantGenTexts_ExtensionOverwrite(t *testing.T) { + type testExtension struct { + Value string + } + + texts := []*AssistantGenText{ + {Text: "Hello ", Extension: &testExtension{Value: "ext1"}}, + {Text: "world", Extension: &testExtension{Value: "ext2"}}, + } + + result, err := concatAssistantGenTexts(texts) + if err != nil { + t.Logf("Concat error (may be expected if ConcatSliceValue doesn't handle this type): %v", err) + t.Skip("Skipping: ConcatSliceValue doesn't support test type") + } + require.NotNil(t, result) + + assert.Equal(t, "Hello world", result.Text) + + if result.Extension != nil { + t.Logf("Extension type: %T, value: %v", result.Extension, result.Extension) + _, isSlice := result.Extension.([]*testExtension) + if isSlice { + t.Log("WARNING: Extension is a raw slice instead of a concatenated value. " + + "Line 1381 in agentic_message.go overwrites the ConcatSliceValue result " + + "with extensions.Interface(), discarding the concatenation.") + } + } +} + +func TestFunctionToolResultBlockString(t *testing.T) { + t.Run("empty type", func(t *testing.T) { + b := &FunctionToolResultContentBlock{Text: &UserInputText{Text: "x"}} + assert.Equal(t, "unknown block type: \n", b.String()) + }) + + t.Run("known type but empty payload", func(t *testing.T) { + b := &FunctionToolResultContentBlock{Type: FunctionToolResultContentBlockTypeText} + assert.Equal(t, "empty text block\n", b.String()) + }) + + t.Run("unknown type value", func(t *testing.T) { + b := &FunctionToolResultContentBlock{Type: FunctionToolResultContentBlockType("weird")} + assert.Equal(t, "unknown block type: weird\n", b.String()) + }) +} + +func TestConcatFunctionToolResults(t *testing.T) { + t.Run("direct append", func(t *testing.T) { + results := []*FunctionToolResult{ + {CallID: "c1", Name: "tool1", Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeText, Text: &UserInputText{Text: "hello"}}, + }}, + {CallID: "c1", Name: "tool1", Content: []*FunctionToolResultContentBlock{ + {Type: FunctionToolResultContentBlockTypeImage, Image: &UserInputImage{URL: "http://img.png"}}, + }}, + } + got, err := concatFunctionToolResults(results) + require.NoError(t, err) + assert.Len(t, got.Content, 2) + assert.Equal(t, "hello", got.Content[0].Text.Text) + assert.Equal(t, "http://img.png", got.Content[1].Image.URL) + }) +} diff --git a/schema/claude/consts.go b/schema/claude/consts.go new file mode 100644 index 0000000..714b036 --- /dev/null +++ b/schema/claude/consts.go @@ -0,0 +1,27 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package claude defines constants for claude. +package claude + +type TextCitationType string + +const ( + TextCitationTypeCharLocation TextCitationType = "char_location" + TextCitationTypePageLocation TextCitationType = "page_location" + TextCitationTypeContentBlockLocation TextCitationType = "content_block_location" + TextCitationTypeWebSearchResultLocation TextCitationType = "web_search_result_location" +) diff --git a/schema/claude/extension.go b/schema/claude/extension.go new file mode 100644 index 0000000..239941b --- /dev/null +++ b/schema/claude/extension.go @@ -0,0 +1,134 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package claude + +import ( + "fmt" +) + +type ResponseMetaExtension struct { + ID string `json:"id,omitempty"` + StopReason string `json:"stop_reason,omitempty"` + StopSequence string `json:"stop_sequence,omitempty"` + StopDetails *StopDetails `json:"stop_details,omitempty"` +} + +type StopDetails struct { + Category string `json:"category,omitempty"` + Explanation string `json:"explanation,omitempty"` +} + +type AssistantGenTextExtension struct { + Citations []*TextCitation `json:"citations,omitempty"` +} + +type TextCitation struct { + Type TextCitationType `json:"type,omitempty"` + + CharLocation *CitationCharLocation `json:"char_location,omitempty"` + PageLocation *CitationPageLocation `json:"page_location,omitempty"` + ContentBlockLocation *CitationContentBlockLocation `json:"content_block_location,omitempty"` + WebSearchResultLocation *CitationWebSearchResultLocation `json:"web_search_result_location,omitempty"` +} + +type CitationCharLocation struct { + CitedText string `json:"cited_text,omitempty"` + + DocumentTitle string `json:"document_title,omitempty"` + DocumentIndex int `json:"document_index,omitempty"` + + StartCharIndex int `json:"start_char_index,omitempty"` + EndCharIndex int `json:"end_char_index,omitempty"` +} + +type CitationPageLocation struct { + CitedText string `json:"cited_text,omitempty"` + + DocumentTitle string `json:"document_title,omitempty"` + DocumentIndex int `json:"document_index,omitempty"` + + StartPageNumber int `json:"start_page_number,omitempty"` + EndPageNumber int `json:"end_page_number,omitempty"` +} + +type CitationContentBlockLocation struct { + CitedText string `json:"cited_text,omitempty"` + + DocumentTitle string `json:"document_title,omitempty"` + DocumentIndex int `json:"document_index,omitempty"` + + StartBlockIndex int `json:"start_block_index,omitempty"` + EndBlockIndex int `json:"end_block_index,omitempty"` +} + +type CitationWebSearchResultLocation struct { + CitedText string `json:"cited_text,omitempty"` + + Title string `json:"title,omitempty"` + URL string `json:"url,omitempty"` + + EncryptedIndex string `json:"encrypted_index,omitempty"` +} + +// ConcatAssistantGenTextExtensions merges multiple AssistantGenTextExtension chunks into one. +func ConcatAssistantGenTextExtensions(chunks []*AssistantGenTextExtension) (*AssistantGenTextExtension, error) { + if len(chunks) == 0 { + return nil, fmt.Errorf("no assistant generated text extension found") + } + if len(chunks) == 1 { + return chunks[0], nil + } + + ret := &AssistantGenTextExtension{ + Citations: make([]*TextCitation, 0, len(chunks)), + } + + for _, ext := range chunks { + ret.Citations = append(ret.Citations, ext.Citations...) + } + + return ret, nil +} + +// ConcatResponseMetaExtensions merges multiple ResponseMetaExtension chunks into one. +func ConcatResponseMetaExtensions(chunks []*ResponseMetaExtension) (*ResponseMetaExtension, error) { + if len(chunks) == 0 { + return nil, fmt.Errorf("no response meta extension found") + } + if len(chunks) == 1 { + return chunks[0], nil + } + + ret := &ResponseMetaExtension{} + + for _, ext := range chunks { + if ext.ID != "" { + ret.ID = ext.ID + } + if ext.StopReason != "" { + ret.StopReason = ext.StopReason + } + if ext.StopSequence != "" { + ret.StopSequence = ext.StopSequence + } + if ext.StopDetails != nil { + ret.StopDetails = ext.StopDetails + } + } + + return ret, nil +} diff --git a/schema/claude/extension_test.go b/schema/claude/extension_test.go new file mode 100644 index 0000000..474fe74 --- /dev/null +++ b/schema/claude/extension_test.go @@ -0,0 +1,190 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package claude + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestConcatAssistantGenTextExtensions(t *testing.T) { + t.Run("multiple extensions - concatenates all citations", func(t *testing.T) { + exts := []*AssistantGenTextExtension{ + { + Citations: []*TextCitation{ + { + Type: "char_location", + CharLocation: &CitationCharLocation{ + CitedText: "citation 1", + DocumentIndex: 0, + }, + }, + }, + }, + { + Citations: []*TextCitation{ + { + Type: "page_location", + PageLocation: &CitationPageLocation{ + CitedText: "citation 2", + StartPageNumber: 1, + EndPageNumber: 2, + }, + }, + { + Type: "web_search_result_location", + WebSearchResultLocation: &CitationWebSearchResultLocation{ + CitedText: "citation 3", + URL: "https://example.com", + }, + }, + }, + }, + { + Citations: []*TextCitation{ + { + Type: "content_block_location", + ContentBlockLocation: &CitationContentBlockLocation{ + CitedText: "citation 4", + StartBlockIndex: 0, + EndBlockIndex: 5, + }, + }, + }, + }, + } + + result, err := ConcatAssistantGenTextExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Citations, 4) + assert.Equal(t, "citation 1", result.Citations[0].CharLocation.CitedText) + assert.Equal(t, "citation 2", result.Citations[1].PageLocation.CitedText) + assert.Equal(t, "citation 3", result.Citations[2].WebSearchResultLocation.CitedText) + assert.Equal(t, "citation 4", result.Citations[3].ContentBlockLocation.CitedText) + }) + + t.Run("mixed empty and non-empty citations", func(t *testing.T) { + exts := []*AssistantGenTextExtension{ + {Citations: nil}, + { + Citations: []*TextCitation{ + { + Type: "char_location", + CharLocation: &CitationCharLocation{ + CitedText: "text1", + }, + }, + }, + }, + {Citations: []*TextCitation{}}, + { + Citations: []*TextCitation{ + { + Type: "page_location", + PageLocation: &CitationPageLocation{ + CitedText: "text2", + }, + }, + }, + }, + } + + result, err := ConcatAssistantGenTextExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Citations, 2) + assert.Equal(t, "text1", result.Citations[0].CharLocation.CitedText) + assert.Equal(t, "text2", result.Citations[1].PageLocation.CitedText) + }) + + t.Run("streaming scenario - citations arrive in chunks", func(t *testing.T) { + // Simulates streaming where citations arrive progressively + exts := []*AssistantGenTextExtension{ + { + Citations: []*TextCitation{ + {Type: "char_location", CharLocation: &CitationCharLocation{CitedText: "chunk1"}}, + }, + }, + { + Citations: []*TextCitation{ + {Type: "char_location", CharLocation: &CitationCharLocation{CitedText: "chunk2"}}, + }, + }, + { + Citations: []*TextCitation{ + {Type: "char_location", CharLocation: &CitationCharLocation{CitedText: "chunk3"}}, + }, + }, + } + + result, err := ConcatAssistantGenTextExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Citations, 3) + assert.Equal(t, "chunk1", result.Citations[0].CharLocation.CitedText) + assert.Equal(t, "chunk2", result.Citations[1].CharLocation.CitedText) + assert.Equal(t, "chunk3", result.Citations[2].CharLocation.CitedText) + }) +} + +func TestConcatResponseMetaExtensions(t *testing.T) { + t.Run("multiple extensions - takes last non-empty values", func(t *testing.T) { + exts := []*ResponseMetaExtension{ + { + ID: "msg_1", + StopReason: "stop_1", + }, + { + ID: "msg_2", + StopReason: "", + }, + { + ID: "", + StopReason: "stop_3", + }, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "msg_2", result.ID) // Last non-empty ID + assert.Equal(t, "stop_3", result.StopReason) // Last non-empty StopReason + }) + + t.Run("all empty fields", func(t *testing.T) { + exts := []*ResponseMetaExtension{ + {ID: "", StopReason: ""}, + {ID: "", StopReason: ""}, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "", result.ID) + assert.Equal(t, "", result.StopReason) + }) + + t.Run("streaming scenario - ID in first chunk, StopReason in last", func(t *testing.T) { + exts := []*ResponseMetaExtension{ + {ID: "msg_stream_123", StopReason: ""}, + {ID: "", StopReason: ""}, + {ID: "", StopReason: "end_turn"}, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "msg_stream_123", result.ID) + assert.Equal(t, "end_turn", result.StopReason) + }) +} diff --git a/schema/doc.go b/schema/doc.go new file mode 100644 index 0000000..c8ba7a3 --- /dev/null +++ b/schema/doc.go @@ -0,0 +1,86 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package schema defines the core data structures and utilities shared across +// all Eino components. +// +// # Key Types +// +// [Message] is the universal unit of communication between users, models, and +// tools. It carries role, text content, multimodal media, tool calls, and +// response metadata. Helper constructors — [UserMessage], [SystemMessage], +// [AssistantMessage], [ToolMessage] — cover the most common cases. +// +// [Document] represents a piece of text with a metadata map. Typed accessors +// (Score, SubIndexes, DenseVector, SparseVector, DSLInfo, ExtraInfo) read and +// write well-known metadata keys so pipeline stages can pass structured data +// without coupling to specific struct types. +// +// [ToolInfo] describes a tool's name, description, and parameter schema. +// Parameters can be declared either as a [ParameterInfo] map (simple, struct- +// like) or as a raw [jsonschema.Schema] (full JSON Schema 2020-12 expressiveness). +// [ToolChoice] controls whether the model must, may, or must not call tools. +// +// # Streaming +// +// [StreamReader] and [StreamWriter] are the building blocks for streaming data +// through Eino pipelines. Create a linked pair with [Pipe]: +// +// sr, sw := schema.Pipe[*schema.Message](10) +// go func() { +// defer sw.Close() +// sw.Send(chunk, nil) +// }() +// defer sr.Close() +// for { +// chunk, err := sr.Recv() +// if errors.Is(err, io.EOF) { break } +// } +// +// Important constraints: +// - A StreamReader is read-once: only one goroutine may call Recv. +// - Always call Close, even when the loop ends on io.EOF, to release resources. +// - To give the same stream to multiple consumers, call [StreamReader.Copy]. +// +// # Four Streaming Paradigms +// +// Eino components and Lambda functions are classified by their input/output +// streaming shape. The framework automatically bridges mismatches: +// +// - Invoke: non-streaming in, non-streaming out (ping-pong). +// - Stream: non-streaming in, StreamReader out (server-streaming). ChatModel +// and Tool support this. +// - Collect: StreamReader in, non-streaming out (client-streaming). Useful +// for branch conditions that decide after the first chunk. +// - Transform: StreamReader in, StreamReader out (bidirectional). +// +// When an upstream node outputs T but a downstream node only accepts +// StreamReader[T], the framework wraps T in a single-chunk StreamReader — +// this is called a "fake stream". It satisfies the interface but does NOT +// reduce time-to-first-chunk. Conversely, when a downstream node only accepts +// T but the upstream outputs StreamReader[T], the framework automatically +// concatenates the stream into a complete T. +// +// Utility functions: +// - [StreamReaderFromArray] wraps a slice as a stream (useful in tests). +// - [MergeStreamReaders] fans-in multiple streams into one. +// - [MergeNamedStreamReaders] like MergeStreamReaders but emits [SourceEOF] +// when each named source ends, useful for tracking per-source completion. +// - [StreamReaderWithConvert] transforms element types; return [ErrNoValue] +// from the convert function to skip an element. +// +// See https://www.cloudwego.io/docs/eino/core_modules/chain_and_graph_orchestration/stream_programming_essentials/ +package schema diff --git a/schema/document.go b/schema/document.go new file mode 100644 index 0000000..09c19d5 --- /dev/null +++ b/schema/document.go @@ -0,0 +1,220 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +const ( + docMetaDataKeySubIndexes = "_sub_indexes" + docMetaDataKeyScore = "_score" + docMetaDataKeyExtraInfo = "_extra_info" + docMetaDataKeyDSL = "_dsl" + docMetaDataKeyDenseVector = "_dense_vector" + docMetaDataKeySparseVector = "_sparse_vector" +) + +// Document is a piece of text with a metadata map. It is the shared currency +// between Loader, Transformer, Indexer, and Retriever components. +// +// Metadata is an open map[string]any that lets pipeline stages attach typed +// values to a document without creating a new struct. Well-known keys are +// managed through typed accessor methods — Score, SubIndexes, DenseVector, +// SparseVector, DSLInfo, ExtraInfo — so callers never need to reference the +// raw key strings. +// +// Transformer implementations should preserve existing metadata and merge new +// keys rather than replacing the map outright, so provenance information +// accumulated by earlier stages is not lost. +type Document struct { + // ID is the unique identifier of the document. + ID string `json:"id"` + // Content is the content of the document. + Content string `json:"content"` + // MetaData is the metadata of the document, can be used to store extra information. + MetaData map[string]any `json:"meta_data"` +} + +// String returns the content of the document. +func (d *Document) String() string { + return d.Content +} + +// WithSubIndexes sets the sub-indexes on the document metadata and returns the +// document for chaining. Sub-indexes let an Indexer route a document into +// multiple logical partitions of a vector store simultaneously. +// Use [Document.SubIndexes] to retrieve them. +func (d *Document) WithSubIndexes(indexes []string) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + + d.MetaData[docMetaDataKeySubIndexes] = indexes + + return d +} + +// SubIndexes returns the sub indexes of the document. +// can use doc.WithSubIndexes() to set the sub indexes. +func (d *Document) SubIndexes() []string { + if d.MetaData == nil { + return nil + } + + indexes, ok := d.MetaData[docMetaDataKeySubIndexes].([]string) + if ok { + return indexes + } + + return nil +} + +// WithScore sets the relevance score on the document, typically written by a +// Retriever after ranking results. A higher score means higher relevance. +// Note: [retriever.WithScoreThreshold] filters by this value, not sort order. +// Use [Document.Score] to retrieve it. +func (d *Document) WithScore(score float64) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + + d.MetaData[docMetaDataKeyScore] = score + + return d +} + +// Score returns the score of the document. +// can use doc.WithScore() to set the score. +func (d *Document) Score() float64 { + if d.MetaData == nil { + return 0 + } + + score, ok := d.MetaData[docMetaDataKeyScore].(float64) + if ok { + return score + } + + return 0 +} + +// WithExtraInfo sets the extra info of the document. +// can use doc.ExtraInfo() to get the extra info. +func (d *Document) WithExtraInfo(extraInfo string) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + + d.MetaData[docMetaDataKeyExtraInfo] = extraInfo + + return d +} + +// ExtraInfo returns the extra info of the document. +// can use doc.WithExtraInfo() to set the extra info. +func (d *Document) ExtraInfo() string { + if d.MetaData == nil { + return "" + } + + extraInfo, ok := d.MetaData[docMetaDataKeyExtraInfo].(string) + if ok { + return extraInfo + } + + return "" +} + +// WithDSLInfo attaches a domain-specific-language query description to the +// document. This is consumed by Retriever implementations that support +// structured queries (e.g., filter expressions) alongside vector search. +// Use [Document.DSLInfo] to retrieve it. +func (d *Document) WithDSLInfo(dslInfo map[string]any) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + + d.MetaData[docMetaDataKeyDSL] = dslInfo + + return d +} + +// DSLInfo returns the dsl info of the document. +// can use doc.WithDSLInfo() to set the dsl info. +func (d *Document) DSLInfo() map[string]any { + if d.MetaData == nil { + return nil + } + + dslInfo, ok := d.MetaData[docMetaDataKeyDSL].(map[string]any) + if ok { + return dslInfo + } + + return nil +} + +// WithDenseVector sets the dense vector of the document. +// can use doc.DenseVector() to get the dense vector. +func (d *Document) WithDenseVector(vector []float64) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + + d.MetaData[docMetaDataKeyDenseVector] = vector + + return d +} + +// DenseVector returns the dense vector of the document. +// can use doc.WithDenseVector() to set the dense vector. +func (d *Document) DenseVector() []float64 { + if d.MetaData == nil { + return nil + } + + vector, ok := d.MetaData[docMetaDataKeyDenseVector].([]float64) + if ok { + return vector + } + + return nil +} + +// WithSparseVector sets the sparse vector of the document, key indices -> value vector. +// can use doc.SparseVector() to get the sparse vector. +func (d *Document) WithSparseVector(sparse map[int]float64) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + + d.MetaData[docMetaDataKeySparseVector] = sparse + + return d +} + +// SparseVector returns the sparse vector of the document, key indices -> value vector. +// can use doc.WithSparseVector() to set the sparse vector. +func (d *Document) SparseVector() map[int]float64 { + if d.MetaData == nil { + return nil + } + + sparse, ok := d.MetaData[docMetaDataKeySparseVector].(map[int]float64) + if ok { + return sparse + } + + return nil +} diff --git a/schema/document_test.go b/schema/document_test.go new file mode 100644 index 0000000..db2ba37 --- /dev/null +++ b/schema/document_test.go @@ -0,0 +1,53 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "testing" + + "github.com/smartystreets/goconvey/convey" +) + +func TestDocument(t *testing.T) { + convey.Convey("test document", t, func() { + var ( + subIndexes = []string{"hello", "bye"} + score = 1.1 + extraInfo = "asd" + dslInfo = map[string]any{"hello": true} + vector = []float64{1.1, 2.2} + ) + + d := &Document{ + ID: "asd", + Content: "qwe", + MetaData: nil, + } + + d.WithSubIndexes(subIndexes). + WithDenseVector(vector). + WithScore(score). + WithExtraInfo(extraInfo). + WithDSLInfo(dslInfo) + + convey.So(d.SubIndexes(), convey.ShouldEqual, subIndexes) + convey.So(d.Score(), convey.ShouldEqual, score) + convey.So(d.ExtraInfo(), convey.ShouldEqual, extraInfo) + convey.So(d.DSLInfo(), convey.ShouldEqual, dslInfo) + convey.So(d.DenseVector(), convey.ShouldEqual, vector) + }) +} diff --git a/schema/gemini/extension.go b/schema/gemini/extension.go new file mode 100644 index 0000000..efbc4f4 --- /dev/null +++ b/schema/gemini/extension.go @@ -0,0 +1,115 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package gemini defines the extension for gemini. +package gemini + +import ( + "fmt" +) + +type ResponseMetaExtension struct { + ID string `json:"id,omitempty"` + FinishReason string `json:"finish_reason,omitempty"` + GroundingMeta *GroundingMetadata `json:"grounding_meta,omitempty"` +} + +type GroundingMetadata struct { + // List of supporting references retrieved from specified grounding source. + GroundingChunks []*GroundingChunk `json:"grounding_chunks,omitempty"` + // Optional. List of grounding support. + GroundingSupports []*GroundingSupport `json:"grounding_supports,omitempty"` + // Optional. Google search entry for the following-up web searches. + SearchEntryPoint *SearchEntryPoint `json:"search_entry_point,omitempty"` + // Optional. Web search queries for the following-up web search. + WebSearchQueries []string `json:"web_search_queries,omitempty"` +} + +type GroundingChunk struct { + // Grounding chunk from the web. + Web *GroundingChunkWeb `json:"web,omitempty"` +} + +// GroundingChunkWeb is the chunk from the web. +type GroundingChunkWeb struct { + // Domain of the (original) URI. This field is not supported in Gemini API. + Domain string `json:"domain,omitempty"` + // Title of the chunk. + Title string `json:"title,omitempty"` + // URI reference of the chunk. + URI string `json:"uri,omitempty"` +} + +type GroundingSupport struct { + // Confidence score of the support references. Ranges from 0 to 1. 1 is the most confident. + // For Gemini 2.0 and before, this list must have the same size as the grounding_chunk_indices. + // For Gemini 2.5 and after, this list will be empty and should be ignored. + ConfidenceScores []float32 `json:"confidence_scores,omitempty"` + // A list of indices (into 'grounding_chunk') specifying the citations associated with + // the claim. For instance [1,3,4] means that grounding_chunk[1], grounding_chunk[3], + // grounding_chunk[4] are the retrieved content attributed to the claim. + GroundingChunkIndices []int `json:"grounding_chunk_indices,omitempty"` + // Segment of the content this support belongs to. + Segment *Segment `json:"segment,omitempty"` +} + +// Segment of the content. +type Segment struct { + // Output only. End index in the given Part, measured in bytes. Offset from the start + // of the Part, exclusive, starting at zero. + EndIndex int `json:"end_index,omitempty"` + // Output only. The index of a Part object within its parent Content object. + PartIndex int `json:"part_index,omitempty"` + // Output only. Start index in the given Part, measured in bytes. Offset from the start + // of the Part, inclusive, starting at zero. + StartIndex int `json:"start_index,omitempty"` + // Output only. The text corresponding to the segment from the response. + Text string `json:"text,omitempty"` +} + +// SearchEntryPoint is the Google search entry point. +type SearchEntryPoint struct { + // Optional. Web content snippet that can be embedded in a web page or an app webview. + RenderedContent string `json:"rendered_content,omitempty"` + // Optional. Base64 encoded JSON representing array of tuple. + SDKBlob []byte `json:"sdk_blob,omitempty"` +} + +// ConcatResponseMetaExtensions concatenates multiple ResponseMetaExtension chunks into a single one. +func ConcatResponseMetaExtensions(chunks []*ResponseMetaExtension) (*ResponseMetaExtension, error) { + if len(chunks) == 0 { + return nil, fmt.Errorf("no response meta extension found") + } + if len(chunks) == 1 { + return chunks[0], nil + } + + ret := &ResponseMetaExtension{} + + for _, ext := range chunks { + if ext.ID != "" { + ret.ID = ext.ID + } + if ext.FinishReason != "" { + ret.FinishReason = ext.FinishReason + } + if ext.GroundingMeta != nil { + ret.GroundingMeta = ext.GroundingMeta + } + } + + return ret, nil +} diff --git a/schema/gemini/extension_test.go b/schema/gemini/extension_test.go new file mode 100644 index 0000000..56f390a --- /dev/null +++ b/schema/gemini/extension_test.go @@ -0,0 +1,79 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package gemini + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestConcatResponseMetaExtensions(t *testing.T) { + t.Run("multiple extensions - takes last non-empty values", func(t *testing.T) { + meta1 := &GroundingMetadata{WebSearchQueries: []string{"query1"}} + meta2 := &GroundingMetadata{WebSearchQueries: []string{"query2"}} + + exts := []*ResponseMetaExtension{ + { + ID: "resp_1", + FinishReason: "STOP", + GroundingMeta: meta1, + }, + { + ID: "resp_2", + FinishReason: "", + GroundingMeta: nil, + }, + { + ID: "", + FinishReason: "MAX_TOKENS", + GroundingMeta: meta2, + }, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "resp_2", result.ID) + assert.Equal(t, "MAX_TOKENS", result.FinishReason) + assert.Equal(t, meta2, result.GroundingMeta) + }) + + t.Run("streaming scenario", func(t *testing.T) { + meta := &GroundingMetadata{ + GroundingChunks: []*GroundingChunk{ + { + Web: &GroundingChunkWeb{ + Title: "Example", + URI: "https://example.com", + }, + }, + }, + } + + exts := []*ResponseMetaExtension{ + {ID: "stream_123", FinishReason: "", GroundingMeta: nil}, + {ID: "", FinishReason: "", GroundingMeta: nil}, + {ID: "", FinishReason: "STOP", GroundingMeta: meta}, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "stream_123", result.ID) + assert.Equal(t, "STOP", result.FinishReason) + assert.Equal(t, meta, result.GroundingMeta) + }) +} diff --git a/schema/message.go b/schema/message.go new file mode 100644 index 0000000..892abe7 --- /dev/null +++ b/schema/message.go @@ -0,0 +1,1937 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "context" + "fmt" + "io" + "reflect" + "sort" + "strings" + "sync" + "text/template" + + "github.com/nikolalohinski/gonja" + "github.com/nikolalohinski/gonja/config" + "github.com/nikolalohinski/gonja/exec" + "github.com/nikolalohinski/gonja/nodes" + "github.com/nikolalohinski/gonja/parser" + "github.com/slongfield/pyfmt" + + "github.com/cloudwego/eino/internal" + "github.com/cloudwego/eino/internal/generic" +) + +func init() { + internal.RegisterStreamChunkConcatFunc(ConcatMessages) + internal.RegisterStreamChunkConcatFunc(ConcatMessageArray) + + internal.RegisterStreamChunkConcatFunc(ConcatAgenticMessages) + internal.RegisterStreamChunkConcatFunc(ConcatAgenticMessagesArray) + + internal.RegisterStreamChunkConcatFunc(ConcatToolResults) +} + +func buildConcatGenericArray[T any](f func([]*T) (*T, error)) func([][]*T) ([]*T, error) { + return func(mas [][]*T) ([]*T, error) { + arrayLen := len(mas[0]) + + ret := make([]*T, arrayLen) + slicesToConcat := make([][]*T, arrayLen) + + for _, ma := range mas { + if len(ma) != arrayLen { + return nil, fmt.Errorf("unexpected array length. "+ + "Got %d, expected %d", len(ma), arrayLen) + } + + for i := 0; i < arrayLen; i++ { + m := ma[i] + if m != nil { + slicesToConcat[i] = append(slicesToConcat[i], m) + } + } + } + + for i, slice := range slicesToConcat { + if len(slice) == 0 { + ret[i] = nil + } else if len(slice) == 1 { + ret[i] = slice[0] + } else { + cm, err := f(slice) + if err != nil { + return nil, err + } + + ret[i] = cm + } + } + + return ret, nil + } +} + +// ConcatMessageArray merges aligned slices of messages into a single slice, +// concatenating messages at the same index across the input arrays. +func ConcatMessageArray(mas [][]*Message) ([]*Message, error) { + return buildConcatGenericArray[Message](ConcatMessages)(mas) +} + +// FormatType used by MessageTemplate.Format +type FormatType uint8 + +const ( + // FString Supported by pyfmt(github.com/slongfield/pyfmt), which is an implementation of https://peps.python.org/pep-3101/. + FString FormatType = 0 + // GoTemplate https://pkg.go.dev/text/template. + GoTemplate FormatType = 1 + // Jinja2 Supported by gonja(github.com/nikolalohinski/gonja), which is a implementation of https://jinja.palletsprojects.com/en/3.1.x/templates/. + Jinja2 FormatType = 2 +) + +// RoleType is the type of the role of a message. +type RoleType string + +const ( + // Assistant is the role of an assistant, means the message is returned by ChatModel. + Assistant RoleType = "assistant" + // User is the role of a user, means the message is a user message. + User RoleType = "user" + // System is the role of a system, means the message is a system message. + System RoleType = "system" + // Tool is the role of a tool, means the message is a tool call output. + Tool RoleType = "tool" +) + +// FunctionCall is the function call in a message. +// It's used in Assistant Message. +type FunctionCall struct { + // Name is the name of the function to call, it can be used to identify the specific function. + Name string `json:"name,omitempty"` + // Arguments is the arguments to call the function with, in JSON format. + Arguments string `json:"arguments,omitempty"` +} + +// ToolCall is the tool call in a message. +// It's used in Assistant Message when there are tool calls should be made. +type ToolCall struct { + // Index is used when there are multiple tool calls in a message. + // In stream mode, it's used to identify the chunk of the tool call for merging. + Index *int `json:"index,omitempty"` + // ID is the id of the tool call, it can be used to identify the specific tool call. + ID string `json:"id"` + // Type is the type of the tool call, default is "function". + Type string `json:"type"` + // Function is the function call to be made. + Function FunctionCall `json:"function"` + // Extra is used to store extra information for the tool call. + Extra map[string]any `json:"extra,omitempty"` +} + +// ImageURLDetail is the detail of the image url. +type ImageURLDetail string + +const ( + // ImageURLDetailHigh means the high quality image url. + ImageURLDetailHigh ImageURLDetail = "high" + // ImageURLDetailLow means the low quality image url. + ImageURLDetailLow ImageURLDetail = "low" + // ImageURLDetailAuto means the auto quality image url. + ImageURLDetailAuto ImageURLDetail = "auto" +) + +// MessagePartCommon represents the common abstract components for input and output of multi-modal types. +type MessagePartCommon struct { + // URL is primarily used for HTTP or HTTPS access links. + // For data in the format 'data:[][;base64],' (the 'data' URL Schema of RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397)), + // it is recommended to use Base64Data and MIMEType fields separately instead. + URL *string `json:"url,omitempty"` + + // Base64Data represents the binary data in Base64 encoded string format. + Base64Data *string `json:"base64data,omitempty"` + + // MIMEType is the mime type , eg."image/png",""audio/wav" etc. + MIMEType string `json:"mime_type,omitempty"` + + // Deprecated: Use MessageOutputPart.Extra or MessageInputPart.Extra to set additional metadata instead. + Extra map[string]any `json:"extra,omitempty"` +} + +// MessageInputImage is used to represent an image part in message. +// Choose either URL or Base64Data. +type MessageInputImage struct { + MessagePartCommon + + // Detail is the quality of the image url. + Detail ImageURLDetail `json:"detail,omitempty"` +} + +// MessageInputAudio is used to represent an audio part in message. +// Choose either URL or Base64Data. +type MessageInputAudio struct { + MessagePartCommon +} + +// MessageInputVideo is used to represent a video part in message. +// Choose either URL or Base64Data. +type MessageInputVideo struct { + MessagePartCommon +} + +// MessageInputFile is used to represent a file part in message. +// Choose either URL or Base64Data. +type MessageInputFile struct { + MessagePartCommon + + // Name represents the filename. + // Optional. + Name string `json:"name,omitempty"` +} + +// MessageInputPart represents the input part of message. +type MessageInputPart struct { + Type ChatMessagePartType `json:"type"` + + Text string `json:"text,omitempty"` + + // Image is the image input of the part, it's used when Type is "image_url". + Image *MessageInputImage `json:"image,omitempty"` + + // Audio is the audio input of the part, it's used when Type is "audio_url". + Audio *MessageInputAudio `json:"audio,omitempty"` + + // Video is the video input of the part, it's used when Type is "video_url". + Video *MessageInputVideo `json:"video,omitempty"` + + // File is the file input of the part, it's used when Type is "file_url". + File *MessageInputFile `json:"file,omitempty"` + + // ToolSearchResult holds the result of a tool search request, containing the matched tool names and their definitions. + ToolSearchResult *ToolSearchResult `json:"tool_search_result,omitempty"` + + // Extra is used to store extra information. + Extra map[string]any `json:"extra,omitempty"` +} + +// MessageOutputImage is used to represent an image part in message. +type MessageOutputImage struct { + MessagePartCommon +} + +// MessageOutputAudio is used to represent an audio part in message. +type MessageOutputAudio struct { + MessagePartCommon +} + +// MessageOutputVideo is used to represent a video part in message. +type MessageOutputVideo struct { + MessagePartCommon +} + +// MessageOutputReasoning represents the reasoning content generated by reasoning models. +// Some models produce reasoning steps before generating the final response. +// This struct captures that reasoning output. +type MessageOutputReasoning struct { + // Text is either the thought summary or the raw reasoning text itself. + Text string `json:"text,omitempty"` + + // Signature contains encrypted reasoning tokens. + // Required by some models when passing reasoning context back in subsequent requests. + Signature string `json:"signature,omitempty"` +} + +// MessageStreamingMeta contains metadata for streaming responses. +// It is used to track position of part when the model outputs multiple parts in a single response. +type MessageStreamingMeta struct { + // Index specifies the index position of this part in the final response. + // This is useful for reassembling multiple reasoning/content parts in correct order. + Index int `json:"index,omitempty"` +} + +// MessageOutputPart represents a part of an assistant-generated message. +// It can contain text, or multimedia content like images, audio, or video. +type MessageOutputPart struct { + // Type is the type of the part, e.g. "text", "image_url", "audio_url", "video_url". + Type ChatMessagePartType `json:"type"` + + // Text is the text of the part, it's used when Type is "text". + Text string `json:"text,omitempty"` + + // Image is the image output of the part, used when Type is ChatMessagePartTypeImageURL. + Image *MessageOutputImage `json:"image,omitempty"` + + // Audio is the audio output of the part, used when Type is ChatMessagePartTypeAudioURL. + Audio *MessageOutputAudio `json:"audio,omitempty"` + + // Video is the video output of the part, used when Type is ChatMessagePartTypeVideoURL. + Video *MessageOutputVideo `json:"video,omitempty"` + + // Reasoning contains the reasoning content generated by the model. + // Used when Type is ChatMessagePartTypeReasoning. + Reasoning *MessageOutputReasoning `json:"reasoning,omitempty"` + + // Extra is used to store extra information. + Extra map[string]any `json:"extra,omitempty"` + + // StreamingMeta contains metadata for streaming responses. + // This field is typically used at runtime and not serialized. + StreamingMeta *MessageStreamingMeta `json:"-"` +} + +// Deprecated: This struct is deprecated as the MultiContent field is deprecated. +// For the image input part of the model, use MessageInputImage. +// For the image output part of the model, use MessageOutputImage. +// Choose either URL or URI. +// If your model implementation supports it, URL could embed inline image data +// as defined in RFC-2397. +type ChatMessageImageURL struct { + // URL can either be a traditional URL or a special URL conforming to RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397). + // double check with model implementations for detailed instructions on how to use this. + URL string `json:"url,omitempty"` + + URI string `json:"uri,omitempty"` + // Detail is the quality of the image url. + Detail ImageURLDetail `json:"detail,omitempty"` + + // MIMEType is the mime type of the image, eg. "image/png". + MIMEType string `json:"mime_type,omitempty"` + // Extra is used to store extra information for the image url. + Extra map[string]any `json:"extra,omitempty"` +} + +// ChatMessagePartType is the type of the part in a chat message. +type ChatMessagePartType string + +const ( + // ChatMessagePartTypeText means the part is a text. + ChatMessagePartTypeText ChatMessagePartType = "text" + // ChatMessagePartTypeImageURL means the part is an image url. + ChatMessagePartTypeImageURL ChatMessagePartType = "image_url" + // ChatMessagePartTypeAudioURL means the part is an audio url. + ChatMessagePartTypeAudioURL ChatMessagePartType = "audio_url" + // ChatMessagePartTypeVideoURL means the part is a video url. + ChatMessagePartTypeVideoURL ChatMessagePartType = "video_url" + // ChatMessagePartTypeFileURL means the part is a file url. + ChatMessagePartTypeFileURL ChatMessagePartType = "file_url" + // ChatMessagePartTypeReasoning means the part is a reasoning block. + ChatMessagePartTypeReasoning ChatMessagePartType = "reasoning" + + // ChatMessagePartTypeToolSearchResult means the part contains tool search results. + ChatMessagePartTypeToolSearchResult ChatMessagePartType = "tool_search_result" +) + +// Deprecated: This struct is deprecated as the MultiContent field is deprecated. +// For the audio input part of the model, use MessageInputAudio. +// For the audio output part of the model, use MessageOutputAudio. +// Choose either URL or URI. +// If supported, URL may embed inline audio data per RFC-2397. +type ChatMessageAudioURL struct { + // URL can either be a traditional URL or a special URL conforming to RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397). + // double check with model implementations for detailed instructions on how to use this. + URL string `json:"url,omitempty"` + URI string `json:"uri,omitempty"` + + // MIMEType is the mime type of the audio, eg. "audio/wav" or "audio/ogg". + MIMEType string `json:"mime_type,omitempty"` + // Extra is used to store extra information for the audio url. + Extra map[string]any `json:"extra,omitempty"` +} + +// Deprecated: This struct is deprecated as the MultiContent field is deprecated. +// For the video input part of the model, use MessageInputVideo. +// For the video output part of the model, use MessageOutputVideo. +// Choose either URL or URI. +// If supported, URL may embed inline video data per RFC-2397. +type ChatMessageVideoURL struct { + // URL can either be a traditional URL or a special URL conforming to RFC-2397 (https://www.rfc-editor.org/rfc/rfc2397). + // double check with model implementations for detailed instructions on how to use this. + URL string `json:"url,omitempty"` + URI string `json:"uri,omitempty"` + + // MIMEType is the mime type of the video, eg. "video/mp4". + MIMEType string `json:"mime_type,omitempty"` + // Extra is used to store extra information for the video url. + Extra map[string]any `json:"extra,omitempty"` +} + +// Deprecated: This struct is deprecated as the MultiContent field is deprecated. +// For the file input part of the model, use MessageInputFile. +// Choose either URL or URI. +type ChatMessageFileURL struct { + URL string `json:"url,omitempty"` + URI string `json:"uri,omitempty"` + + // MIMEType is the mime type of the file, eg. "application/pdf", "text/plain". + MIMEType string `json:"mime_type,omitempty"` + // Name is the name of the file. + Name string `json:"name,omitempty"` + + // Extra is used to store extra information for the file url. + Extra map[string]any `json:"extra,omitempty"` +} + +// Deprecated: This struct is deprecated as the MultiContent field is deprecated. +// For model input, use MessageInputPart. For model output, use MessageOutputPart. +type ChatMessagePart struct { + // Type is the type of the part, eg. "text", "image_url", "audio_url", "video_url", "file_url". + Type ChatMessagePartType `json:"type,omitempty"` + + // Text is the text of the part, it's used when Type is "text". + Text string `json:"text,omitempty"` + + // ImageURL is the image url of the part, it's used when Type is "image_url". + ImageURL *ChatMessageImageURL `json:"image_url,omitempty"` + // AudioURL is the audio url of the part, it's used when Type is "audio_url". + AudioURL *ChatMessageAudioURL `json:"audio_url,omitempty"` + // VideoURL is the video url of the part, it's used when Type is "video_url". + VideoURL *ChatMessageVideoURL `json:"video_url,omitempty"` + // FileURL is the file url of the part, it's used when Type is "file_url". + FileURL *ChatMessageFileURL `json:"file_url,omitempty"` +} + +// LogProbs is the top-level structure containing the log probability information. +type LogProbs struct { + // Content is a list of message content tokens with log probability information. + Content []LogProb `json:"content"` +} + +// LogProb represents the probability information for a token. +type LogProb struct { + // Token represents the text of the token, which is a contiguous sequence of characters + // (e.g., a word, part of a word, or punctuation) as understood by the tokenization process used by the language model. + Token string `json:"token"` + // LogProb is the log probability of this token, if it is within the top 20 most likely tokens. + // Otherwise, the value `-9999.0` is used to signify that the token is very unlikely. + LogProb float64 `json:"logprob"` + // Bytes is a list of integers representing the UTF-8 bytes representation of the token. + // Useful in instances where characters are represented by multiple tokens and + // their byte representations must be combined to generate the correct text + // representation. Can be `null` if there is no bytes representation for the token. + Bytes []int64 `json:"bytes,omitempty"` // Omitting the field if it is null + // TopLogProbs is a list of the most likely tokens and their log probability, at this token position. + // In rare cases, there may be fewer than the number of requested top_logprobs returned. + TopLogProbs []TopLogProb `json:"top_logprobs"` +} + +// TopLogProb describes a likely token and its log probability at a position. +type TopLogProb struct { + // Token represents the text of the token, which is a contiguous sequence of characters + // (e.g., a word, part of a word, or punctuation) as understood by the tokenization process used by the language model. + Token string `json:"token"` + // LogProb is the log probability of this token, if it is within the top 20 most likely tokens. + // Otherwise, the value `-9999.0` is used to signify that the token is very unlikely. + LogProb float64 `json:"logprob"` + // Bytes is a list of integers representing the UTF-8 bytes representation of the token. + // Useful in instances where characters are represented by multiple tokens and + // their byte representations must be combined to generate the correct text + // representation. Can be `null` if there is no bytes representation for the token. + Bytes []int64 `json:"bytes,omitempty"` +} + +// ResponseMeta collects meta information about a chat response. +type ResponseMeta struct { + // FinishReason is the reason why the chat response is finished. + // It's usually "stop", "length", "tool_calls", "content_filter", "null". This is defined by chat model implementation. + FinishReason string `json:"finish_reason,omitempty"` + // Usage is the token usage of the chat response, whether usage exists depends on whether the chat model implementation returns. + Usage *TokenUsage `json:"usage,omitempty"` + // LogProbs is Log probability information. + LogProbs *LogProbs `json:"logprobs,omitempty"` +} + +// Message denotes the data structure for model input and output, originating from either user input or model return. +// It supports both text-only and multimodal content. +// +// For text-only input from a user, use the Content field: +// +// &schema.Message{ +// Role: schema.User, +// Content: "What is the capital of France?", +// } +// +// For multimodal input from a user, use the UserInputMultiContent field. +// This allows combining text with other media like images: +// +// &schema.Message{ +// Role: schema.User, +// UserInputMultiContent: []schema.MessageInputPart{ +// {Type: schema.ChatMessagePartTypeText, Text: "What is in this image?"}, +// {Type: schema.ChatMessagePartTypeImageURL, Image: &schema.MessageInputImage{ +// MessagePartCommon: schema.MessagePartCommon{ +// URL: toPtr("https://example.com/cat.jpg"), +// }, +// Detail: schema.ImageURLDetailHigh, +// }}, +// }, +// } +// +// When the model returns multimodal content, it is available in the AssistantGenMultiContent field: +// +// &schema.Message{ +// Role: schema.Assistant, +// AssistantGenMultiContent: []schema.MessageOutputPart{ +// {Type: schema.ChatMessagePartTypeText, Text: "Here is the generated image:"}, +// {Type: schema.ChatMessagePartTypeImage, Image: &schema.MessageOutputImage{ +// MessagePartCommon: schema.MessagePartCommon{ +// Base64Data: toPtr("base64_image_binary"), +// MIMEType: "image/png", +// }, +// }}, +// }, +// } +type Message struct { + Role RoleType `json:"role"` + + // Content is for user text input and model text output. + Content string `json:"content"` + + // if MultiContent is not empty, use this instead of Content + // if MultiContent is empty, use Content + // Deprecated: Use UserInputMultiContent for user multimodal inputs and AssistantGenMultiContent for model multimodal outputs. + MultiContent []ChatMessagePart `json:"multi_content,omitempty"` + + // UserInputMultiContent passes multimodal content provided by the user to the model. + UserInputMultiContent []MessageInputPart `json:"user_input_multi_content,omitempty"` + + // AssistantGenMultiContent is for receiving multimodal output from the model. + AssistantGenMultiContent []MessageOutputPart `json:"assistant_output_multi_content,omitempty"` + + Name string `json:"name,omitempty"` + + // only for AssistantMessage + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + + // only for ToolMessage + ToolCallID string `json:"tool_call_id,omitempty"` + // only for ToolMessage + ToolName string `json:"tool_name,omitempty"` + + ResponseMeta *ResponseMeta `json:"response_meta,omitempty"` + + // ReasoningContent is the thinking process of the model, which will be included when the model returns reasoning content. + ReasoningContent string `json:"reasoning_content,omitempty"` + + // customized information for model implementation + Extra map[string]any `json:"extra,omitempty"` +} + +// TokenUsage Represents the token usage of chat model request. +type TokenUsage struct { + // PromptTokens is the number of prompt tokens, including all the input tokens of this request. + PromptTokens int `json:"prompt_tokens"` + // PromptTokenDetails is a breakdown of the prompt tokens. + PromptTokenDetails PromptTokenDetails `json:"prompt_token_details"` + // CompletionTokens is the number of completion tokens. + CompletionTokens int `json:"completion_tokens"` + // TotalTokens is the total number of tokens. + TotalTokens int `json:"total_tokens"` + // CompletionTokensDetails is breakdown of completion tokens. + CompletionTokensDetails CompletionTokensDetails `json:"completion_token_details"` +} + +type CompletionTokensDetails struct { + // ReasoningTokens tokens generated by the model for reasoning. + // This is currently supported by OpenAI, Gemini, ARK and Qwen chat models. + // For other models, this field will be 0. + ReasoningTokens int `json:"reasoning_tokens,omitempty"` +} + +// PromptTokenDetails provides a breakdown of prompt token usage. +type PromptTokenDetails struct { + // Cached tokens present in the prompt. + CachedTokens int `json:"cached_tokens"` +} + +var _ MessagesTemplate = &Message{} +var _ MessagesTemplate = MessagesPlaceholder("", false) + +// MessagesTemplate is the interface for messages template. +// It's used to render a template to a list of messages. +// e.g. +// +// chatTemplate := prompt.FromMessages( +// schema.SystemMessage("you are an eino helper"), +// schema.MessagesPlaceholder("history", false), // <= this will use the value of "history" in params +// ) +// msgs, err := chatTemplate.Format(ctx, params) +type MessagesTemplate interface { + Format(ctx context.Context, vs map[string]any, formatType FormatType) ([]*Message, error) +} + +type messagesPlaceholder struct { + key string + optional bool +} + +// MessagesPlaceholder can render a placeholder to a list of messages in params. +// e.g. +// +// placeholder := MessagesPlaceholder("history", false) +// params := map[string]any{ +// "history": []*schema.Message{{Role: "user", Content: "what is eino?"}, {Role: "assistant", Content: "eino is a great framework to build llm apps"}}, +// "query": "how to use eino?", +// } +// chatTemplate := chatTpl := prompt.FromMessages( +// schema.SystemMessage("you are eino helper"), +// schema.MessagesPlaceholder("history", false), // <= this will use the value of "history" in params +// ) +// msgs, err := chatTemplate.Format(ctx, params) +func MessagesPlaceholder(key string, optional bool) MessagesTemplate { + return &messagesPlaceholder{ + key: key, + optional: optional, + } +} + +// Format just return the messages of specified key. +// because it's a placeholder. +// e.g. +// +// placeholder := MessagesPlaceholder("history", false) +// params := map[string]any{ +// "history": []*schema.Message{{Role: "user", Content: "what is eino?"}, {Role: "assistant", Content: "eino is a great freamwork to build llm apps"}}, +// "query": "how to use eino?", +// } +// msgs, err := placeholder.Format(ctx, params) // <= this will return the value of "history" in params +func (p *messagesPlaceholder) Format(_ context.Context, vs map[string]any, _ FormatType) ([]*Message, error) { + v, ok := vs[p.key] + if !ok { + if p.optional { + return []*Message{}, nil + } + + return nil, fmt.Errorf("message placeholder format: %s not found", p.key) + } + + msgs, ok := v.([]*Message) + if !ok { + return nil, fmt.Errorf("only messages can be used to format message placeholder, key: %v, actual type: %v", p.key, reflect.TypeOf(v)) + } + + return msgs, nil +} + +func formatContent(content string, vs map[string]any, formatType FormatType) (string, error) { + switch formatType { + case FString: + return pyfmt.Fmt(content, vs) + case GoTemplate: + parsedTmpl, err := template.New("template"). + Option("missingkey=error"). + Parse(content) + if err != nil { + return "", err + } + sb := new(strings.Builder) + err = parsedTmpl.Execute(sb, vs) + if err != nil { + return "", err + } + return sb.String(), nil + case Jinja2: + env, err := getJinjaEnv() + if err != nil { + return "", err + } + tpl, err := env.FromString(content) + if err != nil { + return "", err + } + out, err := tpl.Execute(vs) + if err != nil { + return "", err + } + return out, nil + default: + return "", fmt.Errorf("unknown format type: %v", formatType) + } +} + +// Format returns the messages after rendering by the given formatType. +// e.g. +// +// msg := schema.UserMessage("hello world, {name}") +// msgs, err := msg.Format(ctx, map[string]any{"name": "eino"}, schema.FString) // <= this will render the content of msg by pyfmt +// // msgs[0].Content will be "hello world, eino" +func (m *Message) Format(_ context.Context, vs map[string]any, formatType FormatType) ([]*Message, error) { + c, err := formatContent(m.Content, vs, formatType) + if err != nil { + return nil, err + } + copied := *m + copied.Content = c + + if len(m.MultiContent) > 0 { + copied.MultiContent, err = formatMultiContent(m.MultiContent, vs, formatType) + if err != nil { + return nil, err + } + } + + if len(m.UserInputMultiContent) > 0 { + copied.UserInputMultiContent, err = formatUserInputMultiContent(m.UserInputMultiContent, vs, formatType) + if err != nil { + return nil, err + } + } + + return []*Message{&copied}, nil +} + +func formatMultiContent(multiContent []ChatMessagePart, vs map[string]any, formatType FormatType) ([]ChatMessagePart, error) { + copiedMC := make([]ChatMessagePart, len(multiContent)) + copy(copiedMC, multiContent) + + for i, mc := range copiedMC { + switch mc.Type { + case ChatMessagePartTypeText: + nmc, err := formatContent(mc.Text, vs, formatType) + if err != nil { + return nil, err + } + copiedMC[i].Text = nmc + case ChatMessagePartTypeImageURL: + if mc.ImageURL == nil { + continue + } + url, err := formatContent(mc.ImageURL.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedMC[i].ImageURL.URL = url + case ChatMessagePartTypeAudioURL: + if mc.AudioURL == nil { + continue + } + url, err := formatContent(mc.AudioURL.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedMC[i].AudioURL.URL = url + case ChatMessagePartTypeVideoURL: + if mc.VideoURL == nil { + continue + } + url, err := formatContent(mc.VideoURL.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedMC[i].VideoURL.URL = url + case ChatMessagePartTypeFileURL: + if mc.FileURL == nil { + continue + } + url, err := formatContent(mc.FileURL.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedMC[i].FileURL.URL = url + } + } + + return copiedMC, nil +} + +func formatUserInputMultiContent(userInputMultiContent []MessageInputPart, vs map[string]any, formatType FormatType) ([]MessageInputPart, error) { + copiedUIMC := make([]MessageInputPart, len(userInputMultiContent)) + copy(copiedUIMC, userInputMultiContent) + + for i, uimc := range copiedUIMC { + switch uimc.Type { + case ChatMessagePartTypeText: + text, err := formatContent(uimc.Text, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Text = text + case ChatMessagePartTypeImageURL: + if uimc.Image == nil { + continue + } + if uimc.Image.URL != nil && *uimc.Image.URL != "" { + url, err := formatContent(*uimc.Image.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Image.URL = &url + } + if uimc.Image.Base64Data != nil && *uimc.Image.Base64Data != "" { + base64data, err := formatContent(*uimc.Image.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Image.Base64Data = &base64data + } + case ChatMessagePartTypeAudioURL: + if uimc.Audio == nil { + continue + } + if uimc.Audio.URL != nil && *uimc.Audio.URL != "" { + url, err := formatContent(*uimc.Audio.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Audio.URL = &url + } + if uimc.Audio.Base64Data != nil && *uimc.Audio.Base64Data != "" { + base64data, err := formatContent(*uimc.Audio.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Audio.Base64Data = &base64data + } + case ChatMessagePartTypeVideoURL: + if uimc.Video == nil { + continue + } + if uimc.Video.URL != nil && *uimc.Video.URL != "" { + url, err := formatContent(*uimc.Video.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Video.URL = &url + } + if uimc.Video.Base64Data != nil && *uimc.Video.Base64Data != "" { + base64data, err := formatContent(*uimc.Video.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].Video.Base64Data = &base64data + } + case ChatMessagePartTypeFileURL: + if uimc.File == nil { + continue + } + if uimc.File.URL != nil && *uimc.File.URL != "" { + url, err := formatContent(*uimc.File.URL, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].File.URL = &url + } + if uimc.File.Base64Data != nil && *uimc.File.Base64Data != "" { + base64data, err := formatContent(*uimc.File.Base64Data, vs, formatType) + if err != nil { + return nil, err + } + copiedUIMC[i].File.Base64Data = &base64data + } + } + } + + return copiedUIMC, nil +} + +// String returns the string representation of the message. +// e.g. +// +// msg := schema.UserMessage("hello world") +// fmt.Println(msg.String()) // Output will be: `user: hello world`` +// +// msg := schema.Message{ +// Role: schema.Tool, +// Content: "{...}", +// ToolCallID: "callxxxx" +// } +// fmt.Println(msg.String()) +// Output will be: +// tool: {...} +// call_id: callxxxx +func (m *Message) String() string { + sb := &strings.Builder{} + sb.WriteString(fmt.Sprintf("%s: %s", m.Role, m.Content)) + + if len(m.UserInputMultiContent) > 0 { + sb.WriteString("\nuser_input_multi_content:") + for i, part := range m.UserInputMultiContent { + sb.WriteString(fmt.Sprintf("\n [%d] %s", i, formatInputPart(part))) + } + } + + if len(m.AssistantGenMultiContent) > 0 { + sb.WriteString("\nassistant_gen_multi_content:") + for i, part := range m.AssistantGenMultiContent { + sb.WriteString(fmt.Sprintf("\n [%d] %s", i, formatOutputPart(part))) + } + } + + if len(m.MultiContent) > 0 { + sb.WriteString("\nmulti_content:") + for i, part := range m.MultiContent { + sb.WriteString(fmt.Sprintf("\n [%d] %s", i, formatChatMessagePart(part))) + } + } + + if len(m.ReasoningContent) > 0 { + sb.WriteString("\nreasoning content:\n") + sb.WriteString(m.ReasoningContent) + } + if len(m.ToolCalls) > 0 { + sb.WriteString("\ntool_calls:\n") + for _, tc := range m.ToolCalls { + if tc.Index != nil { + sb.WriteString(fmt.Sprintf("index[%d]:", *tc.Index)) + } + sb.WriteString(fmt.Sprintf("%+v\n", tc)) + } + } + if m.ToolCallID != "" { + sb.WriteString(fmt.Sprintf("\ntool_call_id: %s", m.ToolCallID)) + } + if m.ToolName != "" { + sb.WriteString(fmt.Sprintf("\ntool_call_name: %s", m.ToolName)) + } + if m.ResponseMeta != nil { + sb.WriteString(fmt.Sprintf("\nfinish_reason: %s", m.ResponseMeta.FinishReason)) + if m.ResponseMeta.Usage != nil { + sb.WriteString(fmt.Sprintf("\nusage: %v", m.ResponseMeta.Usage)) + } + } + + return sb.String() +} + +func formatInputPart(part MessageInputPart) string { + switch part.Type { + case ChatMessagePartTypeText: + return fmt.Sprintf("text: %s", part.Text) + case ChatMessagePartTypeImageURL: + return fmt.Sprintf("image: %s", formatMessageInputMedia(part.Image)) + case ChatMessagePartTypeAudioURL: + return fmt.Sprintf("audio: %s", formatMessageInputMedia(part.Audio)) + case ChatMessagePartTypeVideoURL: + return fmt.Sprintf("video: %s", formatMessageInputMedia(part.Video)) + case ChatMessagePartTypeFileURL: + return fmt.Sprintf("file: %s", formatMessageInputFile(part.File)) + default: + return fmt.Sprintf("unknown type: %s", part.Type) + } +} + +func formatMessageInputMedia[T MessageInputImage | MessageInputAudio | MessageInputVideo](media *T) string { + if media == nil { + return "" + } + var parts []string + switch v := any(media).(type) { + case *MessageInputImage: + if v.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *v.URL)) + } + if v.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*v.Base64Data))) + } + if v.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", v.MIMEType)) + } + if v.Detail != "" { + parts = append(parts, fmt.Sprintf("detail=%s", v.Detail)) + } + if len(v.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", v.Extra)) + } + case *MessageInputAudio: + if v.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *v.URL)) + } + if v.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*v.Base64Data))) + } + if v.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", v.MIMEType)) + } + if len(v.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", v.Extra)) + } + case *MessageInputVideo: + if v.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *v.URL)) + } + if v.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*v.Base64Data))) + } + if v.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", v.MIMEType)) + } + if len(v.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", v.Extra)) + } + } + if len(parts) == 0 { + return "" + } + return strings.Join(parts, ", ") +} + +func formatMessageInputFile(file *MessageInputFile) string { + if file == nil { + return "" + } + var parts []string + if file.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *file.URL)) + } + if file.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*file.Base64Data))) + } + if file.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", file.MIMEType)) + } + if file.Name != "" { + parts = append(parts, fmt.Sprintf("name=%s", file.Name)) + } + if len(file.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", file.Extra)) + } + if len(parts) == 0 { + return "" + } + return strings.Join(parts, ", ") +} + +func formatOutputPart(part MessageOutputPart) string { + switch part.Type { + case ChatMessagePartTypeText: + return fmt.Sprintf("text: %s", part.Text) + case ChatMessagePartTypeImageURL: + return fmt.Sprintf("image: %s", formatMessageOutputMedia(part.Image)) + case ChatMessagePartTypeAudioURL: + return fmt.Sprintf("audio: %s", formatMessageOutputMedia(part.Audio)) + case ChatMessagePartTypeVideoURL: + return fmt.Sprintf("video: %s", formatMessageOutputMedia(part.Video)) + default: + return fmt.Sprintf("unknown type: %s", part.Type) + } +} + +func formatMessageOutputMedia[T MessageOutputImage | MessageOutputAudio | MessageOutputVideo](media *T) string { + if media == nil { + return "" + } + var parts []string + switch v := any(media).(type) { + case *MessageOutputImage: + if v.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *v.URL)) + } + if v.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*v.Base64Data))) + } + if v.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", v.MIMEType)) + } + if len(v.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", v.Extra)) + } + case *MessageOutputAudio: + if v.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *v.URL)) + } + if v.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*v.Base64Data))) + } + if v.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", v.MIMEType)) + } + if len(v.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", v.Extra)) + } + case *MessageOutputVideo: + if v.URL != nil { + parts = append(parts, fmt.Sprintf("url=%s", *v.URL)) + } + if v.Base64Data != nil { + parts = append(parts, fmt.Sprintf("base64[%d bytes]", len(*v.Base64Data))) + } + if v.MIMEType != "" { + parts = append(parts, fmt.Sprintf("mime=%s", v.MIMEType)) + } + if len(v.Extra) > 0 { + parts = append(parts, fmt.Sprintf("extra=%v", v.Extra)) + } + } + if len(parts) == 0 { + return "" + } + return strings.Join(parts, ", ") +} + +func formatChatMessagePart(part ChatMessagePart) string { + switch part.Type { + case ChatMessagePartTypeText: + return fmt.Sprintf("text: %s", part.Text) + case ChatMessagePartTypeImageURL: + if part.ImageURL != nil { + return fmt.Sprintf("image_url: %s", part.ImageURL.URL) + } + return "image_url: " + case ChatMessagePartTypeAudioURL: + if part.AudioURL != nil { + return fmt.Sprintf("audio_url: %s", part.AudioURL.URL) + } + return "audio_url: " + case ChatMessagePartTypeVideoURL: + if part.VideoURL != nil { + return fmt.Sprintf("video_url: %s", part.VideoURL.URL) + } + return "video_url: " + case ChatMessagePartTypeFileURL: + if part.FileURL != nil { + return fmt.Sprintf("file_url: %s", part.FileURL.URL) + } + return "file_url: " + default: + return fmt.Sprintf("unknown type: %s", part.Type) + } +} + +// SystemMessage represents a message with Role "system". +func SystemMessage(content string) *Message { + return &Message{ + Role: System, + Content: content, + } +} + +// AssistantMessage represents a message with Role "assistant". +func AssistantMessage(content string, toolCalls []ToolCall) *Message { + return &Message{ + Role: Assistant, + Content: content, + ToolCalls: toolCalls, + } +} + +// UserMessage represents a message with Role "user". +func UserMessage(content string) *Message { + return &Message{ + Role: User, + Content: content, + } + +} + +type toolMessageOptions struct { + toolName string +} + +// ToolMessageOption defines a option for ToolMessage +type ToolMessageOption func(*toolMessageOptions) + +// WithToolName returns a ToolMessageOption that sets the tool call name. +func WithToolName(name string) ToolMessageOption { + return func(o *toolMessageOptions) { + o.toolName = name + } +} + +// ToolMessage represents a message with Role "tool". +func ToolMessage(content string, toolCallID string, opts ...ToolMessageOption) *Message { + o := &toolMessageOptions{} + for _, opt := range opts { + opt(o) + } + return &Message{ + Role: Tool, + Content: content, + ToolCallID: toolCallID, + ToolName: o.toolName, + } +} + +// ConcatToolResults merges multiple ToolResult chunks into a single ToolResult. +// It collects all ToolOutputParts from the input chunks and merges contiguous text parts within each chunk. +// +// Merge rules: +// - Text parts: Contiguous text parts within each chunk are concatenated into a single text part. +// - Non-text parts (image, audio, video, file): These parts are kept as-is without merging. +// Each non-text part type can only appear in one chunk; if the same non-text type appears +// in multiple chunks, an error is returned. +// +// This function is primarily used in streaming scenarios where tool output is delivered +// in multiple chunks that need to be merged into a complete result. +// +// Parameters: +// - chunks: A slice of ToolResult pointers representing sequential chunks from a stream. +// Nil chunks and chunks with empty Parts are safely ignored. +// +// Returns: +// - *ToolResult: The merged ToolResult containing all content from the chunks. +// Returns an empty ToolResult if chunks is empty or all chunks are nil/empty. +// - error: An error if the same non-text part type appears in multiple chunks. +func ConcatToolResults(chunks []*ToolResult) (*ToolResult, error) { + if len(chunks) == 0 { + return &ToolResult{}, nil + } + + nonTextPartTypes := make(map[ToolPartType]int) + + var allParts []ToolOutputPart + for chunkIdx, chunk := range chunks { + if chunk == nil || len(chunk.Parts) == 0 { + continue + } + + for _, part := range chunk.Parts { + if part.Type != ToolPartTypeText { + // This restricts non-text modal content only appear once. + if prevChunkIdx, exists := nonTextPartTypes[part.Type]; exists { + return nil, fmt.Errorf("conflicting %s parts found in chunk %d and chunk %d: "+ + "non-text modality parts cannot appear in multiple chunks", part.Type, prevChunkIdx, chunkIdx) + } + nonTextPartTypes[part.Type] = chunkIdx + } + } + + allParts = append(allParts, chunk.Parts...) + } + + mergedChunkParts, err := concatToolOutputParts(allParts) + if err != nil { + return nil, fmt.Errorf("failed to merge parts: %w", err) + } + + return &ToolResult{Parts: mergedChunkParts}, nil +} + +func concatToolOutputParts(parts []ToolOutputPart) ([]ToolOutputPart, error) { + if len(parts) == 0 { + return nil, nil + } + + groups := groupToolOutputParts(parts) + + merged := make([]ToolOutputPart, 0, len(groups)) + for _, group := range groups { + if len(group) == 1 { + merged = append(merged, group...) + continue + } + switch group[0].Type { + case ToolPartTypeText: + mergedPart, err := mergeToolTextParts(group) + if err != nil { + return nil, err + } + merged = append(merged, mergedPart) + default: + merged = append(merged, group...) + } + } + + return merged, nil +} + +func groupToolOutputParts(parts []ToolOutputPart) [][]ToolOutputPart { + groups := make([][]ToolOutputPart, 0) + i := 0 + for i < len(parts) { + if parts[i].Type == ToolPartTypeText { + end := i + 1 + for end < len(parts) && parts[end].Type == ToolPartTypeText { + end++ + } + groups = append(groups, parts[i:end]) + i = end + } else { + groups = append(groups, parts[i:i+1]) + i++ + } + } + return groups +} + +func mergeToolTextParts(group []ToolOutputPart) (ToolOutputPart, error) { + var sb strings.Builder + extraList := make([]map[string]any, 0, len(group)) + for _, part := range group { + sb.WriteString(part.Text) + if len(part.Extra) > 0 { + extraList = append(extraList, part.Extra) + } + } + var mergedExtra map[string]any + if len(extraList) > 0 { + var err error + mergedExtra, err = concatExtra(extraList) + if err != nil { + return ToolOutputPart{}, fmt.Errorf("failed to concat tool output text part extra: %w", err) + } + } + return ToolOutputPart{ + Type: ToolPartTypeText, + Text: sb.String(), + Extra: mergedExtra, + }, nil +} + +func concatToolCalls(chunks []ToolCall) ([]ToolCall, error) { + var merged []ToolCall + m := make(map[int][]int) + for i := range chunks { + index := chunks[i].Index + if index == nil { + merged = append(merged, chunks[i]) + } else { + m[*index] = append(m[*index], i) + } + } + + var args strings.Builder + for k, v := range m { + index := k + toolCall := ToolCall{Index: &index} + if len(v) > 0 { + toolCall = chunks[v[0]] + } + + args.Reset() + toolID, toolType, toolName := "", "", "" // these field will output atomically in any chunk + + for _, n := range v { + chunk := chunks[n] + if chunk.ID != "" { + if toolID == "" { + toolID = chunk.ID + } else if toolID != chunk.ID { + return nil, fmt.Errorf("cannot concat ToolCalls with different tool id: '%s' '%s'", toolID, chunk.ID) + } + + } + + if chunk.Type != "" { + if toolType == "" { + toolType = chunk.Type + } else if toolType != chunk.Type { + return nil, fmt.Errorf("cannot concat ToolCalls with different tool type: '%s' '%s'", toolType, chunk.Type) + } + } + + if chunk.Function.Name != "" { + if toolName == "" { + toolName = chunk.Function.Name + } else if toolName != chunk.Function.Name { + return nil, fmt.Errorf("cannot concat ToolCalls with different tool name: '%s' '%s'", toolName, chunk.Function.Name) + } + } + + if chunk.Function.Arguments != "" { + _, err := args.WriteString(chunk.Function.Arguments) + if err != nil { + return nil, err + } + } + } + + toolCall.ID = toolID + toolCall.Type = toolType + toolCall.Function.Name = toolName + toolCall.Function.Arguments = args.String() + + merged = append(merged, toolCall) + } + + if len(merged) > 1 { + sort.SliceStable(merged, func(i, j int) bool { + iVal, jVal := merged[i].Index, merged[j].Index + if iVal == nil && jVal == nil { + return false + } else if iVal == nil && jVal != nil { + return true + } else if iVal != nil && jVal == nil { + return false + } + + return *iVal < *jVal + }) + } + + return merged, nil +} + +func concatAssistantMultiContent(parts []MessageOutputPart) ([]MessageOutputPart, error) { + if len(parts) == 0 { + return parts, nil + } + + groups := groupOutputParts(parts) + + merged := make([]MessageOutputPart, 0, len(groups)) + for _, group := range groups { + mergedPart, err := mergeOutputPartGroup(group) + if err != nil { + return nil, err + } + merged = append(merged, mergedPart) + } + + return merged, nil +} + +func groupOutputParts(parts []MessageOutputPart) [][]MessageOutputPart { + if len(parts) == 0 { + return nil + } + + groups := make([][]MessageOutputPart, 0) + currentGroup := []MessageOutputPart{parts[0]} + + for i := 1; i < len(parts); i++ { + if canMergeOutputParts(currentGroup[0], parts[i]) { + currentGroup = append(currentGroup, parts[i]) + } else { + groups = append(groups, currentGroup) + currentGroup = []MessageOutputPart{parts[i]} + } + } + groups = append(groups, currentGroup) + + return groups +} + +func canMergeOutputParts(current, next MessageOutputPart) bool { + if current.Type != next.Type { + return false + } + + if !isMergeableOutputPartType(current) { + return false + } + + if current.StreamingMeta != nil && next.StreamingMeta != nil { + return current.StreamingMeta.Index == next.StreamingMeta.Index + } + + return current.StreamingMeta == nil && next.StreamingMeta == nil +} + +func isMergeableOutputPartType(part MessageOutputPart) bool { + switch part.Type { + case ChatMessagePartTypeText, ChatMessagePartTypeReasoning: + return true + case ChatMessagePartTypeAudioURL: + return isBase64MessageOutputAudioPart(part) + default: + return false + } +} + +func mergeOutputPartGroup(group []MessageOutputPart) (MessageOutputPart, error) { + if len(group) == 0 { + return MessageOutputPart{}, nil + } + + if len(group) == 1 { + return group[0], nil + } + + first := group[0] + switch first.Type { + case ChatMessagePartTypeText: + return mergeTextParts(group) + case ChatMessagePartTypeReasoning: + return mergeReasoningParts(group) + case ChatMessagePartTypeAudioURL: + if isBase64MessageOutputAudioPart(first) { + return mergeAudioParts(group) + } + } + + return first, nil +} + +func mergeTextParts(group []MessageOutputPart) (MessageOutputPart, error) { + var sb strings.Builder + extraList := make([]map[string]any, 0, len(group)) + for _, part := range group { + sb.WriteString(part.Text) + if len(part.Extra) > 0 { + extraList = append(extraList, part.Extra) + } + } + var mergedExtra map[string]any + if len(extraList) > 0 { + var err error + mergedExtra, err = concatExtra(extraList) + if err != nil { + return MessageOutputPart{}, fmt.Errorf("failed to concat text part extra: %w", err) + } + } + return MessageOutputPart{ + Type: ChatMessagePartTypeText, + Text: sb.String(), + Extra: mergedExtra, + StreamingMeta: group[0].StreamingMeta, + }, nil +} + +func mergeReasoningParts(group []MessageOutputPart) (MessageOutputPart, error) { + var textBuilder strings.Builder + var signature string + extraList := make([]map[string]any, 0, len(group)) + for _, part := range group { + if part.Reasoning != nil { + textBuilder.WriteString(part.Reasoning.Text) + if part.Reasoning.Signature != "" { + signature = part.Reasoning.Signature + } + } + if len(part.Extra) > 0 { + extraList = append(extraList, part.Extra) + } + } + var mergedExtra map[string]any + if len(extraList) > 0 { + var err error + mergedExtra, err = concatExtra(extraList) + if err != nil { + return MessageOutputPart{}, fmt.Errorf("failed to concat reasoning part extra: %w", err) + } + } + return MessageOutputPart{ + Type: ChatMessagePartTypeReasoning, + Reasoning: &MessageOutputReasoning{ + Text: textBuilder.String(), + Signature: signature, + }, + Extra: mergedExtra, + StreamingMeta: group[0].StreamingMeta, + }, nil +} + +func mergeAudioParts(group []MessageOutputPart) (MessageOutputPart, error) { + var b64Builder strings.Builder + var mimeType string + audioExtraList := make([]map[string]any, 0, len(group)) + partExtraList := make([]map[string]any, 0, len(group)) + + for _, part := range group { + audioPart := part.Audio + if audioPart.Base64Data != nil { + b64Builder.WriteString(*audioPart.Base64Data) + } + if mimeType == "" { + mimeType = audioPart.MIMEType + } + if len(audioPart.Extra) > 0 { + audioExtraList = append(audioExtraList, audioPart.Extra) + } + if len(part.Extra) > 0 { + partExtraList = append(partExtraList, part.Extra) + } + } + + var mergedAudioExtra map[string]any + var err error + if len(audioExtraList) > 0 { + mergedAudioExtra, err = concatExtra(audioExtraList) + if err != nil { + return MessageOutputPart{}, fmt.Errorf("failed to concat audio extra: %w", err) + } + } + + var mergedPartExtra map[string]any + if len(partExtraList) > 0 { + mergedPartExtra, err = concatExtra(partExtraList) + if err != nil { + return MessageOutputPart{}, fmt.Errorf("failed to concat audio part extra: %w", err) + } + } + + mergedB64 := b64Builder.String() + return MessageOutputPart{ + Type: ChatMessagePartTypeAudioURL, + Audio: &MessageOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &mergedB64, + MIMEType: mimeType, + Extra: mergedAudioExtra, + }, + }, + Extra: mergedPartExtra, + StreamingMeta: group[0].StreamingMeta, + }, nil +} + +func isBase64MessageOutputAudioPart(part MessageOutputPart) bool { + return part.Type == ChatMessagePartTypeAudioURL && + part.Audio != nil && + part.Audio.Base64Data != nil && + part.Audio.URL == nil +} + +func concatUserMultiContent(parts []MessageInputPart) ([]MessageInputPart, error) { + if len(parts) == 0 { + return parts, nil + } + + merged := make([]MessageInputPart, 0, len(parts)) + i := 0 + for i < len(parts) { + currentPart := parts[i] + + if currentPart.Type == ChatMessagePartTypeText { + end := i + 1 + for end < len(parts) && parts[end].Type == ChatMessagePartTypeText { + end++ + } + + if end == i+1 { + merged = append(merged, currentPart) + } else { + var sb strings.Builder + for k := i; k < end; k++ { + sb.WriteString(parts[k].Text) + } + mergedPart := MessageInputPart{ + Type: ChatMessagePartTypeText, + Text: sb.String(), + } + merged = append(merged, mergedPart) + } + i = end + } else { + + merged = append(merged, currentPart) + i++ + } + } + + return merged, nil +} + +func concatExtra(extraList []map[string]any) (map[string]any, error) { + if len(extraList) == 1 { + return generic.CopyMap(extraList[0]), nil + } + + return internal.ConcatItems(extraList) +} + +// ConcatMessages concat messages with the same role and name. +// It will concat tool calls with the same index. +// It will return an error if the messages have different roles or names. +// It's useful for concatenating messages from a stream. +// e.g. +// +// msgs := []*Message{} +// for { +// msg, err := stream.Recv() +// if errors.Is(err, io.EOF) { +// break +// } +// if err != nil {...} +// msgs = append(msgs, msg) +// } +// +// concatedMsg, err := ConcatMessages(msgs) // concatedMsg.Content will be full content of all messages +func ConcatMessages(msgs []*Message) (*Message, error) { + var ( + contents []string + contentLen int + reasoningContents []string + reasoningContentLen int + toolCalls []ToolCall + multiContentParts []ChatMessagePart + assistantGenMultiContentParts []MessageOutputPart + userInputMultiContentParts []MessageInputPart + ret = Message{} + extraList = make([]map[string]any, 0, len(msgs)) + ) + + for idx, msg := range msgs { + if msg == nil { + return nil, fmt.Errorf("unexpected nil chunk in message stream, index: %d", idx) + } + + if msg.Role != "" { + if ret.Role == "" { + ret.Role = msg.Role + } else if ret.Role != msg.Role { + return nil, fmt.Errorf("cannot concat messages with "+ + "different roles: '%s' '%s'", ret.Role, msg.Role) + } + } + + if msg.Name != "" { + if ret.Name == "" { + ret.Name = msg.Name + } else if ret.Name != msg.Name { + return nil, fmt.Errorf("cannot concat messages with"+ + " different names: '%s' '%s'", ret.Name, msg.Name) + } + } + + if msg.ToolCallID != "" { + if ret.ToolCallID == "" { + ret.ToolCallID = msg.ToolCallID + } else if ret.ToolCallID != msg.ToolCallID { + return nil, fmt.Errorf("cannot concat messages with"+ + " different toolCallIDs: '%s' '%s'", ret.ToolCallID, msg.ToolCallID) + } + } + if msg.ToolName != "" { + if ret.ToolName == "" { + ret.ToolName = msg.ToolName + } else if ret.ToolName != msg.ToolName { + return nil, fmt.Errorf("cannot concat messages with"+ + " different toolNames: '%s' '%s'", ret.ToolCallID, msg.ToolCallID) + } + } + + if msg.Content != "" { + contents = append(contents, msg.Content) + contentLen += len(msg.Content) + } + if msg.ReasoningContent != "" { + reasoningContents = append(reasoningContents, msg.ReasoningContent) + reasoningContentLen += len(msg.ReasoningContent) + } + + if len(msg.ToolCalls) > 0 { + toolCalls = append(toolCalls, msg.ToolCalls...) + } + + if len(msg.Extra) > 0 { + extraList = append(extraList, msg.Extra) + } + + // The 'MultiContent' field is deprecated but is kept for backward compatibility. + if len(msg.MultiContent) > 0 { + multiContentParts = append(multiContentParts, msg.MultiContent...) + } + + if len(msg.AssistantGenMultiContent) > 0 { + assistantGenMultiContentParts = append(assistantGenMultiContentParts, msg.AssistantGenMultiContent...) + } + if len(msg.UserInputMultiContent) > 0 { + userInputMultiContentParts = append(userInputMultiContentParts, msg.UserInputMultiContent...) + } + if msg.ResponseMeta != nil && ret.ResponseMeta == nil { + ret.ResponseMeta = &ResponseMeta{} + } + + if msg.ResponseMeta != nil && ret.ResponseMeta != nil { + // keep the last FinishReason with a valid value. + if msg.ResponseMeta.FinishReason != "" { + ret.ResponseMeta.FinishReason = msg.ResponseMeta.FinishReason + } + + if msg.ResponseMeta.Usage != nil { + if ret.ResponseMeta.Usage == nil { + ret.ResponseMeta.Usage = &TokenUsage{} + } + + if msg.ResponseMeta.Usage.PromptTokens > ret.ResponseMeta.Usage.PromptTokens { + ret.ResponseMeta.Usage.PromptTokens = msg.ResponseMeta.Usage.PromptTokens + } + if msg.ResponseMeta.Usage.CompletionTokens > ret.ResponseMeta.Usage.CompletionTokens { + ret.ResponseMeta.Usage.CompletionTokens = msg.ResponseMeta.Usage.CompletionTokens + } + + if msg.ResponseMeta.Usage.TotalTokens > ret.ResponseMeta.Usage.TotalTokens { + ret.ResponseMeta.Usage.TotalTokens = msg.ResponseMeta.Usage.TotalTokens + } + + if msg.ResponseMeta.Usage.PromptTokenDetails.CachedTokens > ret.ResponseMeta.Usage.PromptTokenDetails.CachedTokens { + ret.ResponseMeta.Usage.PromptTokenDetails.CachedTokens = msg.ResponseMeta.Usage.PromptTokenDetails.CachedTokens + } + + if msg.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens > ret.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens { + ret.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens = msg.ResponseMeta.Usage.CompletionTokensDetails.ReasoningTokens + } + } + + if msg.ResponseMeta.LogProbs != nil { + if ret.ResponseMeta.LogProbs == nil { + ret.ResponseMeta.LogProbs = &LogProbs{} + } + + ret.ResponseMeta.LogProbs.Content = append(ret.ResponseMeta.LogProbs.Content, msg.ResponseMeta.LogProbs.Content...) + } + + } + } + + if len(contents) > 0 { + var sb strings.Builder + sb.Grow(contentLen) + for _, content := range contents { + _, err := sb.WriteString(content) + if err != nil { + return nil, err + } + } + + ret.Content = sb.String() + } + if len(reasoningContents) > 0 { + var sb strings.Builder + sb.Grow(reasoningContentLen) + for _, rc := range reasoningContents { + _, err := sb.WriteString(rc) + if err != nil { + return nil, err + } + } + + ret.ReasoningContent = sb.String() + } + + if len(toolCalls) > 0 { + merged, err := concatToolCalls(toolCalls) + if err != nil { + return nil, err + } + + ret.ToolCalls = merged + } + + if len(extraList) > 0 { + extra, err := concatExtra(extraList) + if err != nil { + return nil, fmt.Errorf("failed to concat message's extra: %w", err) + } + + if len(extra) > 0 { + ret.Extra = extra + } + } + + if len(multiContentParts) > 0 { + ret.MultiContent = multiContentParts + } + + if len(assistantGenMultiContentParts) > 0 { + merged, err := concatAssistantMultiContent(assistantGenMultiContentParts) + if err != nil { + return nil, fmt.Errorf("failed to concat message's assistant multicontent: %w", err) + } + ret.AssistantGenMultiContent = merged + } + + if len(userInputMultiContentParts) > 0 { + merged, err := concatUserMultiContent(userInputMultiContentParts) + if err != nil { + return nil, fmt.Errorf("failed to concat message's user multicontent: %w", err) + } + ret.UserInputMultiContent = merged + } + + return &ret, nil +} + +// ConcatMessageStream drains a stream of messages and returns a single +// concatenated message representing the merged content. +func ConcatMessageStream(s *StreamReader[*Message]) (*Message, error) { + defer s.Close() + + var msgs []*Message + for { + msg, err := s.Recv() + if err != nil { + if err == io.EOF { + break + } + + return nil, err + } + + msgs = append(msgs, msg) + } + + return ConcatMessages(msgs) +} + +// custom jinja env +var jinjaEnvOnce sync.Once +var jinjaEnv *gonja.Environment +var envInitErr error + +const ( + jinjaInclude = "include" + jinjaExtends = "extends" + jinjaImport = "import" + jinjaFrom = "from" + jinjaFile = "file" + jinjaFileSet = "fileset" +) + +func getJinjaEnv() (*gonja.Environment, error) { + jinjaEnvOnce.Do(func() { + jinjaEnv = gonja.NewEnvironment(config.DefaultConfig, gonja.DefaultLoader) + formatInitError := "init jinja env fail: %w" + var err error + if jinjaEnv.Statements.Exists(jinjaInclude) { + err = jinjaEnv.Statements.Replace(jinjaInclude, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) { + return nil, fmt.Errorf("keyword[include] has been disabled") + }) + if err != nil { + envInitErr = fmt.Errorf(formatInitError, err) + return + } + } + if jinjaEnv.Statements.Exists(jinjaExtends) { + err = jinjaEnv.Statements.Replace(jinjaExtends, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) { + return nil, fmt.Errorf("keyword[extends] has been disabled") + }) + if err != nil { + envInitErr = fmt.Errorf(formatInitError, err) + return + } + } + if jinjaEnv.Statements.Exists(jinjaFrom) { + err = jinjaEnv.Statements.Replace(jinjaFrom, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) { + return nil, fmt.Errorf("keyword[from] has been disabled") + }) + if err != nil { + envInitErr = fmt.Errorf(formatInitError, err) + return + } + } + if jinjaEnv.Statements.Exists(jinjaImport) { + err = jinjaEnv.Statements.Replace(jinjaImport, func(parser *parser.Parser, args *parser.Parser) (nodes.Statement, error) { + return nil, fmt.Errorf("keyword[import] has been disabled") + }) + if err != nil { + envInitErr = fmt.Errorf(formatInitError, err) + return + } + } + if jinjaEnv.Filters.Exists(jinjaFile) { + err = jinjaEnv.Filters.Replace(jinjaFile, func(e *exec.Evaluator, in *exec.Value, params *exec.VarArgs) *exec.Value { + return exec.AsValue(fmt.Errorf("keyword[file] has been disabled")) + }) + if err != nil { + envInitErr = fmt.Errorf(formatInitError, err) + return + } + } + if jinjaEnv.Filters.Exists(jinjaFileSet) { + err = jinjaEnv.Filters.Replace(jinjaFileSet, func(e *exec.Evaluator, in *exec.Value, params *exec.VarArgs) *exec.Value { + return exec.AsValue(fmt.Errorf("keyword[fileset] has been disabled")) + }) + if err != nil { + envInitErr = fmt.Errorf(formatInitError, err) + return + } + } + }) + return jinjaEnv, envInitErr +} diff --git a/schema/message_parser.go b/schema/message_parser.go new file mode 100644 index 0000000..7bd95b8 --- /dev/null +++ b/schema/message_parser.go @@ -0,0 +1,141 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "context" + "fmt" + "strings" + + "github.com/bytedance/sonic" +) + +// MessageParser parses a Message into a strongly typed value. +type MessageParser[T any] interface { + Parse(ctx context.Context, m *Message) (T, error) +} + +// MessageParseFrom determines the source of the data to be parsed. default is content (Message.Content). +type MessageParseFrom string + +// MessageParseFrom indicates the source data used by the parser. +const ( + MessageParseFromContent MessageParseFrom = "content" + MessageParseFromToolCall MessageParseFrom = "tool_call" +) + +// MessageJSONParseConfig configures JSON parsing behavior for Message. +type MessageJSONParseConfig struct { + // parse from content or tool call, default is content. + ParseFrom MessageParseFrom `json:"parse_from,omitempty"` + + // parse key path, default is empty. + // must be a valid json path expression, eg: field.sub_field + ParseKeyPath string `json:"parse_key_path,omitempty"` +} + +// NewMessageJSONParser creates a new MessageJSONParser. +func NewMessageJSONParser[T any](config *MessageJSONParseConfig) MessageParser[T] { + if config == nil { + config = &MessageJSONParseConfig{} + } + + if config.ParseFrom == "" { + config.ParseFrom = MessageParseFromContent + } + + return &MessageJSONParser[T]{ + ParseFrom: config.ParseFrom, + ParseKeyPath: config.ParseKeyPath, + } +} + +// MessageJSONParser is a parser that parses a message into an object T, using json unmarshal. +// eg of parse to single struct: +// +// config := &MessageJSONParseConfig{ +// ParseFrom: MessageParseFromToolCall, +// } +// parser := NewMessageJSONParser[GetUserParam](config) +// param, err := parser.Parse(ctx, message) +// +// eg of parse to slice of struct: +// +// config := &MessageJSONParseConfig{ +// ParseFrom: MessageParseFromToolCall, +// } +// +// parser := NewMessageJSONParser[GetUserParam](config) +// param, err := parser.Parse(ctx, message) +type MessageJSONParser[T any] struct { + ParseFrom MessageParseFrom + ParseKeyPath string +} + +// Parse parses a message into an object T. +func (p *MessageJSONParser[T]) Parse(ctx context.Context, m *Message) (parsed T, err error) { + if p.ParseFrom == MessageParseFromContent { + return p.parse(m.Content) + } else if p.ParseFrom == MessageParseFromToolCall { + if len(m.ToolCalls) == 0 { + return parsed, fmt.Errorf("no tool call found") + } + + return p.parse(m.ToolCalls[0].Function.Arguments) + } + + return parsed, fmt.Errorf("invalid parse from type: %s", p.ParseFrom) +} + +// extractData extracts data from a string using the parse key path. +func (p *MessageJSONParser[T]) extractData(data string) (string, error) { + if p.ParseKeyPath == "" { + return data, nil + } + + keys := strings.Split(p.ParseKeyPath, ".") + interfaceKeys := make([]any, len(keys)) + for i, key := range keys { + interfaceKeys[i] = key + } + + node, err := sonic.GetFromString(data, interfaceKeys...) + if err != nil { + return "", fmt.Errorf("failed to get parse key path: %w", err) + } + + bytes, err := node.MarshalJSON() + if err != nil { + return "", fmt.Errorf("failed to marshal node: %w", err) + } + + return string(bytes), nil +} + +// parse parses a string into an object T. +func (p *MessageJSONParser[T]) parse(data string) (parsed T, err error) { + parsedData, err := p.extractData(data) + if err != nil { + return parsed, err + } + + if err := sonic.UnmarshalString(parsedData, &parsed); err != nil { + return parsed, fmt.Errorf("failed to unmarshal content: %w", err) + } + + return parsed, nil +} diff --git a/schema/message_parser_test.go b/schema/message_parser_test.go new file mode 100644 index 0000000..e645f87 --- /dev/null +++ b/schema/message_parser_test.go @@ -0,0 +1,179 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +type TestStructForParse struct { + ID int `json:"id"` + Name string `json:"name"` + XX struct { + YY int `json:"yy"` + } `json:"xx"` +} + +func TestMessageJSONParser(t *testing.T) { + ctx := context.Background() + + t.Run("parse from content", func(t *testing.T) { + parser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromContent, + }) + + parsed, err := parser.Parse(ctx, &Message{ + Content: `{"id": 1, "name": "test", "xx": {"yy": 2}}`, + }) + assert.Nil(t, err) + assert.Equal(t, 1, parsed.ID) + }) + + t.Run("parse from tool call", func(t *testing.T) { + t.Run("only one tool call, default use first tool call", func(t *testing.T) { + parser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromToolCall, + }) + + parsed, err := parser.Parse(ctx, &Message{ + ToolCalls: []ToolCall{ + {Function: FunctionCall{Arguments: `{"id": 1, "name": "test", "xx": {"yy": 2}}`}}, + }, + }) + assert.Nil(t, err) + assert.Equal(t, 1, parsed.ID) + }) + + t.Run("parse key path", func(t *testing.T) { + type TestStructForParse2 struct { + YY int `json:"yy"` + } + + parser := NewMessageJSONParser[TestStructForParse2](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromToolCall, + ParseKeyPath: "xx", + }) + + parsed, err := parser.Parse(ctx, &Message{ + ToolCalls: []ToolCall{ + {Function: FunctionCall{Arguments: `{"id": 1, "name": "test", "xx": {"yy": 2}}`}}, + }, + }) + assert.Nil(t, err) + assert.Equal(t, 2, parsed.YY) + }) + + t.Run("parse key path, deep level", func(t *testing.T) { + type TestStructForParse3 struct { + ZZ int `json:"zz"` + } + + parser := NewMessageJSONParser[TestStructForParse3](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromToolCall, + ParseKeyPath: "xx.yy", + }) + + parsed, err := parser.Parse(ctx, &Message{ + ToolCalls: []ToolCall{ + {Function: FunctionCall{Arguments: `{"id": 1, "name": "test", "xx": {"yy": {"zz": 3}}}`}}, + }, + }) + assert.Nil(t, err) + assert.Equal(t, 3, parsed.ZZ) + }) + + t.Run("parse key with pointer", func(t *testing.T) { + type TestStructForParse4 struct { + ZZ *int `json:"zz"` + } + + parser := NewMessageJSONParser[**TestStructForParse4](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromToolCall, + }) + + parsed, err := parser.Parse(ctx, &Message{ + ToolCalls: []ToolCall{{Function: FunctionCall{Arguments: `{"zz": 3}`}}}, + }) + assert.Nil(t, err) + assert.Equal(t, 3, *((**parsed).ZZ)) + }) + }) + + t.Run("parse of slice", func(t *testing.T) { + t.Run("valid slice string, not multiple tool calls", func(t *testing.T) { + parser := NewMessageJSONParser[[]map[string]any](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromToolCall, + }) + + parsed, err := parser.Parse(ctx, &Message{ + ToolCalls: []ToolCall{{Function: FunctionCall{Arguments: `[{"id": 1}, {"id": 2}]`}}}, + }) + assert.Nil(t, err) + assert.Equal(t, 2, len(parsed)) + }) + + t.Run("invalid slice string, not multiple tool calls", func(t *testing.T) { + parser := NewMessageJSONParser[[]map[string]any](&MessageJSONParseConfig{ + ParseFrom: MessageParseFromToolCall, + }) + + _, err := parser.Parse(ctx, &Message{ + ToolCalls: []ToolCall{ + {Function: FunctionCall{Arguments: `{"id": 1}`}}, + {Function: FunctionCall{Arguments: `{"id": 2}`}}, + }, + }) + assert.NotNil(t, err) + }) + }) + + t.Run("invalid configs", func(t *testing.T) { + parser := NewMessageJSONParser[TestStructForParse](nil) + _, err := parser.Parse(ctx, &Message{ + Content: "", + }) + assert.NotNil(t, err) + }) + + t.Run("invalid parse key path", func(t *testing.T) { + parser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{ + ParseKeyPath: "...invalid", + }) + _, err := parser.Parse(ctx, &Message{}) + assert.NotNil(t, err) + }) + + t.Run("invalid parse from", func(t *testing.T) { + parser := NewMessageJSONParser[TestStructForParse](&MessageJSONParseConfig{ + ParseFrom: "invalid", + }) + _, err := parser.Parse(ctx, &Message{}) + assert.NotNil(t, err) + }) + + t.Run("invalid parse from type", func(t *testing.T) { + parser := NewMessageJSONParser[int](&MessageJSONParseConfig{ + ParseFrom: MessageParseFrom("invalid"), + }) + _, err := parser.Parse(ctx, &Message{}) + assert.NotNil(t, err) + }) + +} diff --git a/schema/message_test.go b/schema/message_test.go new file mode 100644 index 0000000..bb086b8 --- /dev/null +++ b/schema/message_test.go @@ -0,0 +1,2189 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "context" + "reflect" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/generic" +) + +func TestMessageTemplate(t *testing.T) { + pyFmtMessage := UserMessage("input: {question}") + jinja2Message := UserMessage("input: {{question}}") + goTemplateMessage := UserMessage("input: {{.question}}") + ctx := context.Background() + question := "what's the weather today" + expected := []*Message{UserMessage("input: " + question)} + + ms, err := pyFmtMessage.Format(ctx, map[string]any{"question": question}, FString) + assert.Nil(t, err) + assert.True(t, reflect.DeepEqual(expected, ms)) + ms, err = jinja2Message.Format(ctx, map[string]any{"question": question}, Jinja2) + assert.Nil(t, err) + assert.True(t, reflect.DeepEqual(expected, ms)) + ms, err = goTemplateMessage.Format(ctx, map[string]any{"question": question}, GoTemplate) + assert.Nil(t, err) + assert.True(t, reflect.DeepEqual(expected, ms)) + + mp := MessagesPlaceholder("chat_history", false) + m1 := UserMessage("how are you?") + m2 := AssistantMessage("I'm good. how about you?", nil) + ms, err = mp.Format(ctx, map[string]any{"chat_history": []*Message{m1, m2}}, FString) + assert.Nil(t, err) + + // len(ms) == 2 + assert.Equal(t, 2, len(ms)) + assert.Equal(t, ms[0], m1) + assert.Equal(t, ms[1], m2) +} + +func TestConcatMessage(t *testing.T) { + t.Run("tool_call_normal_append", func(t *testing.T) { + expectMsg := &Message{ + Role: "assistant", + Content: "", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "i_am_a_too_call_id", + Type: "function", + Function: FunctionCall{ + Name: "i_am_a_tool_name", + Arguments: "{}", + }, + }, + }, + } + givenMsgList := []*Message{ + { + Role: "", + Content: "", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + }, + }, + { + Role: "assistant", + Content: "", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "i_am_a_too_call_id", + Type: "function", + Function: FunctionCall{ + Name: "i_am_a_tool_name", + }, + }, + }, + }, + { + Role: "", + Content: "", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "", + Type: "", + Function: FunctionCall{ + Name: "", + Arguments: "{}", + }, + }, + }, + }, + } + + msg, err := ConcatMessages(givenMsgList) + assert.NoError(t, err) + assert.EqualValues(t, expectMsg, msg) + }) + + t.Run("exist_nil_message", func(t *testing.T) { + givenMsgList := []*Message{ + nil, + { + Role: "assistant", + Content: "", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "i_am_a_too_call_id", + Type: "function", + Function: FunctionCall{ + Name: "i_am_a_tool_name", + }, + }, + }, + }, + } + + _, err := ConcatMessages(givenMsgList) + assert.ErrorContains(t, err, "unexpected nil chunk in message stream") + }) + + t.Run("response_meta", func(t *testing.T) { + expectedMsg := &Message{ + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + Usage: &TokenUsage{ + CompletionTokens: 15, + PromptTokens: 30, + PromptTokenDetails: PromptTokenDetails{ + CachedTokens: 15, + }, + CompletionTokensDetails: CompletionTokensDetails{ + ReasoningTokens: 8, + }, + TotalTokens: 45, + }, + }, + } + + givenMsgList := []*Message{ + { + Role: "assistant", + }, + { + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "", + Usage: &TokenUsage{ + CompletionTokens: 10, + PromptTokens: 20, + PromptTokenDetails: PromptTokenDetails{ + CachedTokens: 10, + }, + CompletionTokensDetails: CompletionTokensDetails{ + ReasoningTokens: 5, + }, + TotalTokens: 30, + }, + }, + }, + { + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + }, + }, + { + Role: "assistant", + ResponseMeta: &ResponseMeta{ + Usage: &TokenUsage{ + CompletionTokens: 15, + PromptTokens: 30, + PromptTokenDetails: PromptTokenDetails{ + CachedTokens: 15, + }, + CompletionTokensDetails: CompletionTokensDetails{ + ReasoningTokens: 8, + }, + TotalTokens: 45, + }, + }, + }, + } + + msg, err := ConcatMessages(givenMsgList) + assert.NoError(t, err) + assert.Equal(t, expectedMsg, msg) + + givenMsgList = append(givenMsgList, &Message{ + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "tool_calls", + }, + }) + msg, err = ConcatMessages(givenMsgList) + assert.NoError(t, err) + expectedMsg.ResponseMeta.FinishReason = "tool_calls" + assert.Equal(t, expectedMsg, msg) + + }) + + t.Run("err: different roles", func(t *testing.T) { + msgs := []*Message{ + {Role: User}, + {Role: Assistant}, + } + + msg, err := ConcatMessages(msgs) + if assert.Error(t, err) { + assert.ErrorContains(t, err, "cannot concat messages with different roles") + assert.Nil(t, msg) + } + }) + + t.Run("err: different name", func(t *testing.T) { + msgs := []*Message{ + {Role: Assistant, Name: "n", Content: "1"}, + {Role: Assistant, Name: "a", Content: "2"}, + } + + msg, err := ConcatMessages(msgs) + if assert.Error(t, err) { + assert.ErrorContains(t, err, "cannot concat messages with different names") + assert.Nil(t, msg) + } + }) + + t.Run("err: different tool name", func(t *testing.T) { + msgs := []*Message{ + { + Role: "", + Content: "", + ToolCallID: "123", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "abc", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + }, + }, + { + Role: "assistant", + Content: "", + ToolCallID: "321", + ToolCalls: []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "abc", + Type: "function", + Function: FunctionCall{ + Name: "i_am_a_tool_name", + }, + }, + }, + }, + } + + msg, err := ConcatMessages(msgs) + if assert.Error(t, err) { + assert.ErrorContains(t, err, "cannot concat messages with different toolCallIDs") + assert.Nil(t, msg) + } + }) + + t.Run("first response meta usage is nil", func(t *testing.T) { + exp := &Message{ + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + Usage: &TokenUsage{ + CompletionTokens: 15, + PromptTokens: 30, + TotalTokens: 45, + }, + }, + } + + msgs := []*Message{ + { + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "", + Usage: nil, + }, + }, + { + Role: "assistant", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + }, + }, + { + Role: "assistant", + ResponseMeta: &ResponseMeta{ + Usage: &TokenUsage{ + CompletionTokens: 15, + PromptTokens: 30, + TotalTokens: 45, + }, + }, + }, + } + + msg, err := ConcatMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, exp, msg) + }) + + t.Run("concurrent concat", func(t *testing.T) { + content := "i_am_a_good_concat_message" + exp := &Message{Role: Assistant, Content: content} + var msgs []*Message + for i := 0; i < len(content); i++ { + msgs = append(msgs, &Message{Role: Assistant, Content: content[i : i+1]}) + } + + wg := sync.WaitGroup{} + size := 100 + wg.Add(size) + for i := 0; i < size; i++ { + go func() { + defer wg.Done() + msg, err := ConcatMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, exp, msg) + }() + } + + wg.Wait() + }) + + t.Run("concat logprobs", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + Content: "🚀", + ResponseMeta: &ResponseMeta{ + LogProbs: &LogProbs{ + Content: []LogProb{ + { + Token: "\\xf0\\x9f\\x9a", + LogProb: -0.0000073458323, + Bytes: []int64{240, 159, 154}, + }, + { + Token: "\\x80", + LogProb: 0, + Bytes: []int64{128}, + }, + }, + }, + }, + }, + { + Role: "", + Content: "❤️", + ResponseMeta: &ResponseMeta{ + LogProbs: &LogProbs{ + Content: []LogProb{ + { + Token: "❤️", + LogProb: -0.0011431955, + Bytes: []int64{226, 157, 164, 239, 184, 143}, + }, + }, + }, + }, + }, + { + Role: "", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + Usage: &TokenUsage{ + PromptTokens: 7, + CompletionTokens: 3, + TotalTokens: 10, + }, + }, + }, + } + + msg, err := ConcatMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, 3, len(msg.ResponseMeta.LogProbs.Content)) + assert.Equal(t, msgs[0].ResponseMeta.LogProbs.Content[0], msg.ResponseMeta.LogProbs.Content[0]) + assert.Equal(t, msgs[0].ResponseMeta.LogProbs.Content[1], msg.ResponseMeta.LogProbs.Content[1]) + assert.Equal(t, msgs[1].ResponseMeta.LogProbs.Content[0], msg.ResponseMeta.LogProbs.Content[2]) + }) + + t.Run("fix unexpected setting ResponseMeta of the first element in slice after ConcatMessages", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + Content: "🚀", + // ResponseMeta: &ResponseMeta{}, + }, + { + Role: "", + Content: "❤️", + ResponseMeta: &ResponseMeta{}, + }, + { + Role: "", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + Usage: &TokenUsage{ + PromptTokens: 7, + CompletionTokens: 3, + TotalTokens: 10, + }, + }, + }, + } + + msg, err := ConcatMessages(msgs) + assert.NoError(t, err) + assert.Equal(t, msgs[2].ResponseMeta, msg.ResponseMeta) + assert.Nil(t, msgs[0].ResponseMeta) + }) + + t.Run("concat assistant multi content", func(t *testing.T) { + base64Audio1 := "dGVzdF9hdWRpb18x" + base64Audio2 := "dGVzdF9hdWRpb18y" + imageURL1 := "https://example.com/image1.png" + imageURL2 := "https://example.com/image2.png" + + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello, "}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "world!"}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1}}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2, MIMEType: "audio/wav"}}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL1}}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL2}}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + mergedBase64Audio := base64Audio1 + base64Audio2 + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello, world!"}, + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio, MIMEType: "audio/wav"}}}, + {Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL1}}}, + {Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{MessagePartCommon: MessagePartCommon{URL: &imageURL2}}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat assistant multi content with extra", func(t *testing.T) { + base64Audio1 := "dGVzdF9hdWRpb18x" + base64Audio2 := "dGVzdF9hdWRpb18y" + + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1, Extra: map[string]any{"key1": "val1"}}}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2, Extra: map[string]any{"key2": "val2"}}}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + mergedBase64Audio := base64Audio1 + base64Audio2 + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio, Extra: map[string]any{"key1": "val1", "key2": "val2"}}}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat assistant multi content with single extra", func(t *testing.T) { + base64Audio1 := "dGVzdF9hdWRpb18x" + base64Audio2 := "dGVzdF9hdWRpb18y" + + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1, Extra: map[string]any{"key1": "val1"}}}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2}}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + mergedBase64Audio := base64Audio1 + base64Audio2 + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio, Extra: map[string]any{"key1": "val1"}}}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat text parts with extra", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello ", Extra: map[string]any{"key1": "val1"}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "World", Extra: map[string]any{"key2": "val2"}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello World", Extra: map[string]any{"key1": "val1", "key2": "val2"}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat text parts with single extra", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello ", Extra: map[string]any{"key1": "val1"}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "World"}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello World", Extra: map[string]any{"key1": "val1"}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat reasoning parts with extra", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "First, "}, Extra: map[string]any{"key1": "val1"}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "I need to think."}, Extra: map[string]any{"key2": "val2"}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "First, I need to think."}, Extra: map[string]any{"key1": "val1", "key2": "val2"}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat audio parts with outer extra", func(t *testing.T) { + base64Audio1 := "dGVzdF9hdWRpb18x" + base64Audio2 := "dGVzdF9hdWRpb18y" + + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio1}}, Extra: map[string]any{"outer1": "val1"}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &base64Audio2}}, Extra: map[string]any{"outer2": "val2"}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + mergedBase64Audio := base64Audio1 + base64Audio2 + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{MessagePartCommon: MessagePartCommon{Base64Data: &mergedBase64Audio}}, Extra: map[string]any{"outer1": "val1", "outer2": "val2"}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat reasoning parts", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "First, "}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "I need to think."}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Here is my answer."}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "First, I need to think."}}, + {Type: ChatMessagePartTypeText, Text: "Here is my answer."}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat reasoning parts with signature", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "Step 1: "}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "analyze.", Signature: "sig_abc"}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: " Step 2: ", Signature: "sig_xyz"}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "conclude."}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "Step 1: analyze. Step 2: conclude.", Signature: "sig_xyz"}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat with streaming meta index grouping", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "Think "}, StreamingMeta: &MessageStreamingMeta{Index: 0}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "more."}, StreamingMeta: &MessageStreamingMeta{Index: 0}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Hello ", StreamingMeta: &MessageStreamingMeta{Index: 1}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "world!", StreamingMeta: &MessageStreamingMeta{Index: 1}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeReasoning, Reasoning: &MessageOutputReasoning{Text: "Think more."}, StreamingMeta: &MessageStreamingMeta{Index: 0}}, + {Type: ChatMessagePartTypeText, Text: "Hello world!", StreamingMeta: &MessageStreamingMeta{Index: 1}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat with different streaming meta index should not merge", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "First block ", StreamingMeta: &MessageStreamingMeta{Index: 0}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Second block ", StreamingMeta: &MessageStreamingMeta{Index: 1}}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "continues.", StreamingMeta: &MessageStreamingMeta{Index: 0}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "First block ", StreamingMeta: &MessageStreamingMeta{Index: 0}}, + {Type: ChatMessagePartTypeText, Text: "Second block ", StreamingMeta: &MessageStreamingMeta{Index: 1}}, + {Type: ChatMessagePartTypeText, Text: "continues.", StreamingMeta: &MessageStreamingMeta{Index: 0}}, + } + + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat empty parts", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{}, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + assert.Empty(t, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat single part no merge needed", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Single"}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Single"}, + } + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat multiple consecutive text parts", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "One "}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Two "}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Three "}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Four"}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "One Two Three Four"}, + } + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat without streaming meta should not merge with streaming meta parts", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "No meta "}, + }, + }, + { + Role: Assistant, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "With meta", StreamingMeta: &MessageStreamingMeta{Index: 0}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedContent := []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "No meta "}, + {Type: ChatMessagePartTypeText, Text: "With meta", StreamingMeta: &MessageStreamingMeta{Index: 0}}, + } + assert.Equal(t, expectedContent, mergedMsg.AssistantGenMultiContent) + }) + + t.Run("concat multi content (deprecated)", func(t *testing.T) { + msgs := []*Message{ + { + Role: Assistant, + MultiContent: []ChatMessagePart{ + {Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: "image1.jpg"}}, + }, + }, + { + Role: Assistant, + MultiContent: []ChatMessagePart{ + {Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: "image2.jpg"}}, + }, + }, + } + + mergedMsg, err := ConcatMessages(msgs) + assert.NoError(t, err) + + expectedMultiContent := []ChatMessagePart{ + {Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: "image1.jpg"}}, + {Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: "image2.jpg"}}, + } + + assert.Equal(t, expectedMultiContent, mergedMsg.MultiContent) + }) +} + +func TestConcatToolCalls(t *testing.T) { + t.Run("atomic_field_in_first_chunk", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + }, + }, + { + Index: generic.PtrOf(0), + Function: FunctionCall{ + Arguments: "call me please", + }, + }, + } + + expectedToolCall := ToolCall{ + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + Arguments: "call me please", + }, + } + + tc, err := concatToolCalls(givenToolCalls) + assert.NoError(t, err) + assert.Len(t, tc, 1) + assert.EqualValues(t, expectedToolCall, tc[0]) + }) + + t.Run("atomic_field_in_every_chunk", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + Arguments: "call me please", + }, + }, + } + + expectedToolCall := ToolCall{ + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + Arguments: "call me please", + }, + } + + tc, err := concatToolCalls(givenToolCalls) + assert.NoError(t, err) + assert.Len(t, tc, 1) + assert.EqualValues(t, expectedToolCall, tc[0]) + }) + + t.Run("atomic_field_in_interval", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + { + Index: generic.PtrOf(0), + ID: "", + Type: "function", + Function: FunctionCall{ + Name: "", + Arguments: "call me please", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "", + Function: FunctionCall{ + Name: "", + Arguments: "", + }, + }, + } + + expectedToolCall := ToolCall{ + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "", + Arguments: "call me please", + }, + } + + tc, err := concatToolCalls(givenToolCalls) + assert.NoError(t, err) + assert.Len(t, tc, 1) + assert.EqualValues(t, expectedToolCall, tc[0]) + }) + + t.Run("different_tool_id", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id_1", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + Arguments: "call me please", + }, + }, + } + + _, err := concatToolCalls(givenToolCalls) + assert.ErrorContains(t, err, "cannot concat ToolCalls with different tool id") + }) + + t.Run("different_tool_type", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function_1", + Function: FunctionCall{ + Name: "tool_name", + Arguments: "call me please", + }, + }, + } + + _, err := concatToolCalls(givenToolCalls) + assert.ErrorContains(t, err, "cannot concat ToolCalls with different tool type") + }) + + t.Run("different_tool_name", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "tool_name_1", + Arguments: "call me please", + }, + }, + } + + _, err := concatToolCalls(givenToolCalls) + assert.ErrorContains(t, err, "cannot concat ToolCalls with different tool name") + }) + + t.Run("multi_tool_call", func(t *testing.T) { + givenToolCalls := []ToolCall{ + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + { + Index: generic.PtrOf(0), + ID: "", + Type: "function", + Function: FunctionCall{ + Name: "", + Arguments: "call me please", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "", + Function: FunctionCall{ + Name: "", + Arguments: "", + }, + }, + { + Index: generic.PtrOf(1), + ID: "tool_call_id", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + { + Index: generic.PtrOf(1), + ID: "", + Type: "function", + Function: FunctionCall{ + Name: "", + Arguments: "call me please", + }, + }, + { + Index: generic.PtrOf(1), + ID: "tool_call_id", + Type: "", + Function: FunctionCall{ + Name: "", + Arguments: "", + }, + }, + { + Index: nil, + ID: "22", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + { + Index: nil, + ID: "44", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + } + + expectedToolCall := []ToolCall{ + { + Index: nil, + ID: "22", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + { + Index: nil, + ID: "44", + Type: "", + Function: FunctionCall{ + Name: "", + }, + }, + { + Index: generic.PtrOf(0), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "", + Arguments: "call me please", + }, + }, + { + Index: generic.PtrOf(1), + ID: "tool_call_id", + Type: "function", + Function: FunctionCall{ + Name: "", + Arguments: "call me please", + }, + }, + } + + tc, err := concatToolCalls(givenToolCalls) + assert.NoError(t, err) + assert.EqualValues(t, expectedToolCall, tc) + }) +} + +func TestFormatMultiContent(t *testing.T) { + vs := map[string]any{ + "name": "eino", + "url": "https://example.com/img.png", + "id": "42", + } + + t.Run("empty input", func(t *testing.T) { + out, err := formatMultiContent(nil, vs, FString) + assert.NoError(t, err) + assert.Equal(t, []ChatMessagePart{}, out) + }) + + t.Run("render text and urls with FString", func(t *testing.T) { + in := []ChatMessagePart{ + {Type: ChatMessagePartTypeText, Text: "hello {name}"}, + {Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: "{url}"}}, + {Type: ChatMessagePartTypeAudioURL, AudioURL: &ChatMessageAudioURL{URL: "http://audio/{id}.wav"}}, + {Type: ChatMessagePartTypeVideoURL, VideoURL: &ChatMessageVideoURL{URL: "http://video/{id}.mp4"}}, + {Type: ChatMessagePartTypeFileURL, FileURL: &ChatMessageFileURL{URL: "http://file/{id}.txt"}}, + } + + out, err := formatMultiContent(in, vs, FString) + assert.NoError(t, err) + if assert.Len(t, out, len(in)) { + assert.Equal(t, "hello eino", out[0].Text) + assert.Equal(t, "https://example.com/img.png", out[1].ImageURL.URL) + assert.Equal(t, "http://audio/42.wav", out[2].AudioURL.URL) + assert.Equal(t, "http://video/42.mp4", out[3].VideoURL.URL) + assert.Equal(t, "http://file/42.txt", out[4].FileURL.URL) + } + }) + + t.Run("nil nested pointer should be skipped", func(t *testing.T) { + in := []ChatMessagePart{ + {Type: ChatMessagePartTypeImageURL, ImageURL: nil}, + {Type: ChatMessagePartTypeAudioURL, AudioURL: nil}, + {Type: ChatMessagePartTypeVideoURL, VideoURL: nil}, + {Type: ChatMessagePartTypeFileURL, FileURL: nil}, + } + out, err := formatMultiContent(in, vs, FString) + assert.NoError(t, err) + assert.Equal(t, in, out) + }) + + t.Run("missing var should error in GoTemplate", func(t *testing.T) { + in := []ChatMessagePart{{Type: ChatMessagePartTypeText, Text: "hi {{.who}}"}} + _, err := formatMultiContent(in, map[string]any{"name": "x"}, GoTemplate) + assert.Error(t, err) + }) + +} + +func TestFormatUserInputMultiContent(t *testing.T) { + makeStrPtr := func(s string) *string { return &s } + + vs := map[string]any{ + "x": "world", + "img": "https://example.com/i.png", + "b64": "YmFzZTY0", + "aid": "99", + "vid": "77", + "file": "abc", + } + + t.Run("empty input", func(t *testing.T) { + out, err := formatUserInputMultiContent(nil, vs, FString) + assert.NoError(t, err) + assert.Equal(t, []MessageInputPart{}, out) + }) + + t.Run("render text and both URL/Base64 for each type", func(t *testing.T) { + in := []MessageInputPart{ + {Type: ChatMessagePartTypeText, Text: "hello {x}"}, + {Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{MessagePartCommon: MessagePartCommon{URL: makeStrPtr("{img}"), Base64Data: makeStrPtr("{b64}")}}}, + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageInputAudio{MessagePartCommon: MessagePartCommon{URL: makeStrPtr("http://a/{aid}.wav"), Base64Data: makeStrPtr("{b64}")}}}, + {Type: ChatMessagePartTypeVideoURL, Video: &MessageInputVideo{MessagePartCommon: MessagePartCommon{URL: makeStrPtr("http://v/{vid}.mp4"), Base64Data: makeStrPtr("{b64}")}}}, + {Type: ChatMessagePartTypeFileURL, File: &MessageInputFile{MessagePartCommon: MessagePartCommon{URL: makeStrPtr("/f/{file}.txt"), Base64Data: makeStrPtr("{b64}")}}}, + } + + out, err := formatUserInputMultiContent(in, vs, FString) + assert.NoError(t, err) + if assert.Len(t, out, len(in)) { + assert.Equal(t, "hello world", out[0].Text) + assert.Equal(t, "https://example.com/i.png", *out[1].Image.URL) + assert.Equal(t, "YmFzZTY0", *out[1].Image.Base64Data) + assert.Equal(t, "http://a/99.wav", *out[2].Audio.URL) + assert.Equal(t, "YmFzZTY0", *out[2].Audio.Base64Data) + assert.Equal(t, "http://v/77.mp4", *out[3].Video.URL) + assert.Equal(t, "YmFzZTY0", *out[3].Video.Base64Data) + assert.Equal(t, "/f/abc.txt", *out[4].File.URL) + assert.Equal(t, "YmFzZTY0", *out[4].File.Base64Data) + } + }) + + t.Run("empty string pointer should not be formatted", func(t *testing.T) { + empty := "" + in := []MessageInputPart{ + {Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{MessagePartCommon: MessagePartCommon{URL: &empty, Base64Data: &empty}}}, + } + out, err := formatUserInputMultiContent(in, vs, FString) + assert.NoError(t, err) + if assert.Len(t, out, 1) { + assert.NotNil(t, out[0].Image.URL) + assert.NotNil(t, out[0].Image.Base64Data) + assert.Equal(t, "", *out[0].Image.URL) + assert.Equal(t, "", *out[0].Image.Base64Data) + } + }) +} + +func TestConcatToolResults(t *testing.T) { + t.Run("empty_chunks", func(t *testing.T) { + result, err := ConcatToolResults([]*ToolResult{}) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Empty(t, result.Parts) + }) + + t.Run("nil_chunks", func(t *testing.T) { + result, err := ConcatToolResults([]*ToolResult{nil, nil}) + assert.NoError(t, err) + assert.NotNil(t, result) + assert.Empty(t, result.Parts) + }) + + t.Run("single_text_part", func(t *testing.T) { + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Hello World"}, + }, + }, + } + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "Hello World", result.Parts[0].Text) + }) + + t.Run("multiple_text_parts_merge", func(t *testing.T) { + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Hello "}, + }, + }, + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "World"}, + }, + }, + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "!"}, + }, + }, + } + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "Hello World!", result.Parts[0].Text) + }) + + t.Run("multiple_text_parts_merge_with_extra", func(t *testing.T) { + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Hello ", Extra: map[string]any{"k1": "v1"}}, + {Type: ToolPartTypeText, Text: "World", Extra: map[string]any{"k2": "v2"}}, + }, + }, + } + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, "Hello World", result.Parts[0].Text) + assert.Equal(t, map[string]any{"k1": "v1", "k2": "v2"}, result.Parts[0].Extra) + }) + + t.Run("multiple_text_parts_merge_with_single_extra", func(t *testing.T) { + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Hello ", Extra: map[string]any{"k1": "v1"}}, + {Type: ToolPartTypeText, Text: "World"}, + }, + }, + } + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, "Hello World", result.Parts[0].Text) + assert.Equal(t, map[string]any{"k1": "v1"}, result.Parts[0].Extra) + }) + + t.Run("cross_chunk_audio_conflict_error", func(t *testing.T) { + base64Data1 := "YXVkaW8x" + base64Data2 := "YXVkaW8y" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data1, + MIMEType: "audio/wav", + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data2, + MIMEType: "audio/wav", + }, + }, + }, + }, + }, + } + + _, err := ConcatToolResults(chunks) + assert.Error(t, err) + assert.Contains(t, err.Error(), "conflicting") + assert.Contains(t, err.Error(), "audio") + }) + + t.Run("mixed_types_no_merge", func(t *testing.T) { + imageURL := "https://example.com/image.png" + videoURL := "https://example.com/video.mp4" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Text part"}, + { + Type: ToolPartTypeImage, + Image: &ToolOutputImage{ + MessagePartCommon: MessagePartCommon{ + URL: &imageURL, + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeVideo, + Video: &ToolOutputVideo{ + MessagePartCommon: MessagePartCommon{ + URL: &videoURL, + }, + }, + }, + }, + }, + } + + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 3) + assert.Equal(t, ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, ToolPartTypeImage, result.Parts[1].Type) + assert.Equal(t, ToolPartTypeVideo, result.Parts[2].Type) + }) + + t.Run("mixed_text_and_single_audio", func(t *testing.T) { + base64Data1 := "YXVkaW8x" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Part 1 "}, + {Type: ToolPartTypeText, Text: "Part 2"}, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data1, + MIMEType: "audio/wav", + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: " Part 3"}, + }, + }, + } + + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 3) + + assert.Equal(t, ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "Part 1 Part 2", result.Parts[0].Text) + + assert.Equal(t, ToolPartTypeAudio, result.Parts[1].Type) + assert.NotNil(t, result.Parts[1].Audio) + assert.NotNil(t, result.Parts[1].Audio.Base64Data) + assert.Equal(t, "YXVkaW8x", *result.Parts[1].Audio.Base64Data) + + assert.Equal(t, ToolPartTypeText, result.Parts[2].Type) + assert.Equal(t, " Part 3", result.Parts[2].Text) + }) + + t.Run("cross_chunk_audio_url_and_base64_conflict_error", func(t *testing.T) { + audioURL := "https://example.com/audio.wav" + base64Data := "YXVkaW8x" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + URL: &audioURL, + MIMEType: "audio/wav", + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data, + MIMEType: "audio/wav", + }, + }, + }, + }, + }, + } + + _, err := ConcatToolResults(chunks) + assert.Error(t, err) + assert.Contains(t, err.Error(), "conflicting") + assert.Contains(t, err.Error(), "audio") + }) + + t.Run("single_audio_with_extra_fields", func(t *testing.T) { + base64Data1 := "YXVkaW8x" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data1, + MIMEType: "audio/wav", + Extra: map[string]any{ + "key1": "value1", + }, + }, + }, + }, + }, + }, + } + + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, ToolPartTypeAudio, result.Parts[0].Type) + assert.NotNil(t, result.Parts[0].Audio) + assert.NotNil(t, result.Parts[0].Audio.Base64Data) + assert.Equal(t, "YXVkaW8x", *result.Parts[0].Audio.Base64Data) + assert.NotNil(t, result.Parts[0].Audio.Extra) + assert.Equal(t, "value1", result.Parts[0].Audio.Extra["key1"]) + }) + + t.Run("cross_chunk_image_conflict_error", func(t *testing.T) { + imageURL1 := "https://example.com/image1.png" + imageURL2 := "https://example.com/image2.png" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeImage, + Image: &ToolOutputImage{ + MessagePartCommon: MessagePartCommon{ + URL: &imageURL1, + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeImage, + Image: &ToolOutputImage{ + MessagePartCommon: MessagePartCommon{ + URL: &imageURL2, + }, + }, + }, + }, + }, + } + + _, err := ConcatToolResults(chunks) + assert.Error(t, err) + assert.Contains(t, err.Error(), "conflicting") + assert.Contains(t, err.Error(), "image") + }) + + t.Run("cross_chunk_video_conflict_error", func(t *testing.T) { + videoURL1 := "https://example.com/video1.mp4" + videoURL2 := "https://example.com/video2.mp4" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeVideo, + Video: &ToolOutputVideo{ + MessagePartCommon: MessagePartCommon{ + URL: &videoURL1, + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeVideo, + Video: &ToolOutputVideo{ + MessagePartCommon: MessagePartCommon{ + URL: &videoURL2, + }, + }, + }, + }, + }, + } + + _, err := ConcatToolResults(chunks) + assert.Error(t, err) + assert.Contains(t, err.Error(), "conflicting") + assert.Contains(t, err.Error(), "video") + }) + + t.Run("cross_chunk_file_conflict_error", func(t *testing.T) { + fileURL1 := "https://example.com/file1.pdf" + fileURL2 := "https://example.com/file2.pdf" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeFile, + File: &ToolOutputFile{ + MessagePartCommon: MessagePartCommon{ + URL: &fileURL1, + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeFile, + File: &ToolOutputFile{ + MessagePartCommon: MessagePartCommon{ + URL: &fileURL2, + }, + }, + }, + }, + }, + } + + _, err := ConcatToolResults(chunks) + assert.Error(t, err) + assert.Contains(t, err.Error(), "conflicting") + assert.Contains(t, err.Error(), "file") + }) + + t.Run("same_chunk_text_merged", func(t *testing.T) { + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + {Type: ToolPartTypeText, Text: "Hello "}, + {Type: ToolPartTypeText, Text: "World"}, + }, + }, + } + + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 1) + assert.Equal(t, ToolPartTypeText, result.Parts[0].Type) + assert.Equal(t, "Hello World", result.Parts[0].Text) + }) + + t.Run("different_non_text_types_across_chunks_allowed", func(t *testing.T) { + imageURL := "https://example.com/image.png" + videoURL := "https://example.com/video.mp4" + base64Audio := "YXVkaW8x" + + chunks := []*ToolResult{ + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeImage, + Image: &ToolOutputImage{ + MessagePartCommon: MessagePartCommon{ + URL: &imageURL, + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeVideo, + Video: &ToolOutputVideo{ + MessagePartCommon: MessagePartCommon{ + URL: &videoURL, + }, + }, + }, + }, + }, + { + Parts: []ToolOutputPart{ + { + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Audio, + MIMEType: "audio/wav", + }, + }, + }, + }, + }, + } + + result, err := ConcatToolResults(chunks) + assert.NoError(t, err) + assert.Len(t, result.Parts, 3) + assert.Equal(t, ToolPartTypeImage, result.Parts[0].Type) + assert.Equal(t, ToolPartTypeVideo, result.Parts[1].Type) + assert.Equal(t, ToolPartTypeAudio, result.Parts[2].Type) + }) +} + +func TestMessageString(t *testing.T) { + t.Run("basic message", func(t *testing.T) { + msg := &Message{ + Role: User, + Content: "Hello, world!", + } + result := msg.String() + assert.Contains(t, result, "user: Hello, world!") + }) + + t.Run("message with UserInputMultiContent", func(t *testing.T) { + imageURL := "https://example.com/image.png" + msg := &Message{ + Role: User, + Content: "", + UserInputMultiContent: []MessageInputPart{ + {Type: ChatMessagePartTypeText, Text: "Describe this image:"}, + {Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{ + MessagePartCommon: MessagePartCommon{URL: &imageURL}, + }}, + }, + } + result := msg.String() + assert.Contains(t, result, "user_input_multi_content:") + assert.Contains(t, result, "[0] text: Describe this image:") + assert.Contains(t, result, "[1] image: url=https://example.com/image.png") + }) + + t.Run("message with AssistantGenMultiContent", func(t *testing.T) { + base64Data := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + msg := &Message{ + Role: Assistant, + Content: "", + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Here is the generated image:"}, + {Type: ChatMessagePartTypeImageURL, Image: &MessageOutputImage{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data, + MIMEType: "image/png", + }, + }}, + }, + } + result := msg.String() + assert.Contains(t, result, "assistant_gen_multi_content:") + assert.Contains(t, result, "[0] text: Here is the generated image:") + assert.Contains(t, result, "[1] image: base64[") + assert.Contains(t, result, "mime=image/png") + }) + + t.Run("message with MultiContent (deprecated)", func(t *testing.T) { + msg := &Message{ + Role: User, + Content: "", + MultiContent: []ChatMessagePart{ + {Type: ChatMessagePartTypeText, Text: "What is this?"}, + {Type: ChatMessagePartTypeImageURL, ImageURL: &ChatMessageImageURL{URL: "https://example.com/photo.jpg"}}, + }, + } + result := msg.String() + assert.Contains(t, result, "multi_content:") + assert.Contains(t, result, "[0] text: What is this?") + assert.Contains(t, result, "[1] image_url: https://example.com/photo.jpg") + }) + + t.Run("message with ToolCalls", func(t *testing.T) { + idx := 0 + msg := &Message{ + Role: Assistant, + Content: "", + ToolCalls: []ToolCall{ + { + Index: &idx, + ID: "call_123", + Type: "function", + Function: FunctionCall{ + Name: "get_weather", + Arguments: `{"location": "Beijing"}`, + }, + }, + }, + } + result := msg.String() + assert.Contains(t, result, "tool_calls:") + assert.Contains(t, result, "index[0]:") + assert.Contains(t, result, "get_weather") + }) + + t.Run("tool message", func(t *testing.T) { + msg := &Message{ + Role: Tool, + Content: `{"temperature": 25}`, + ToolCallID: "call_123", + ToolName: "get_weather", + } + result := msg.String() + assert.Contains(t, result, "tool: {\"temperature\": 25}") + assert.Contains(t, result, "tool_call_id: call_123") + assert.Contains(t, result, "tool_call_name: get_weather") + }) + + t.Run("message with reasoning content", func(t *testing.T) { + msg := &Message{ + Role: Assistant, + Content: "The answer is 42.", + ReasoningContent: "Let me think about this step by step...", + } + result := msg.String() + assert.Contains(t, result, "reasoning content:") + assert.Contains(t, result, "Let me think about this step by step...") + }) + + t.Run("message with response meta", func(t *testing.T) { + msg := &Message{ + Role: Assistant, + Content: "Hello!", + ResponseMeta: &ResponseMeta{ + FinishReason: "stop", + Usage: &TokenUsage{ + PromptTokens: 10, + CompletionTokens: 5, + TotalTokens: 15, + }, + }, + } + result := msg.String() + assert.Contains(t, result, "finish_reason: stop") + assert.Contains(t, result, "usage:") + }) + + t.Run("message with audio input", func(t *testing.T) { + audioURL := "https://example.com/audio.wav" + msg := &Message{ + Role: User, + UserInputMultiContent: []MessageInputPart{ + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageInputAudio{ + MessagePartCommon: MessagePartCommon{URL: &audioURL}, + }}, + }, + } + result := msg.String() + assert.Contains(t, result, "[0] audio: url=https://example.com/audio.wav") + }) + + t.Run("message with video input", func(t *testing.T) { + videoURL := "https://example.com/video.mp4" + msg := &Message{ + Role: User, + UserInputMultiContent: []MessageInputPart{ + {Type: ChatMessagePartTypeVideoURL, Video: &MessageInputVideo{ + MessagePartCommon: MessagePartCommon{URL: &videoURL}, + }}, + }, + } + result := msg.String() + assert.Contains(t, result, "[0] video: url=https://example.com/video.mp4") + }) + + t.Run("message with file input", func(t *testing.T) { + fileURL := "https://example.com/document.pdf" + msg := &Message{ + Role: User, + UserInputMultiContent: []MessageInputPart{ + {Type: ChatMessagePartTypeFileURL, File: &MessageInputFile{ + MessagePartCommon: MessagePartCommon{URL: &fileURL}, + }}, + }, + } + result := msg.String() + assert.Contains(t, result, "[0] file: url=https://example.com/document.pdf") + }) + + t.Run("nil media parts", func(t *testing.T) { + msg := &Message{ + Role: User, + UserInputMultiContent: []MessageInputPart{ + {Type: ChatMessagePartTypeImageURL, Image: nil}, + }, + } + result := msg.String() + assert.Contains(t, result, "[0] image: ") + }) + + t.Run("combined multi-content types", func(t *testing.T) { + imageURL := "https://example.com/image.png" + base64Audio := "YXVkaW9kYXRh" + msg := &Message{ + Role: User, + Content: "Main content", + UserInputMultiContent: []MessageInputPart{ + {Type: ChatMessagePartTypeText, Text: "User input text"}, + {Type: ChatMessagePartTypeImageURL, Image: &MessageInputImage{ + MessagePartCommon: MessagePartCommon{URL: &imageURL}, + }}, + }, + AssistantGenMultiContent: []MessageOutputPart{ + {Type: ChatMessagePartTypeText, Text: "Assistant output text"}, + {Type: ChatMessagePartTypeAudioURL, Audio: &MessageOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Audio, + MIMEType: "audio/wav", + }, + }}, + }, + } + result := msg.String() + assert.Contains(t, result, "user: Main content") + assert.Contains(t, result, "user_input_multi_content:") + assert.Contains(t, result, "assistant_gen_multi_content:") + }) +} + +func TestConvToolOutputPartToMessageInputPart(t *testing.T) { + t.Run("text part", func(t *testing.T) { + toolPart := ToolOutputPart{ + Type: ToolPartTypeText, + Text: "test text", + Extra: map[string]any{"key": "value"}, + } + result, err := convToolOutputPartToMessageInputPart(toolPart) + assert.NoError(t, err) + assert.Equal(t, ChatMessagePartTypeText, result.Type) + assert.Equal(t, "test text", result.Text) + assert.Equal(t, map[string]any{"key": "value"}, result.Extra) + }) + + t.Run("image part", func(t *testing.T) { + url := "https://example.com/image.png" + toolPart := ToolOutputPart{ + Type: ToolPartTypeImage, + Image: &ToolOutputImage{ + MessagePartCommon: MessagePartCommon{ + URL: &url, + MIMEType: "image/png", + }, + }, + Extra: map[string]any{"img_key": "img_value"}, + } + result, err := convToolOutputPartToMessageInputPart(toolPart) + assert.NoError(t, err) + assert.Equal(t, ChatMessagePartTypeImageURL, result.Type) + assert.NotNil(t, result.Image) + assert.Equal(t, url, *result.Image.URL) + assert.Equal(t, "image/png", result.Image.MIMEType) + assert.Equal(t, map[string]any{"img_key": "img_value"}, result.Extra) + }) + + t.Run("image part nil content", func(t *testing.T) { + toolPart := ToolOutputPart{ + Type: ToolPartTypeImage, + Image: nil, + } + result, err := convToolOutputPartToMessageInputPart(toolPart) + assert.Error(t, err) + assert.ErrorContains(t, err, "image content is nil") + assert.Equal(t, MessageInputPart{}, result) + }) + + t.Run("audio part", func(t *testing.T) { + base64Data := "dGVzdF9hdWRpbw==" + toolPart := ToolOutputPart{ + Type: ToolPartTypeAudio, + Audio: &ToolOutputAudio{ + MessagePartCommon: MessagePartCommon{ + Base64Data: &base64Data, + MIMEType: "audio/wav", + }, + }, + } + result, err := convToolOutputPartToMessageInputPart(toolPart) + assert.NoError(t, err) + assert.Equal(t, ChatMessagePartTypeAudioURL, result.Type) + assert.NotNil(t, result.Audio) + assert.Equal(t, base64Data, *result.Audio.Base64Data) + assert.Equal(t, "audio/wav", result.Audio.MIMEType) + }) + + t.Run("audio part nil content", func(t *testing.T) { + toolPart := ToolOutputPart{ + Type: ToolPartTypeAudio, + Audio: nil, + } + _, err := convToolOutputPartToMessageInputPart(toolPart) + assert.Error(t, err) + assert.ErrorContains(t, err, "audio content is nil") + }) + + t.Run("video part", func(t *testing.T) { + url := "https://example.com/video.mp4" + toolPart := ToolOutputPart{ + Type: ToolPartTypeVideo, + Video: &ToolOutputVideo{ + MessagePartCommon: MessagePartCommon{ + URL: &url, + MIMEType: "video/mp4", + }, + }, + } + result, err := convToolOutputPartToMessageInputPart(toolPart) + assert.NoError(t, err) + assert.Equal(t, ChatMessagePartTypeVideoURL, result.Type) + assert.NotNil(t, result.Video) + assert.Equal(t, url, *result.Video.URL) + assert.Equal(t, "video/mp4", result.Video.MIMEType) + }) + + t.Run("video part nil content", func(t *testing.T) { + toolPart := ToolOutputPart{ + Type: ToolPartTypeVideo, + Video: nil, + } + _, err := convToolOutputPartToMessageInputPart(toolPart) + assert.Error(t, err) + assert.ErrorContains(t, err, "video content is nil") + }) + + t.Run("file part", func(t *testing.T) { + url := "https://example.com/file.pdf" + toolPart := ToolOutputPart{ + Type: ToolPartTypeFile, + File: &ToolOutputFile{ + MessagePartCommon: MessagePartCommon{ + URL: &url, + MIMEType: "application/pdf", + }, + }, + Extra: map[string]any{"file_key": "file_value"}, + } + result, err := convToolOutputPartToMessageInputPart(toolPart) + assert.NoError(t, err) + assert.Equal(t, ChatMessagePartTypeFileURL, result.Type) + assert.NotNil(t, result.File) + assert.Equal(t, url, *result.File.URL) + assert.Equal(t, "application/pdf", result.File.MIMEType) + assert.Equal(t, map[string]any{"file_key": "file_value"}, result.Extra) + }) + + t.Run("file part nil content", func(t *testing.T) { + toolPart := ToolOutputPart{ + Type: ToolPartTypeFile, + File: nil, + } + _, err := convToolOutputPartToMessageInputPart(toolPart) + assert.Error(t, err) + assert.ErrorContains(t, err, "file content is nil") + }) + + t.Run("unknown type", func(t *testing.T) { + toolPart := ToolOutputPart{ + Type: "unknown_type", + } + _, err := convToolOutputPartToMessageInputPart(toolPart) + assert.Error(t, err) + assert.ErrorContains(t, err, "unknown tool part type") + }) +} diff --git a/schema/openai/consts.go b/schema/openai/consts.go new file mode 100644 index 0000000..5958cef --- /dev/null +++ b/schema/openai/consts.go @@ -0,0 +1,95 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package openai defines constants for openai. +package openai + +type TextAnnotationType string + +const ( + TextAnnotationTypeFileCitation TextAnnotationType = "file_citation" + TextAnnotationTypeURLCitation TextAnnotationType = "url_citation" + TextAnnotationTypeContainerFileCitation TextAnnotationType = "container_file_citation" + TextAnnotationTypeFilePath TextAnnotationType = "file_path" +) + +type ReasoningEffort string + +const ( + ReasoningEffortMinimal ReasoningEffort = "minimal" + ReasoningEffortLow ReasoningEffort = "low" + ReasoningEffortMedium ReasoningEffort = "medium" + ReasoningEffortHigh ReasoningEffort = "high" +) + +type ReasoningSummary string + +const ( + ReasoningSummaryAuto ReasoningSummary = "auto" + ReasoningSummaryConcise ReasoningSummary = "concise" + ReasoningSummaryDetailed ReasoningSummary = "detailed" +) + +type ServiceTier string + +const ( + ServiceTierAuto ServiceTier = "auto" + ServiceTierDefault ServiceTier = "default" + ServiceTierFlex ServiceTier = "flex" + ServiceTierScale ServiceTier = "scale" + ServiceTierPriority ServiceTier = "priority" +) + +type PromptCacheRetention string + +const ( + PromptCacheRetentionInMemory PromptCacheRetention = "in-memory" + PromptCacheRetention24h PromptCacheRetention = "24h" +) + +type ResponseStatus string + +const ( + ResponseStatusCompleted ResponseStatus = "completed" + ResponseStatusFailed ResponseStatus = "failed" + ResponseStatusInProgress ResponseStatus = "in_progress" + ResponseStatusCancelled ResponseStatus = "cancelled" + ResponseStatusQueued ResponseStatus = "queued" + ResponseStatusIncomplete ResponseStatus = "incomplete" +) + +type ResponseErrorCode string + +const ( + ResponseErrorCodeServerError ResponseErrorCode = "server_error" + ResponseErrorCodeRateLimitExceeded ResponseErrorCode = "rate_limit_exceeded" + ResponseErrorCodeInvalidPrompt ResponseErrorCode = "invalid_prompt" + ResponseErrorCodeVectorStoreTimeout ResponseErrorCode = "vector_store_timeout" + ResponseErrorCodeInvalidImage ResponseErrorCode = "invalid_image" + ResponseErrorCodeInvalidImageFormat ResponseErrorCode = "invalid_image_format" + ResponseErrorCodeInvalidBase64Image ResponseErrorCode = "invalid_base64_image" + ResponseErrorCodeInvalidImageURL ResponseErrorCode = "invalid_image_url" + ResponseErrorCodeImageTooLarge ResponseErrorCode = "image_too_large" + ResponseErrorCodeImageTooSmall ResponseErrorCode = "image_too_small" + ResponseErrorCodeImageParseError ResponseErrorCode = "image_parse_error" + ResponseErrorCodeImageContentPolicyViolation ResponseErrorCode = "image_content_policy_violation" + ResponseErrorCodeInvalidImageMode ResponseErrorCode = "invalid_image_mode" + ResponseErrorCodeImageFileTooLarge ResponseErrorCode = "image_file_too_large" + ResponseErrorCodeUnsupportedImageMediaType ResponseErrorCode = "unsupported_image_media_type" + ResponseErrorCodeEmptyImageFile ResponseErrorCode = "empty_image_file" + ResponseErrorCodeFailedToDownloadImage ResponseErrorCode = "failed_to_download_image" + ResponseErrorCodeImageFileNotFound ResponseErrorCode = "image_file_not_found" +) diff --git a/schema/openai/extension.go b/schema/openai/extension.go new file mode 100644 index 0000000..3bee170 --- /dev/null +++ b/schema/openai/extension.go @@ -0,0 +1,283 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package openai + +import ( + "fmt" + "sort" +) + +type ResponseMetaExtension struct { + ID string `json:"id,omitempty"` + Status ResponseStatus `json:"status,omitempty"` + Error *ResponseError `json:"error,omitempty"` + IncompleteDetails *IncompleteDetails `json:"incomplete_details,omitempty"` + PreviousResponseID string `json:"previous_response_id,omitempty"` + Reasoning *Reasoning `json:"reasoning,omitempty"` + ServiceTier ServiceTier `json:"service_tier,omitempty"` + CreatedAt int64 `json:"created_at,omitempty"` + PromptCacheRetention PromptCacheRetention `json:"prompt_cache_retention,omitempty"` +} + +type AssistantGenTextExtension struct { + Refusal *OutputRefusal `json:"refusal,omitempty"` + Annotations []*TextAnnotation `json:"annotations,omitempty"` +} + +type ReasoningExtension struct { + // Content is the reasoning text content. + Content []*ReasoningContent `json:"content,omitempty"` +} + +type ReasoningContent struct { + // Index specifies the index position of this content in the final response. + // Only available in streaming response. + Index *int `json:"index,omitempty"` + + Text string `json:"text,omitempty"` +} + +type ResponseError struct { + Code ResponseErrorCode `json:"code,omitempty"` + Message string `json:"message,omitempty"` +} + +type IncompleteDetails struct { + Reason string `json:"reason,omitempty"` +} + +type Reasoning struct { + Effort ReasoningEffort `json:"effort,omitempty"` + Summary ReasoningSummary `json:"summary,omitempty"` +} + +type OutputRefusal struct { + Reason string `json:"reason,omitempty"` +} + +type TextAnnotation struct { + // Index specifies the index position of this annotation in the final response. + // Only available in streaming response. + Index int `json:"index,omitempty"` + + Type TextAnnotationType `json:"type,omitempty"` + + FileCitation *TextAnnotationFileCitation `json:"file_citation,omitempty"` + URLCitation *TextAnnotationURLCitation `json:"url_citation,omitempty"` + ContainerFileCitation *TextAnnotationContainerFileCitation `json:"container_file_citation,omitempty"` + FilePath *TextAnnotationFilePath `json:"file_path,omitempty"` +} + +type TextAnnotationFileCitation struct { + // The ID of the file. + FileID string `json:"file_id,omitempty"` + // The filename of the file cited. + Filename string `json:"filename,omitempty"` + + // The index of the file in the list of files. + Index int `json:"index,omitempty"` +} + +type TextAnnotationURLCitation struct { + // The title of the web resource. + Title string `json:"title,omitempty"` + // The URL of the web resource. + URL string `json:"url,omitempty"` + + // The index of the first character of the URL citation in the message. + StartIndex int `json:"start_index,omitempty"` + // The index of the last character of the URL citation in the message. + EndIndex int `json:"end_index,omitempty"` +} + +type TextAnnotationContainerFileCitation struct { + // The ID of the container file. + ContainerID string `json:"container_id,omitempty"` + + // The ID of the file. + FileID string `json:"file_id,omitempty"` + // The filename of the container file cited. + Filename string `json:"filename,omitempty"` + + // The index of the first character of the container file citation in the message. + StartIndex int `json:"start_index,omitempty"` + // The index of the last character of the container file citation in the message. + EndIndex int `json:"end_index,omitempty"` +} + +type TextAnnotationFilePath struct { + // The ID of the file. + FileID string `json:"file_id,omitempty"` + + // The index of the file in the list of files. + Index int `json:"index,omitempty"` +} + +// ConcatAssistantGenTextExtensions concatenates multiple AssistantGenTextExtension chunks into a single one. +func ConcatAssistantGenTextExtensions(chunks []*AssistantGenTextExtension) (*AssistantGenTextExtension, error) { + if len(chunks) == 0 { + return nil, fmt.Errorf("no assistant generated text extension found") + } + + ret := &AssistantGenTextExtension{} + + var allAnnotations []*TextAnnotation + for _, ext := range chunks { + allAnnotations = append(allAnnotations, ext.Annotations...) + } + + var ( + indices []int + indexToAnnotation = map[int]*TextAnnotation{} + ) + + for _, an := range allAnnotations { + if an == nil { + continue + } + if indexToAnnotation[an.Index] == nil { + indexToAnnotation[an.Index] = an + indices = append(indices, an.Index) + } else { + return nil, fmt.Errorf("duplicate annotation index %d", an.Index) + } + } + + sort.Slice(indices, func(i, j int) bool { + return indices[i] < indices[j] + }) + + ret.Annotations = make([]*TextAnnotation, 0, len(indices)) + for _, idx := range indices { + an := *indexToAnnotation[idx] + an.Index = 0 // clear index + ret.Annotations = append(ret.Annotations, &an) + } + + for _, ext := range chunks { + if ext.Refusal == nil { + continue + } + if ret.Refusal == nil { + ret.Refusal = ext.Refusal + } else { + ret.Refusal.Reason += ext.Refusal.Reason + } + } + + return ret, nil +} + +// ConcatReasoningExtensions concatenates multiple ReasoningExtension chunks into a single one. +func ConcatReasoningExtensions(chunks []*ReasoningExtension) (*ReasoningExtension, error) { + if len(chunks) == 0 { + return nil, fmt.Errorf("no reasoning extension found") + } + + ret := &ReasoningExtension{} + + var ( + indices []int + indexToContent = map[int]*ReasoningContent{} + hasIndexed bool + hasUnindexed bool + ) + + for _, ext := range chunks { + if ext == nil { + continue + } + for _, c := range ext.Content { + if c == nil { + continue + } + + if c.Index == nil { + hasUnindexed = true + ret.Content = append(ret.Content, &ReasoningContent{Text: c.Text}) + continue + } + + hasIndexed = true + idx := *c.Index + if existing, ok := indexToContent[idx]; ok { + existing.Text += c.Text + } else { + indexToContent[idx] = &ReasoningContent{Text: c.Text} + indices = append(indices, idx) + } + } + } + + if hasIndexed && hasUnindexed { + return nil, fmt.Errorf("reasoning content chunks mix indexed and non-indexed content") + } + + if hasIndexed { + sort.Ints(indices) + ret.Content = make([]*ReasoningContent, 0, len(indices)) + for _, idx := range indices { + ret.Content = append(ret.Content, indexToContent[idx]) + } + } + + return ret, nil +} + +// ConcatResponseMetaExtensions concatenates multiple ResponseMetaExtension chunks into a single one. +func ConcatResponseMetaExtensions(chunks []*ResponseMetaExtension) (*ResponseMetaExtension, error) { + if len(chunks) == 0 { + return nil, fmt.Errorf("no response meta extension found") + } + if len(chunks) == 1 { + return chunks[0], nil + } + + ret := &ResponseMetaExtension{} + + for _, ext := range chunks { + if ext.ID != "" { + ret.ID = ext.ID + } + if ext.Status != "" { + ret.Status = ext.Status + } + if ext.Error != nil { + ret.Error = ext.Error + } + if ext.IncompleteDetails != nil { + ret.IncompleteDetails = ext.IncompleteDetails + } + if ext.PreviousResponseID != "" { + ret.PreviousResponseID = ext.PreviousResponseID + } + if ext.Reasoning != nil { + ret.Reasoning = ext.Reasoning + } + if ext.ServiceTier != "" { + ret.ServiceTier = ext.ServiceTier + } + if ext.CreatedAt != 0 { + ret.CreatedAt = ext.CreatedAt + } + if ext.PromptCacheRetention != "" { + ret.PromptCacheRetention = ext.PromptCacheRetention + } + } + + return ret, nil +} diff --git a/schema/openai/extension_test.go b/schema/openai/extension_test.go new file mode 100644 index 0000000..5ee57ca --- /dev/null +++ b/schema/openai/extension_test.go @@ -0,0 +1,268 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package openai + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/generic" +) + +func TestConcatResponseMetaExtensions(t *testing.T) { + t.Run("multiple extensions - takes last non-empty values", func(t *testing.T) { + err1 := &ResponseError{Code: "err1", Message: "msg1"} + incomplete := &IncompleteDetails{Reason: "max_tokens"} + + exts := []*ResponseMetaExtension{ + { + ID: "id_1", + Status: "in_progress", + Error: err1, + IncompleteDetails: nil, + }, + { + ID: "id_2", + Status: "", + Error: nil, + IncompleteDetails: nil, + }, + { + ID: "", + Status: "completed", + Error: nil, + IncompleteDetails: incomplete, + }, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "id_2", result.ID) + assert.Equal(t, ResponseStatus("completed"), result.Status) + assert.Equal(t, err1, result.Error) + assert.Equal(t, incomplete, result.IncompleteDetails) + }) + + t.Run("streaming scenario", func(t *testing.T) { + exts := []*ResponseMetaExtension{ + {ID: "chatcmpl_stream", Status: "", Error: nil, IncompleteDetails: nil}, + {ID: "", Status: ResponseStatus("in_progress"), Error: nil, IncompleteDetails: nil}, + {ID: "", Status: ResponseStatus("completed"), Error: nil, IncompleteDetails: nil}, + } + + result, err := ConcatResponseMetaExtensions(exts) + assert.NoError(t, err) + assert.Equal(t, "chatcmpl_stream", result.ID) + assert.Equal(t, ResponseStatus("completed"), result.Status) + }) +} + +func TestConcatAssistantGenTextExtensions(t *testing.T) { + t.Run("single extension with annotations", func(t *testing.T) { + ext := &AssistantGenTextExtension{ + Annotations: []*TextAnnotation{ + { + Index: 0, + Type: "file_citation", + FileCitation: &TextAnnotationFileCitation{ + FileID: "file_123", + Filename: "doc.pdf", + }, + }, + }, + } + + result, err := ConcatAssistantGenTextExtensions([]*AssistantGenTextExtension{ext}) + assert.NoError(t, err) + assert.Len(t, result.Annotations, 1) + assert.Equal(t, "file_123", result.Annotations[0].FileCitation.FileID) + }) + + t.Run("multiple extensions - merges annotations by index", func(t *testing.T) { + exts := []*AssistantGenTextExtension{ + { + Annotations: []*TextAnnotation{ + { + Index: 0, + Type: "file_citation", + FileCitation: &TextAnnotationFileCitation{ + FileID: "file_1", + }, + }, + }, + }, + { + Annotations: []*TextAnnotation{ + { + Index: 2, + Type: "url_citation", + URLCitation: &TextAnnotationURLCitation{ + URL: "https://example.com", + }, + }, + }, + }, + { + Annotations: []*TextAnnotation{ + { + Index: 1, + Type: "file_path", + FilePath: &TextAnnotationFilePath{ + FileID: "file_2", + }, + }, + }, + }, + } + + result, err := ConcatAssistantGenTextExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Annotations, 3) + assert.Equal(t, "file_1", result.Annotations[0].FileCitation.FileID) + assert.Equal(t, "file_2", result.Annotations[1].FilePath.FileID) + assert.Equal(t, "https://example.com", result.Annotations[2].URLCitation.URL) + }) + + t.Run("streaming scenario - annotations arrive in chunks", func(t *testing.T) { + exts := []*AssistantGenTextExtension{ + { + Annotations: []*TextAnnotation{ + {Index: 0, Type: "file_citation", FileCitation: &TextAnnotationFileCitation{FileID: "f1"}}, + }, + }, + { + Annotations: []*TextAnnotation{ + {Index: 1, Type: "url_citation", URLCitation: &TextAnnotationURLCitation{URL: "url1"}}, + }, + }, + { + Annotations: []*TextAnnotation{ + {Index: 2, Type: "file_path", FilePath: &TextAnnotationFilePath{FileID: "f2"}}, + }, + }, + } + + result, err := ConcatAssistantGenTextExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Annotations, 3) + assert.Equal(t, "f1", result.Annotations[0].FileCitation.FileID) + assert.Equal(t, "url1", result.Annotations[1].URLCitation.URL) + assert.Equal(t, "f2", result.Annotations[2].FilePath.FileID) + }) + + t.Run("multiple extensions - concatenates refusal reason", func(t *testing.T) { + ext1 := &AssistantGenTextExtension{Refusal: &OutputRefusal{Reason: "A"}} + ext2 := &AssistantGenTextExtension{Refusal: &OutputRefusal{Reason: "B"}} + + result, err := ConcatAssistantGenTextExtensions([]*AssistantGenTextExtension{ext1, ext2}) + assert.NoError(t, err) + assert.NotNil(t, result.Refusal) + assert.Equal(t, "AB", result.Refusal.Reason) + }) + + t.Run("duplicate index - error occurrence", func(t *testing.T) { + exts := []*AssistantGenTextExtension{ + { + Annotations: []*TextAnnotation{ + {Index: 0, Type: "file_citation", FileCitation: &TextAnnotationFileCitation{FileID: "first"}}, + }, + }, + { + Annotations: []*TextAnnotation{ + {Index: 0, Type: "url_citation", URLCitation: &TextAnnotationURLCitation{URL: "second"}}, + }, + }, + } + + _, err := ConcatAssistantGenTextExtensions(exts) + assert.Error(t, err) + }) +} + +func TestConcatReasoningExtensions(t *testing.T) { + t.Run("empty chunks - error", func(t *testing.T) { + _, err := ConcatReasoningExtensions(nil) + assert.Error(t, err) + }) + + t.Run("single extension", func(t *testing.T) { + ext := &ReasoningExtension{ + Content: []*ReasoningContent{ + {Text: "hello", Index: generic.PtrOf(0)}, + }, + } + + result, err := ConcatReasoningExtensions([]*ReasoningExtension{ext}) + assert.NoError(t, err) + assert.Len(t, result.Content, 1) + assert.Equal(t, "hello", result.Content[0].Text) + }) + + t.Run("streaming scenario - same index concatenated, ordered by index", func(t *testing.T) { + exts := []*ReasoningExtension{ + {Content: []*ReasoningContent{{Text: "he", Index: generic.PtrOf(0)}}}, + {Content: []*ReasoningContent{{Text: "first", Index: generic.PtrOf(1)}}}, + {Content: []*ReasoningContent{{Text: "llo", Index: generic.PtrOf(0)}}}, + } + + result, err := ConcatReasoningExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Content, 2) + assert.Nil(t, result.Content[0].Index) + assert.Equal(t, "hello", result.Content[0].Text) + assert.Nil(t, result.Content[1].Index) + assert.Equal(t, "first", result.Content[1].Text) + }) + + t.Run("nil extension and nil content skipped", func(t *testing.T) { + exts := []*ReasoningExtension{ + nil, + {Content: []*ReasoningContent{nil, {Text: "ok", Index: generic.PtrOf(0)}}}, + } + + result, err := ConcatReasoningExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Content, 1) + assert.Equal(t, "ok", result.Content[0].Text) + }) + + t.Run("all unindexed - appended in arrival order", func(t *testing.T) { + exts := []*ReasoningExtension{ + {Content: []*ReasoningContent{{Text: "he"}}}, + {Content: []*ReasoningContent{{Text: "llo"}}}, + } + + result, err := ConcatReasoningExtensions(exts) + assert.NoError(t, err) + assert.Len(t, result.Content, 2) + assert.Nil(t, result.Content[0].Index) + assert.Equal(t, "he", result.Content[0].Text) + assert.Nil(t, result.Content[1].Index) + assert.Equal(t, "llo", result.Content[1].Text) + }) + + t.Run("mixed indexed and unindexed - error", func(t *testing.T) { + exts := []*ReasoningExtension{ + {Content: []*ReasoningContent{{Text: "indexed", Index: generic.PtrOf(0)}}}, + {Content: []*ReasoningContent{{Text: "no meta"}}}, + } + + _, err := ConcatReasoningExtensions(exts) + assert.Error(t, err) + }) +} diff --git a/schema/select.go b/schema/select.go new file mode 100644 index 0000000..ed22dfc --- /dev/null +++ b/schema/select.go @@ -0,0 +1,73 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +const maxSelectNum = 5 + +func receiveN[T any](chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) { + return []func(chosenList []int, ss []*stream[T]) (index int, item *streamItem[T], ok bool){ + nil, + func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) { + item, ok := <-ss[chosenList[0]].items + return chosenList[0], &item, ok + }, + func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) { + select { + case item, ok := <-ss[chosenList[0]].items: + return chosenList[0], &item, ok + case item, ok := <-ss[chosenList[1]].items: + return chosenList[1], &item, ok + } + }, + func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) { + select { + case item, ok := <-ss[chosenList[0]].items: + return chosenList[0], &item, ok + case item, ok := <-ss[chosenList[1]].items: + return chosenList[1], &item, ok + case item, ok := <-ss[chosenList[2]].items: + return chosenList[2], &item, ok + } + }, + func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) { + select { + case item, ok := <-ss[chosenList[0]].items: + return chosenList[0], &item, ok + case item, ok := <-ss[chosenList[1]].items: + return chosenList[1], &item, ok + case item, ok := <-ss[chosenList[2]].items: + return chosenList[2], &item, ok + case item, ok := <-ss[chosenList[3]].items: + return chosenList[3], &item, ok + } + }, + func(chosenList []int, ss []*stream[T]) (int, *streamItem[T], bool) { + select { + case item, ok := <-ss[chosenList[0]].items: + return chosenList[0], &item, ok + case item, ok := <-ss[chosenList[1]].items: + return chosenList[1], &item, ok + case item, ok := <-ss[chosenList[2]].items: + return chosenList[2], &item, ok + case item, ok := <-ss[chosenList[3]].items: + return chosenList[3], &item, ok + case item, ok := <-ss[chosenList[4]].items: + return chosenList[4], &item, ok + } + }, + }[len(chosenList)](chosenList, ss) +} diff --git a/schema/serialization.go b/schema/serialization.go new file mode 100644 index 0000000..169bf9e --- /dev/null +++ b/schema/serialization.go @@ -0,0 +1,148 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "encoding/gob" + "reflect" + + "github.com/cloudwego/eino/internal/generic" + "github.com/cloudwego/eino/internal/serialization" +) + +func init() { + RegisterName[*Message]("_eino_message") + RegisterName[[]*Message]("_eino_message_slice") + RegisterName[*AgenticMessage]("_eino_agentic_message") + RegisterName[[]*AgenticMessage]("_eino_agentic_message_slice") + RegisterName[Document]("_eino_document") + RegisterName[RoleType]("_eino_role_type") + RegisterName[ToolCall]("_eino_tool_call") + RegisterName[FunctionCall]("_eino_function_call") + RegisterName[ResponseMeta]("_eino_response_meta") + RegisterName[TokenUsage]("_eino_token_usage") + RegisterName[LogProbs]("_eino_log_probs") + RegisterName[ChatMessagePart]("_eino_chat_message_part") + RegisterName[ChatMessagePartType]("_eino_chat_message_type") + RegisterName[ChatMessageImageURL]("_eino_chat_message_image_url") + RegisterName[ChatMessageAudioURL]("_eino_chat_message_audio_url") + RegisterName[ChatMessageVideoURL]("_eino_chat_message_video_url") + RegisterName[ChatMessageFileURL]("_eino_chat_message_file_url") + RegisterName[MessageInputPart]("_eino_message_input_part") + RegisterName[MessageInputImage]("_eino_message_input_image") + RegisterName[MessageInputAudio]("_eino_message_input_audio") + RegisterName[MessageInputVideo]("_eino_message_input_video") + RegisterName[MessageInputFile]("_eino_message_input_file") + RegisterName[MessageOutputPart]("_eino_message_output_part") + RegisterName[MessageOutputImage]("_eino_message_output_image") + RegisterName[MessageOutputAudio]("_eino_message_output_audio") + RegisterName[MessageOutputVideo]("_eino_message_output_video") + RegisterName[MessagePartCommon]("_eino_message_part_common") + RegisterName[ImageURLDetail]("_eino_image_url_detail") + RegisterName[PromptTokenDetails]("_eino_prompt_token_details") +} + +// RegisterName registers a type with a specific name for serialization. This is +// required for any type you intend to persist in a graph or ADK checkpoint. +// Use this function to maintain backward compatibility by mapping a type to a +// previously used name. For new types, `Register` is preferred. +// +// It is recommended to call this in an `init()` function in the file where the +// type is declared. +// +// What to Register: +// - Top-level types used as state (e.g., structs). +// - Concrete types that are assigned to interface fields. +// +// What NOT to Register: +// - Struct fields with concrete types (e.g., `string`, `int`, other structs). +// These are inferred via reflection. +// +// Serialization Rules: +// +// The serialization behavior is based on Go's standard `encoding/gob` package. +// See https://pkg.go.dev/encoding/gob for detailed rules. +// - Only exported struct fields are serialized. +// - Functions and channels are not supported and will be ignored. +// +// This function panics if registration fails. +func RegisterName[T any](name string) { + gob.RegisterName(name, generic.NewInstance[T]()) + + err := serialization.GenericRegister[T](name) + if err != nil { + panic(err) + } +} + +func getTypeName(rt reflect.Type) string { + name := rt.String() + + // But for named types (or pointers to them), qualify with import path. + // Dereference one pointer looking for a named type. + star := "" + if rt.Name() == "" { + if pt := rt; pt.Kind() == reflect.Pointer { + star = "*" + rt = pt.Elem() + } + } + if rt.Name() != "" { + if rt.PkgPath() == "" { + name = star + rt.Name() + } else { + name = star + rt.PkgPath() + "." + rt.Name() + } + } + return name +} + +// Register registers a type for serialization. This is required for any type +// you intend to persist in a graph or ADK checkpoint. It automatically determines +// the type name and is the recommended method for registering new types. +// +// It is recommended to call this in an `init()` function in the file where the +// type is declared. +// +// What to Register: +// - Top-level types used as state (e.g., structs). +// - Concrete types that are assigned to interface fields. +// +// What NOT to Register: +// - Struct fields with concrete types (e.g., `string`, `int`, other structs). +// These are inferred via reflection. +// +// Serialization Rules: +// +// The serialization behavior is based on Go's standard `encoding/gob` package. +// See https://pkg.go.dev/encoding/gob for detailed rules. +// - Only exported struct fields are serialized. +// - Functions and channels are not supported and will be ignored. +// +// This function panics if registration fails. +func Register[T any]() { + value := generic.NewInstance[T]() + + gob.Register(value) + + name := getTypeName(reflect.TypeOf(value)) + + err := serialization.GenericRegister[T](name) + if err != nil { + panic(err) + } +} diff --git a/schema/serialization_test.go b/schema/serialization_test.go new file mode 100644 index 0000000..d17cc40 --- /dev/null +++ b/schema/serialization_test.go @@ -0,0 +1,191 @@ +/* + * Copyright 2025 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "bytes" + "encoding/gob" + "fmt" + "reflect" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/internal/serialization" +) + +type testStruct struct{} + +func TestGetTypeName(t *testing.T) { + type localNamedType struct{} + + testCases := []struct { + name string + input reflect.Type + expected string + }{ + { + name: "named type from current package", + input: reflect.TypeOf(testStruct{}), + expected: "github.com/cloudwego/eino/schema.testStruct", + }, + { + name: "pointer to named type from current package", + input: reflect.TypeOf(&testStruct{}), + expected: "*github.com/cloudwego/eino/schema.testStruct", + }, + { + name: "unnamed map type", + input: reflect.TypeOf(map[string]int{}), + expected: "map[string]int", + }, + { + name: "pointer to unnamed map type", + input: reflect.TypeOf(new(map[string]int)), + expected: "*map[string]int", + }, + { + name: "built-in type", + input: reflect.TypeOf(0), + expected: "int", + }, + { + name: "pointer to built-in type", + input: reflect.TypeOf(new(int)), + expected: "*int", + }, + { + name: "named type from standard library", + input: reflect.TypeOf(bytes.Buffer{}), + expected: "bytes.Buffer", + }, + { + name: "pointer to named type from standard library", + input: reflect.TypeOf(&bytes.Buffer{}), + expected: "*bytes.Buffer", + }, + { + name: "local named type", + input: reflect.TypeOf(localNamedType{}), + expected: "github.com/cloudwego/eino/schema.localNamedType", + }, + { + name: "pointer to local named type", + input: reflect.TypeOf(&localNamedType{}), + expected: "*github.com/cloudwego/eino/schema.localNamedType", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + actual := getTypeName(tc.input) + if actual != tc.expected { + t.Errorf("getTypeName() got %q, want %q", actual, tc.expected) + } + }) + } +} + +func TestRegister(t *testing.T) { + type testStruct1 struct { + A any + B any + C any + D any + E any + F any + } + + type testStruct2 struct{} + + Register[*testStruct1]() + Register[*testStruct2]() + Register[[]Message]() + Register[[]*testStruct2]() + Register[[]testStruct2]() + + t1 := testStruct1{A: []*Message{{}}, B: []Message{{}}, C: []*testStruct2{{}}, D: []testStruct2{{}}, + E: &testStruct1{}, F: []int{1}} + + in := &serialization.InternalSerializer{} + mar, err := in.Marshal(t1) + if err != nil { + panic(err) + } + var t2 testStruct1 + err = in.Unmarshal(mar, &t2) + if err != nil { + panic(err) + } + assert.Equal(t, t1, t2) + + buf := new(bytes.Buffer) + err = gob.NewEncoder(buf).Encode(t1) + if err != nil { + panic(err) + } + err = gob.NewDecoder(buf).Decode(&t2) + if err != nil { + panic(err) + } + assert.Equal(t, t1, t2) + + f := func() (err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("panic: %v", r) + } + }() + + Register[[]int]() + Register[map[string]any]() + Register[[]*testStruct1]() + Register[[]testStruct1]() + + return nil + } + + err = f() + assert.NoError(t, err) +} + +// TestRegisterStructWithUUIDField reproduces issue #607 +// uuid.UUID is a [16]byte array. Prior to the fix, calling schema.RegisterName on +// a struct with a uuid.UUID field would panic during deserialization. +func TestRegisterStructWithUUIDField(t *testing.T) { + type Item struct { + ID uuid.UUID + } + + RegisterName[Item]("test_item") + + original := Item{ + ID: uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8"), + } + + s := &serialization.InternalSerializer{} + + data, err := s.Marshal(original) + assert.NoError(t, err) + + var result Item + err = s.Unmarshal(data, &result) + assert.NoError(t, err) + + assert.Equal(t, original.ID, result.ID) +} diff --git a/schema/stream.go b/schema/stream.go new file mode 100644 index 0000000..5625efe --- /dev/null +++ b/schema/stream.go @@ -0,0 +1,1024 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "errors" + "fmt" + "io" + "reflect" + "runtime" + "runtime/debug" + "sync" + "sync/atomic" + + "github.com/cloudwego/eino/internal/safe" +) + +// ErrNoValue is a sentinel returned from the convert function passed to +// [StreamReaderWithConvert] to skip a stream element — the element is dropped +// and the next one is read without surfacing an error to the caller. +// +// Use it to filter out empty or irrelevant chunks: +// +// outStream = schema.StreamReaderWithConvert(s, +// func(src string) (string, error) { +// if len(src) == 0 { +// return "", schema.ErrNoValue // skip empty chunks +// } +// return src, nil +// }) +// +// DO NOT use ErrNoValue in any other context. +var ErrNoValue = errors.New("no value") + +// ErrRecvAfterClosed indicates that StreamReader.Recv was unexpectedly called after StreamReader.Close. +// This error should not occur during normal use of StreamReader.Recv. If it does, please check your application code. +var ErrRecvAfterClosed = errors.New("recv after stream closed") + +// SourceEOF represents an EOF error from a specific source stream. +// It is only returned by the method Recv() of StreamReader created +// with MergeNamedStreamReaders when one of its source streams reaches EOF. +type SourceEOF struct { + sourceName string +} + +func (e *SourceEOF) Error() string { + return fmt.Sprintf("EOF from source stream: %s", e.sourceName) +} + +// GetSourceName extracts the source stream name from a SourceEOF error. +// It returns the source name and a boolean indicating whether the error was a SourceEOF. +// If the error is not a SourceEOF, it returns an empty string and false. +func GetSourceName(err error) (string, bool) { + var sErr *SourceEOF + if errors.As(err, &sErr) { + return sErr.sourceName, true + } + + return "", false +} + +// Pipe creates a new stream with the given capacity that represented with StreamWriter and StreamReader. +// The capacity is the maximum number of items that can be buffered in the stream. +// e.g. +// +// sr, sw := schema.Pipe[string](3) +// go func() { // send data +// defer sw.Close() +// for i := 0; i < 10; i++ { +// sw.Send(i, nil) +// } +// } +// +// defer sr.Close() +// for { +// chunk, err := sr.Recv() +// if errors.Is(err, io.EOF) { +// break +// } +// if err != nil { +// panic(err) +// } +// fmt.Println(chunk) +// } +func Pipe[T any](cap int) (*StreamReader[T], *StreamWriter[T]) { + stm := newStream[T](cap) + return stm.asReader(), &StreamWriter[T]{stm: stm} +} + +// StreamWriter the sender of a stream. +// created by Pipe function. +// eg. +// +// sr, sw := schema.Pipe[string](3) +// go func() { // send data +// defer sw.Close() +// for i := 0; i < 10; i++ { +// sw.Send(i, nil) +// } +// } +type StreamWriter[T any] struct { + stm *stream[T] +} + +// Send sends a value to the stream. +// e.g. +// +// closed := sw.Send(i, nil) +// if closed { +// // the stream is closed +// } +func (sw *StreamWriter[T]) Send(chunk T, err error) (closed bool) { + return sw.stm.send(chunk, err) +} + +// Close notify the receiver that the stream sender has finished. +// The stream receiver will get an error of io.EOF from StreamReader.Recv(). +// Notice: always remember to call Close() after sending all data. +// eg. +// +// defer sw.Close() +// for i := 0; i < 10; i++ { +// sw.Send(i, nil) +// } +func (sw *StreamWriter[T]) Close() { + sw.stm.closeSend() +} + +// StreamReader is the consumer side of an Eino stream. +// +// A StreamReader is read-once: only one goroutine should call Recv, and the +// reader must be closed exactly once (whether the loop finishes normally or +// exits early via break or return). +// +// Typical usage: +// +// defer sr.Close() // always close, even after io.EOF +// for { +// chunk, err := sr.Recv() +// if errors.Is(err, io.EOF) { +// break +// } +// if err != nil { +// return err +// } +// process(chunk) +// } +// +// To fan-out a single stream to N independent consumers, call [StreamReader.Copy] +// before any Recv; the original reader becomes unusable after the call. +// +// StreamReaders are created by [Pipe], [StreamReaderFromArray], +// [MergeStreamReaders], [MergeNamedStreamReaders], and [StreamReaderWithConvert]. +type StreamReader[T any] struct { + typ readerType + + st *stream[T] + + ar *arrayReader[T] + + msr *multiStreamReader[T] + + srw *streamReaderWithConvert[T] + + csr *childStreamReader[T] +} + +// Recv receives a value from the stream. +// eg. +// +// for { +// chunk, err := sr.Recv() +// if errors.Is(err, io.EOF) { +// break +// } +// if err != nil { +// panic(err) +// } +// fmt.Println(chunk) +// } +func (sr *StreamReader[T]) Recv() (T, error) { + switch sr.typ { + case readerTypeStream: + return sr.st.recv() + case readerTypeArray: + return sr.ar.recv() + case readerTypeMultiStream: + return sr.msr.recv() + case readerTypeWithConvert: + return sr.srw.recv() + case readerTypeChild: + return sr.csr.recv() + default: + panic("impossible") + } +} + +// Close safely closes the StreamReader. +// It should be called only once, as multiple calls may not work as expected. +// Notice: always remember to call Close() after using Recv(). +// e.g. +// +// defer sr.Close() +// +// for { +// chunk, err := sr.Recv() +// if errors.Is(err, io.EOF) { +// break +// } +// if err != nil { +// panic(err) +// } +// fmt.Println(chunk) +// } +func (sr *StreamReader[T]) Close() { + switch sr.typ { + case readerTypeStream: + sr.st.closeRecv() + case readerTypeArray: + + case readerTypeMultiStream: + sr.msr.close() + case readerTypeWithConvert: + sr.srw.close() + case readerTypeChild: + sr.csr.close() + default: + panic("impossible") + } +} + +// Copy creates n independent StreamReaders that each receive every element of +// the original stream. The original StreamReader becomes unusable after Copy. +// +// Use Copy when two or more pipeline branches need the same stream — +// for example, when a stream must be fed to both a callback handler and the +// next node in a graph: +// +// copies := sr.Copy(2) +// sr1, sr2 := copies[0], copies[1] +// defer sr1.Close() +// defer sr2.Close() +// +// // sr1 and sr2 independently read the same elements +// +// n must be at least 1. If n < 2, the original reader is returned unchanged. +func (sr *StreamReader[T]) Copy(n int) []*StreamReader[T] { + if n < 2 { + return []*StreamReader[T]{sr} + } + + if sr.typ == readerTypeArray { + ret := make([]*StreamReader[T], n) + for i, ar := range sr.ar.copy(n) { + ret[i] = &StreamReader[T]{typ: readerTypeArray, ar: ar} + } + return ret + } + + return copyStreamReaders[T](sr, n) +} + +// SetAutomaticClose sets the StreamReader to automatically close when it's no longer reachable and ready to be GCed. +// NOT concurrency safe. +func (sr *StreamReader[T]) SetAutomaticClose() { + switch sr.typ { + case readerTypeStream: + if !sr.st.automaticClose { + sr.st.automaticClose = true + var flag uint32 + sr.st.closedFlag = &flag + runtime.SetFinalizer(sr, func(s *StreamReader[T]) { + s.Close() + }) + } + case readerTypeMultiStream: + for _, s := range sr.msr.nonClosedStreams() { + if !s.automaticClose { + s.automaticClose = true + var flag uint32 + s.closedFlag = &flag + runtime.SetFinalizer(s, func(st *stream[T]) { + st.closeRecv() + }) + } + } + case readerTypeChild: + parent := sr.csr.parent.sr + parent.SetAutomaticClose() + case readerTypeWithConvert: + sr.srw.sr.SetAutomaticClose() + case readerTypeArray: + // no need to clean up + default: + } +} + +func (sr *StreamReader[T]) recvAny() (any, error) { + return sr.Recv() +} + +func (sr *StreamReader[T]) copyAny(n int) []iStreamReader { + ret := make([]iStreamReader, n) + + srs := sr.Copy(n) + + for i := 0; i < n; i++ { + ret[i] = srs[i] + } + + return ret +} + +func arrToStream[T any](arr []T) *stream[T] { + s := newStream[T](len(arr)) + for i := range arr { + s.send(arr[i], nil) + } + s.closeSend() + + return s +} + +func (sr *StreamReader[T]) toStream() *stream[T] { + switch sr.typ { + case readerTypeStream: + return sr.st + case readerTypeArray: + return sr.ar.toStream() + case readerTypeMultiStream: + return sr.msr.toStream() + case readerTypeWithConvert: + return sr.srw.toStream() + case readerTypeChild: + return sr.csr.toStream() + default: + panic("impossible") + } +} + +type readerType int + +const ( + readerTypeStream readerType = iota + readerTypeArray + readerTypeMultiStream + readerTypeWithConvert + readerTypeChild +) + +type iStreamReader interface { + recvAny() (any, error) + copyAny(int) []iStreamReader + Close() + SetAutomaticClose() +} + +// stream is a channel-based stream with 1 sender and 1 receiver. +// The sender calls closeSend() to notify the receiver that the stream sender has finished. +// The receiver calls closeRecv() to notify the sender that the receiver stop receiving. +type stream[T any] struct { + items chan streamItem[T] + + closed chan struct{} + + automaticClose bool + closedFlag *uint32 // 0 = not closed, 1 = closed, only used when automaticClose is set +} + +type streamItem[T any] struct { + chunk T + err error +} + +func newStream[T any](cap int) *stream[T] { + return &stream[T]{ + items: make(chan streamItem[T], cap), + closed: make(chan struct{}), + } +} + +func (s *stream[T]) asReader() *StreamReader[T] { + return &StreamReader[T]{typ: readerTypeStream, st: s} +} + +func (s *stream[T]) recv() (chunk T, err error) { + item, ok := <-s.items + + if !ok { + item.err = io.EOF + } + + return item.chunk, item.err +} + +func (s *stream[T]) send(chunk T, err error) (closed bool) { + // if the stream is closed, return immediately + select { + case <-s.closed: + return true + default: + } + + item := streamItem[T]{chunk, err} + + select { + case <-s.closed: + return true + case s.items <- item: + return false + } +} + +func (s *stream[T]) closeSend() { + close(s.items) +} + +func (s *stream[T]) closeRecv() { + if s.automaticClose { + if atomic.CompareAndSwapUint32(s.closedFlag, 0, 1) { + close(s.closed) + } + return + } + + close(s.closed) +} + +// StreamReaderFromArray creates a StreamReader from a given slice of elements. +// It takes an array of type T and returns a pointer to a StreamReader[T]. +// This allows for streaming the elements of the array in a controlled manner. +// eg. +// +// sr := schema.StreamReaderFromArray([]int{1, 2, 3}) +// defer sr.Close() +// +// for { +// chunk, err := sr.Recv() +// if errors.Is(err, io.EOF) { +// break +// } +// if err != nil { +// panic(err) +// } +// fmt.Println(chunk) +// } +func StreamReaderFromArray[T any](arr []T) *StreamReader[T] { + return &StreamReader[T]{ar: &arrayReader[T]{arr: arr}, typ: readerTypeArray} +} + +type arrayReader[T any] struct { + arr []T + index int +} + +func (ar *arrayReader[T]) recv() (T, error) { + if ar.index < len(ar.arr) { + ret := ar.arr[ar.index] + ar.index++ + + return ret, nil + } + + var t T + return t, io.EOF +} + +func (ar *arrayReader[T]) copy(n int) []*arrayReader[T] { + ret := make([]*arrayReader[T], n) + + for i := 0; i < n; i++ { + ret[i] = &arrayReader[T]{ + arr: ar.arr, + index: ar.index, + } + } + + return ret +} + +func (ar *arrayReader[T]) toStream() *stream[T] { + return arrToStream(ar.arr[ar.index:]) +} + +type multiArrayReader[T any] struct { + ars []*arrayReader[T] + index int +} + +type multiStreamReader[T any] struct { + sts []*stream[T] + + itemsCases []reflect.SelectCase + + nonClosed []int + + sourceReaderNames []string +} + +func newMultiStreamReader[T any](sts []*stream[T]) *multiStreamReader[T] { + var itemsCases []reflect.SelectCase + if len(sts) > maxSelectNum { + itemsCases = make([]reflect.SelectCase, len(sts)) + for i, st := range sts { + itemsCases[i] = reflect.SelectCase{ + Dir: reflect.SelectRecv, + Chan: reflect.ValueOf(st.items), + } + } + } + + nonClosed := make([]int, len(sts)) + for i := range sts { + nonClosed[i] = i + } + + return &multiStreamReader[T]{ + sts: sts, + itemsCases: itemsCases, + nonClosed: nonClosed, + } +} + +func (msr *multiStreamReader[T]) recv() (T, error) { + for len(msr.nonClosed) > 0 { + var chosen int + var ok bool + if len(msr.nonClosed) > maxSelectNum { + var recv reflect.Value + chosen, recv, ok = reflect.Select(msr.itemsCases) + if ok { + item := recv.Interface().(streamItem[T]) + return item.chunk, item.err + } + msr.itemsCases[chosen].Chan = reflect.Value{} + } else { + var item *streamItem[T] + chosen, item, ok = receiveN(msr.nonClosed, msr.sts) + if ok { + return item.chunk, item.err + } + } + + // delete the closed stream + for i := range msr.nonClosed { + if msr.nonClosed[i] == chosen { + msr.nonClosed = append(msr.nonClosed[:i], msr.nonClosed[i+1:]...) + break + } + } + + if len(msr.sourceReaderNames) > 0 { + var t T + return t, &SourceEOF{msr.sourceReaderNames[chosen]} + } + } + + var t T + return t, io.EOF +} + +func (msr *multiStreamReader[T]) nonClosedStreams() []*stream[T] { + ret := make([]*stream[T], len(msr.nonClosed)) + + for i, idx := range msr.nonClosed { + ret[i] = msr.sts[idx] + } + + return ret +} + +func (msr *multiStreamReader[T]) close() { + for _, s := range msr.sts { + s.closeRecv() + } +} + +func (msr *multiStreamReader[T]) toStream() *stream[T] { + return toStream[T, *multiStreamReader[T]](msr) +} + +type streamReaderWithConvert[T any] struct { + sr iStreamReader + + convert func(any) (T, error) + + errWrapper func(error) error + onEOF func() (T, error) + eofDone bool +} + +func newStreamReaderWithConvert[T any](origin iStreamReader, convert func(any) (T, error), opts ...ConvertOption) *StreamReader[T] { + opt := &convertOptions{} + for _, o := range opts { + o(opt) + } + + srw := &streamReaderWithConvert[T]{ + sr: origin, + convert: convert, + errWrapper: opt.ErrWrapper, + } + + if opt.OnEOF != nil { + typedOnEOF := opt.OnEOF + srw.onEOF = func() (T, error) { + v, err := typedOnEOF() + if err != nil { + var t T + return t, err + } + if v == nil { + var t T + return t, nil + } + return v.(T), nil + } + } + + return &StreamReader[T]{ + typ: readerTypeWithConvert, + srw: srw, + } +} + +type convertOptions struct { + ErrWrapper func(error) error + OnEOF func() (any, error) +} + +type ConvertOption func(*convertOptions) + +// WithErrWrapper wraps non-EOF errors from the underlying stream reader during +// conversion by StreamReaderWithConvert. Errors returned by the convert function +// itself are not wrapped. +// If the wrapper returns nil, the errored chunk is skipped and the next chunk +// is read. If the wrapper returns a non-nil error, that error is surfaced to +// the caller. +func WithErrWrapper(wrapper func(error) error) ConvertOption { + return func(o *convertOptions) { + o.ErrWrapper = wrapper + } +} + +// WithOnEOF registers a callback that fires once when the stream reaches EOF. +// The callback can inject an error or a value before the final io.EOF is returned. +// If the callback returns (nil, io.EOF), the stream ends normally. +// If it returns a non-EOF error, that error is delivered first, then subsequent Recv returns io.EOF. +// If it returns a non-nil value with nil error, that value is delivered first, then io.EOF. +func WithOnEOF(fn func() (any, error)) ConvertOption { + return func(o *convertOptions) { + o.OnEOF = fn + } +} + +// StreamReaderWithConvert returns a new StreamReader[D] that wraps sr and +// applies convert to every element. The original reader sr must not be used +// after calling this function. +// +// Filtering: if convert returns [ErrNoValue], the element is silently dropped +// and the next element is read. This lets you strip empty or irrelevant chunks +// without surfacing an error to the caller. +// +// Error wrapping: use [WithErrWrapper] to wrap non-convert errors (e.g. those +// arriving from an upstream source) before they reach the caller. +// +// intReader := schema.StreamReaderFromArray([]int{0, 1, 2, 3}) +// strReader := schema.StreamReaderWithConvert(intReader, +// func(i int) (string, error) { +// if i == 0 { +// return "", schema.ErrNoValue // skip zero +// } +// return fmt.Sprintf("val_%d", i), nil +// }) +// defer strReader.Close() +// // Recv yields "val_1", "val_2", "val_3" +func StreamReaderWithConvert[T, D any](sr *StreamReader[T], convert func(T) (D, error), opts ...ConvertOption) *StreamReader[D] { + c := func(a any) (D, error) { + return convert(a.(T)) + } + + return newStreamReaderWithConvert(sr, c, opts...) +} + +func (srw *streamReaderWithConvert[T]) recv() (T, error) { + for { + out, err := srw.sr.recvAny() + + if err != nil { + var t T + if err == io.EOF { + if srw.onEOF != nil && !srw.eofDone { + srw.eofDone = true + val, onEOFErr := srw.onEOF() + if onEOFErr != io.EOF { + return val, onEOFErr + } + } + return t, io.EOF + } + if srw.errWrapper != nil { + err = srw.errWrapper(err) + if err != nil { + return t, err + } + + continue + } + + return t, err + } + + t, err := srw.convert(out) + if err == nil { + return t, nil + } + + if !errors.Is(err, ErrNoValue) { + return t, err + } + } +} + +func (srw *streamReaderWithConvert[T]) close() { + srw.sr.Close() +} + +type reader[T any] interface { + recv() (T, error) + close() +} + +func toStream[T any, Reader reader[T]](r Reader) *stream[T] { + ret := newStream[T](5) + + go func() { + defer func() { + panicErr := recover() + if panicErr != nil { + e := safe.NewPanicErr(panicErr, debug.Stack()) + + var chunk T + _ = ret.send(chunk, e) + } + + ret.closeSend() + r.close() + }() + + for { + out, err := r.recv() + if err == io.EOF { + break + } + + closed := ret.send(out, err) + if closed { + break + } + } + }() + + return ret +} + +func (srw *streamReaderWithConvert[T]) toStream() *stream[T] { + return toStream[T, *streamReaderWithConvert[T]](srw) +} + +type cpStreamElement[T any] struct { + once sync.Once + next *cpStreamElement[T] + item streamItem[T] +} + +// copyStreamReaders creates multiple independent StreamReaders from a single StreamReader. +// Each child StreamReader can read from the original stream independently. +func copyStreamReaders[T any](sr *StreamReader[T], n int) []*StreamReader[T] { + cpsr := &parentStreamReader[T]{ + sr: sr, + subStreamList: make([]*cpStreamElement[T], n), + closedNum: 0, + } + + // Initialize subStreamList with an empty element, which acts like a tail node. + // A nil element (used for dereference) represents that the child has been closed. + // It is challenging to link the previous and current elements when the length of the original channel is unknown. + // Additionally, using a previous pointer complicates dereferencing elements, possibly requiring reference counting. + elem := &cpStreamElement[T]{} + + for i := range cpsr.subStreamList { + cpsr.subStreamList[i] = elem + } + + ret := make([]*StreamReader[T], n) + for i := range ret { + ret[i] = &StreamReader[T]{ + csr: &childStreamReader[T]{ + parent: cpsr, + index: i, + }, + typ: readerTypeChild, + } + } + + return ret +} + +type parentStreamReader[T any] struct { + // sr is the original StreamReader. + sr *StreamReader[T] + + // subStreamList maps each child's index to its latest read chunk. + // Each value comes from a hidden linked list of cpStreamElement. + subStreamList []*cpStreamElement[T] + + // closedNum is the count of closed children. + closedNum uint32 +} + +// peek is not safe for concurrent use with the same idx but is safe for different idx. +// Ensure that each child StreamReader uses a for-loop in a single goroutine. +func (p *parentStreamReader[T]) peek(idx int) (t T, err error) { + elem := p.subStreamList[idx] + if elem == nil { + // Unexpected call to receive after the child has been closed. + return t, ErrRecvAfterClosed + } + + // The sync.Once here is used to: + // 1. Write the content of this cpStreamElement. + // 2. Initialize the 'next' field of this cpStreamElement with an empty cpStreamElement, + // similar to the initialization in copyStreamReaders. + elem.once.Do(func() { + t, err = p.sr.Recv() + elem.item = streamItem[T]{chunk: t, err: err} + if err != io.EOF { + elem.next = &cpStreamElement[T]{} + p.subStreamList[idx] = elem.next + } + }) + + // The element has been set and will not be modified again. + // Therefore, children can read this element's content and 'next' pointer concurrently. + t = elem.item.chunk + err = elem.item.err + if err != io.EOF { + p.subStreamList[idx] = elem.next + } + + return t, err +} + +func (p *parentStreamReader[T]) close(idx int) { + if p.subStreamList[idx] == nil { + return // avoid close multiple times + } + + p.subStreamList[idx] = nil + + curClosedNum := atomic.AddUint32(&p.closedNum, 1) + + allClosed := int(curClosedNum) == len(p.subStreamList) + if allClosed { + p.sr.Close() + } +} + +type childStreamReader[T any] struct { + parent *parentStreamReader[T] + index int +} + +func (csr *childStreamReader[T]) recv() (T, error) { + return csr.parent.peek(csr.index) +} + +func (csr *childStreamReader[T]) toStream() *stream[T] { + return toStream[T, *childStreamReader[T]](csr) +} + +func (csr *childStreamReader[T]) close() { + csr.parent.close(csr.index) +} + +// MergeStreamReaders fans in multiple StreamReaders into a single StreamReader. +// Elements from all source streams are interleaved in arrival order (non-deterministic). +// The merged reader reaches EOF only after every source stream has been exhausted. +// +// Callers must still close the merged reader; it propagates the close signal +// to all underlying sources. +// +// Use [MergeNamedStreamReaders] instead when you need to know which source +// stream ended first (it emits a [SourceEOF] per-source EOF rather than +// silently discarding them). +// +// Returns nil if srs is empty. +func MergeStreamReaders[T any](srs []*StreamReader[T]) *StreamReader[T] { + if len(srs) < 1 { + return nil + } + + if len(srs) < 2 { + return srs[0] + } + + var arr []T + var ss []*stream[T] + + for _, sr := range srs { + switch sr.typ { + case readerTypeStream: + ss = append(ss, sr.st) + case readerTypeArray: + arr = append(arr, sr.ar.arr[sr.ar.index:]...) + case readerTypeMultiStream: + ss = append(ss, sr.msr.nonClosedStreams()...) + case readerTypeWithConvert: + ss = append(ss, sr.srw.toStream()) + case readerTypeChild: + ss = append(ss, sr.csr.toStream()) + default: + panic("impossible") + } + } + + if len(ss) == 0 { + return &StreamReader[T]{ + typ: readerTypeArray, + ar: &arrayReader[T]{ + arr: arr, + index: 0, + }, + } + } + + if len(arr) != 0 { + s := arrToStream(arr) + ss = append(ss, s) + } + + return &StreamReader[T]{ + typ: readerTypeMultiStream, + msr: newMultiStreamReader(ss), + } +} + +// MergeNamedStreamReaders merges multiple named StreamReaders into one. +// Unlike [MergeStreamReaders], when a source stream reaches EOF the merged +// reader emits a [SourceEOF] error (instead of silently continuing) so you can +// detect exactly which source finished. Use [GetSourceName] to retrieve the +// name from a SourceEOF error. The merged reader itself signals io.EOF only +// after all named sources are exhausted. +// +// This is useful when downstream logic must react differently to each source +// completing — for example, draining one agent's output before proceeding: +// +// namedStreams := map[string]*schema.StreamReader[string]{ +// "agent_a": srA, +// "agent_b": srB, +// } +// merged := schema.MergeNamedStreamReaders(namedStreams) +// defer merged.Close() +// for { +// chunk, err := merged.Recv() +// if errors.Is(err, io.EOF) { break } +// if name, ok := schema.GetSourceName(err); ok { +// fmt.Printf("%s finished\n", name) +// continue +// } +// if err != nil { return err } +// process(chunk) +// } +// +// Returns nil if srs is empty. +func MergeNamedStreamReaders[T any](srs map[string]*StreamReader[T]) *StreamReader[T] { + if len(srs) < 1 { + return nil + } + + ss := make([]*StreamReader[T], len(srs)) + names := make([]string, len(srs)) + + i := 0 + for name, sr := range srs { + ss[i] = sr + names[i] = name + i++ + } + + return InternalMergeNamedStreamReaders(ss, names) +} + +// InternalMergeNamedStreamReaders merges multiple readers with their names +// into a single multi-stream reader. +func InternalMergeNamedStreamReaders[T any](srs []*StreamReader[T], names []string) *StreamReader[T] { + ss := make([]*stream[T], len(srs)) + + for i, sr := range srs { + ss[i] = sr.toStream() + } + + msr := newMultiStreamReader(ss) + msr.sourceReaderNames = names + + return &StreamReader[T]{ + typ: readerTypeMultiStream, + msr: msr, + } +} diff --git a/schema/stream_copy_external_test.go b/schema/stream_copy_external_test.go new file mode 100644 index 0000000..fc9cc47 --- /dev/null +++ b/schema/stream_copy_external_test.go @@ -0,0 +1,155 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "fmt" + "io" + "runtime" + "sort" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestStream1(t *testing.T) { + runtime.GOMAXPROCS(1) + + sr, sw := Pipe[int](0) + go func() { + for i := 0; i < 100; i++ { + sw.Send(i, nil) + time.Sleep(3 * time.Millisecond) + } + sw.Close() + }() + copied := sr.Copy(2) + var ( + now = time.Now().UnixMilli() + ts = []int64{now, now} + tsOld = []int64{now, now} + ) + var count int32 + wg := sync.WaitGroup{} + wg.Add(2) + go func() { + i := 0 + s := copied[0] + for { + n, e := s.Recv() + if e != nil { + if e == io.EOF { + break + } + } + tsOld[0] = ts[0] + ts[0] = time.Now().UnixMilli() + interval := ts[0] - tsOld[0] + if interval >= 6 { + atomic.AddInt32(&count, 1) + } + assert.Equal(t, i, n) + i++ + } + wg.Done() + }() + go func() { + i := 0 + s := copied[1] + for { + n, e := s.Recv() + if e != nil { + if e == io.EOF { + break + } + } + tsOld[1] = ts[1] + ts[1] = time.Now().UnixMilli() + interval := ts[1] - tsOld[1] + if interval >= 6 { + atomic.AddInt32(&count, 1) + } + assert.Equal(t, i, n) + i++ + } + wg.Done() + }() + wg.Wait() + t.Logf("count= %d", count) +} + +type info struct { + idx int + ts int64 + after int64 + content string +} + +func TestCopyDelay(t *testing.T) { + runtime.GOMAXPROCS(10) + n := 3 + //m := 100 + s := newStream[string](0) + scp := s.asReader().Copy(n) + go func() { + s.send("1", nil) + s.send("2", nil) + time.Sleep(time.Second) + s.send("3", nil) + s.closeSend() + }() + wg := sync.WaitGroup{} + wg.Add(n) + infoList := make([][]info, n) + for i := 0; i < n; i++ { + j := i + go func() { + defer func() { + scp[j].Close() + wg.Done() + }() + for { + lastTime := time.Now() + str, err := scp[j].Recv() + if err == io.EOF { + break + } + now := time.Now() + infoList[j] = append(infoList[j], info{ + idx: j, + ts: now.UnixMicro(), + after: now.Sub(lastTime).Milliseconds(), + content: str, + }) + } + }() + } + wg.Wait() + infos := make([]info, 0) + for _, infoL := range infoList { + infos = append(infos, infoL...) + } + sort.Slice(infos, func(i, j int) bool { + return infos[i].ts < infos[j].ts + }) + for _, info := range infos { + fmt.Printf("child[%d] ts[%d] after[%5dms] content[%s]\n", info.idx, info.ts, info.after, info.content) + } +} diff --git a/schema/stream_oneof_test.go b/schema/stream_oneof_test.go new file mode 100644 index 0000000..740836d --- /dev/null +++ b/schema/stream_oneof_test.go @@ -0,0 +1,324 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema_test + +import ( + "errors" + "io" + "testing" + "time" + + "github.com/cloudwego/eino/schema" +) + +func recvAll(t *testing.T, sr *schema.StreamReader[string]) ([]string, []error) { + t.Helper() + var vals []string + var errs []error + for { + v, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + errs = append(errs, err) + continue + } + vals = append(vals, v) + } + return vals, errs +} + +func makeStream(items []string, opts ...schema.ConvertOption) *schema.StreamReader[string] { + return schema.StreamReaderWithConvert( + schema.StreamReaderFromArray(items), + func(s string) (string, error) { return s, nil }, + opts..., + ) +} + +func TestWithOnEOF_PassThroughEOF(t *testing.T) { + items := []string{"a", "b", "c", "d"} + sr := makeStream(items, schema.WithOnEOF(func() (any, error) { + return nil, io.EOF + })) + defer sr.Close() + + vals, errs := recvAll(t, sr) + if len(errs) != 0 { + t.Fatalf("expected no errors, got %v", errs) + } + if len(vals) != 4 { + t.Fatalf("expected 4 values, got %d: %v", len(vals), vals) + } + for i, want := range items { + if vals[i] != want { + t.Errorf("vals[%d] = %q, want %q", i, vals[i], want) + } + } +} + +func TestWithOnEOF_InjectError(t *testing.T) { + items := []string{"a", "b", "c", "d"} + customErr := errors.New("validation failed") + sr := makeStream(items, schema.WithOnEOF(func() (any, error) { + return nil, customErr + })) + defer sr.Close() + + var vals []string + var gotCustomErr bool + for { + v, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + if errors.Is(err, customErr) { + gotCustomErr = true + continue + } + t.Fatalf("unexpected error: %v", err) + } + vals = append(vals, v) + } + + if len(vals) != 4 { + t.Fatalf("expected 4 values, got %d: %v", len(vals), vals) + } + if !gotCustomErr { + t.Fatalf("expected custom error from onEOF, did not receive it") + } +} + +func TestWithOnEOF_InjectValue(t *testing.T) { + items := []string{"a", "b", "c", "d"} + sr := makeStream(items, schema.WithOnEOF(func() (any, error) { + return "extra", nil + })) + defer sr.Close() + + var vals []string + for { + v, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + vals = append(vals, v) + } + + if len(vals) != 5 { + t.Fatalf("expected 5 values, got %d: %v", len(vals), vals) + } + if vals[4] != "extra" { + t.Errorf("vals[4] = %q, want %q", vals[4], "extra") + } +} + +func TestWithOnEOF_BlockingCallback(t *testing.T) { + sr, sw := schema.Pipe[string](0) + + unblock := make(chan struct{}) + converted := schema.StreamReaderWithConvert(sr, + func(s string) (string, error) { return s, nil }, + schema.WithOnEOF(func() (any, error) { + <-unblock + return "after-block", nil + }), + ) + defer converted.Close() + + go func() { + sw.Send("x", nil) + sw.Close() + }() + + v, err := converted.Recv() + if err != nil { + t.Fatalf("first Recv error: %v", err) + } + if v != "x" { + t.Fatalf("first Recv = %q, want %q", v, "x") + } + + done := make(chan struct{}) + var recvVal string + var recvErr error + go func() { + recvVal, recvErr = converted.Recv() + close(done) + }() + + select { + case <-done: + t.Fatal("Recv returned before unblock signal") + case <-time.After(50 * time.Millisecond): + } + + close(unblock) + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Recv did not return after unblock signal") + } + + if recvErr != nil { + t.Fatalf("second Recv error: %v", recvErr) + } + if recvVal != "after-block" { + t.Errorf("second Recv = %q, want %q", recvVal, "after-block") + } + + v3, err3 := converted.Recv() + if !errors.Is(err3, io.EOF) { + t.Fatalf("third Recv: got (%q, %v), want EOF", v3, err3) + } +} + +func TestWithOnEOF_EmptyStream(t *testing.T) { + customErr := errors.New("empty stream error") + sr := makeStream(nil, schema.WithOnEOF(func() (any, error) { + return nil, customErr + })) + defer sr.Close() + + v, err := sr.Recv() + if !errors.Is(err, customErr) { + t.Fatalf("first Recv: got (%q, %v), want customErr", v, err) + } + + v2, err2 := sr.Recv() + if !errors.Is(err2, io.EOF) { + t.Fatalf("second Recv: got (%q, %v), want EOF", v2, err2) + } +} + +func TestWithOnEOF_WithErrWrapper_ErrorPath(t *testing.T) { + sr, sw := schema.Pipe[string](0) + + streamErr := errors.New("stream error") + onEOFCalled := false + + converted := schema.StreamReaderWithConvert(sr, + func(s string) (string, error) { return s, nil }, + schema.WithErrWrapper(func(err error) error { + return err + }), + schema.WithOnEOF(func() (any, error) { + onEOFCalled = true + return nil, errors.New("should not happen") + }), + ) + defer converted.Close() + + go func() { + sw.Send("a", nil) + sw.Send("", streamErr) + sw.Close() + }() + + v, err := converted.Recv() + if err != nil { + t.Fatalf("first Recv error: %v", err) + } + if v != "a" { + t.Fatalf("first Recv = %q, want %q", v, "a") + } + + _, err = converted.Recv() + if !errors.Is(err, streamErr) { + t.Fatalf("second Recv: got %v, want streamErr", err) + } + + if onEOFCalled { + t.Fatal("onEOF should not have been called when stream errored") + } +} + +func TestWithOnEOF_WithErrWrapper_EOFPath(t *testing.T) { + items := []string{"a", "b", "c"} + errWrapperCalled := false + + sr := schema.StreamReaderWithConvert( + schema.StreamReaderFromArray(items), + func(s string) (string, error) { return s, nil }, + schema.WithErrWrapper(func(err error) error { + errWrapperCalled = true + return err + }), + schema.WithOnEOF(func() (any, error) { + return "oneof-val", nil + }), + ) + defer sr.Close() + + var vals []string + for { + v, err := sr.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + vals = append(vals, v) + } + + if len(vals) != 4 { + t.Fatalf("expected 4 values, got %d: %v", len(vals), vals) + } + if vals[3] != "oneof-val" { + t.Errorf("vals[3] = %q, want %q", vals[3], "oneof-val") + } + if errWrapperCalled { + t.Fatal("errWrapper should not have been called for clean stream") + } +} + +func TestWithOnEOF_MultipleRecvAfterEOF(t *testing.T) { + items := []string{"a"} + customErr := errors.New("oneof error") + + sr := makeStream(items, schema.WithOnEOF(func() (any, error) { + return nil, customErr + })) + defer sr.Close() + + v, err := sr.Recv() + if err != nil { + t.Fatalf("first Recv error: %v", err) + } + if v != "a" { + t.Fatalf("first Recv = %q, want %q", v, "a") + } + + _, err = sr.Recv() + if !errors.Is(err, customErr) { + t.Fatalf("second Recv: got %v, want customErr", err) + } + + for i := 0; i < 5; i++ { + _, err = sr.Recv() + if !errors.Is(err, io.EOF) { + t.Fatalf("Recv #%d after onEOF: got %v, want io.EOF", i+3, err) + } + } +} diff --git a/schema/stream_test.go b/schema/stream_test.go new file mode 100644 index 0000000..9282bf9 --- /dev/null +++ b/schema/stream_test.go @@ -0,0 +1,905 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "errors" + "fmt" + "io" + "math/rand" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestStream(t *testing.T) { + s := newStream[int](0) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 10; i++ { + closed := s.send(i, nil) + if closed { + break + } + } + s.closeSend() + }() + + i := 0 + for { + i++ + if i > 5 { + s.closeRecv() + break + } + v, err := s.recv() + if err != nil { + assert.ErrorIs(t, err, io.EOF) + break + } + t.Log(v) + } + + wg.Wait() +} + +func TestStreamCopy(t *testing.T) { + s := newStream[string](10) + srs := s.asReader().Copy(2) + + s.send("a", nil) + s.send("b", nil) + s.send("c", nil) + s.closeSend() + + defer func() { + for _, sr := range srs { + sr.Close() + } + }() + + for { + v, err := srs[0].Recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + t.Fatal(err) + } + + t.Log("copy 01 recv", v) + } + + for { + v, err := srs[1].Recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + t.Fatal(err) + } + + t.Log("copy 02 recv", v) + } + + for { + v, err := s.recv() + if errors.Is(err, io.EOF) { + break + } + + if err != nil { + t.Fatal(err) + } + + t.Log("recv origin", v) + } + + t.Log("done") +} + +func TestNewStreamCopy(t *testing.T) { + t.Run("test one index recv channel blocked while other indexes could recv", func(t *testing.T) { + s := newStream[string](1) + scp := s.asReader().Copy(2) + + var t1, t2 time.Time + + go func() { + s.send("a", nil) + t1 = time.Now() + time.Sleep(time.Millisecond * 200) + s.send("a", nil) + s.closeSend() + }() + wg := sync.WaitGroup{} + wg.Add(2) + + go func() { + defer func() { + scp[0].Close() + wg.Done() + }() + + for { + str, err := scp[0].Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, str, "a") + } + }() + + go func() { + defer func() { + scp[1].Close() + wg.Done() + }() + + time.Sleep(time.Millisecond * 100) + for { + str, err := scp[1].Recv() + if err == io.EOF { + break + } + + if t2.IsZero() { + t2 = time.Now() + } + + assert.NoError(t, err) + assert.Equal(t, str, "a") + } + }() + + wg.Wait() + + assert.True(t, t2.Sub(t1) < time.Millisecond*200) + }) + + t.Run("test one index recv channel blocked and other index closed", func(t *testing.T) { + s := newStream[string](1) + scp := s.asReader().Copy(2) + + go func() { + s.send("a", nil) + time.Sleep(time.Millisecond * 200) + s.send("a", nil) + s.closeSend() + }() + + wg := sync.WaitGroup{} + wg.Add(2) + + //buf := scp[0].csr.parent.mem.buf + go func() { + defer func() { + scp[0].Close() + wg.Done() + }() + + for { + str, err := scp[0].Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, str, "a") + } + }() + + go func() { + time.Sleep(time.Millisecond * 100) + scp[1].Close() + scp[1].Close() // try close multiple times + wg.Done() + }() + + wg.Wait() + + //assert.Equal(t, 0, buf.Len()) + }) + + t.Run("test long time recv", func(t *testing.T) { + s := newStream[int](2) + n := 1000 + go func() { + for i := 0; i < n; i++ { + s.send(i, nil) + } + + s.closeSend() + }() + + m := 100 + wg := sync.WaitGroup{} + wg.Add(m) + copies := s.asReader().Copy(m) + for i := 0; i < m; i++ { + idx := i + go func() { + cp := copies[idx] + l := 0 + defer func() { + assert.Equal(t, 1000, l) + cp.Close() + wg.Done() + }() + + for { + exp, err := cp.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, exp, l) + l++ + } + }() + } + + wg.Wait() + //memo := copies[0].csr.parent.mem + //assert.Equal(t, true, memo.hasFinished) + //assert.Equal(t, 0, memo.buf.Len()) + }) + + t.Run("test closes", func(t *testing.T) { + s := newStream[int](20) + n := 1000 + go func() { + for i := 0; i < n; i++ { + s.send(i, nil) + } + + s.closeSend() + }() + + m := 100 + wg := sync.WaitGroup{} + wg.Add(m) + + wgEven := sync.WaitGroup{} + wgEven.Add(m / 2) + + sr := s.asReader() + sr.SetAutomaticClose() + copies := sr.Copy(m) + for i := 0; i < m; i++ { + idx := i + go func() { + cp := copies[idx] + l := 0 + defer func() { + cp.Close() + wg.Done() + if idx%2 == 0 { + wgEven.Done() + } + }() + + for { + if idx%2 == 0 && l == idx { + break + } + + exp, err := cp.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, exp, l) + l++ + } + }() + } + + wgEven.Wait() + wg.Wait() + assert.Equal(t, m, int(copies[0].csr.parent.closedNum)) + }) + + t.Run("test reader do no close", func(t *testing.T) { + s := newStream[int](20) + n := 1000 + go func() { + for i := 0; i < n; i++ { + s.send(i, nil) + } + + s.closeSend() + }() + + m := 4 + wg := sync.WaitGroup{} + wg.Add(m) + + copies := s.asReader().Copy(m) + for i := 0; i < m; i++ { + idx := i + cp := copies[idx] + cp.SetAutomaticClose() + go func() { + l := 0 + defer func() { + wg.Done() + }() + + for { + exp, err := cp.Recv() + if err == io.EOF { + break + } + + assert.NoError(t, err) + assert.Equal(t, exp, l) + l++ + } + }() + } + + wg.Wait() + assert.Equal(t, 0, int(copies[0].csr.parent.closedNum)) // not closed + }) + +} + +func checkStream(s *StreamReader[int]) error { + defer s.Close() + + for i := 0; i < 10; i++ { + chunk, err := s.Recv() + if err != nil { + return err + } + if chunk != i { + return fmt.Errorf("receive err, expected:%d, actual: %d", i, chunk) + } + } + _, err := s.Recv() + if err != io.EOF { + return fmt.Errorf("close chan fail") + } + return nil +} + +func testStreamN(cap, n int) error { + s := newStream[int](cap) + go func() { + for i := 0; i < 10; i++ { + s.send(i, nil) + } + s.closeSend() + }() + + vs := s.asReader().Copy(n) + err := checkStream(vs[0]) + if err != nil { + return err + } + + vs = vs[1].Copy(n) + err = checkStream(vs[0]) + if err != nil { + return err + } + vs = vs[1].Copy(n) + err = checkStream(vs[0]) + if err != nil { + return err + } + return nil +} + +func TestCopy(t *testing.T) { + for i := 0; i < 10; i++ { + for j := 2; j < 10; j++ { + err := testStreamN(i, j) + if err != nil { + t.Fatal(err) + } + } + } +} + +func TestCopy5(t *testing.T) { + s := newStream[int](0) + go func() { + for i := 0; i < 10; i++ { + closed := s.send(i, nil) + if closed { + fmt.Printf("has closed") + } + } + s.closeSend() + }() + vs := s.asReader().Copy(5) + time.Sleep(time.Second) + defer func() { + for _, v := range vs { + v.Close() + } + }() + for i := 0; i < 10; i++ { + chunk, err := vs[0].Recv() + if err != nil { + t.Fatal(err) + } + if chunk != i { + t.Fatalf("receive err, expected:%d, actual: %d", i, chunk) + } + } + _, err := vs[0].Recv() + if err != io.EOF { + t.Fatalf("copied stream reader cannot return EOF") + } + _, err = vs[0].Recv() + if err != io.EOF { + t.Fatalf("copied stream reader cannot return EOF repeatedly") + } +} + +func TestStreamReaderWithConvert(t *testing.T) { + s := newStream[int](2) + + var cntA int + var e error + + convA := func(src int) (int, error) { + if src == 1 { + return 0, fmt.Errorf("mock err") + } + + return src, nil + } + + sta := StreamReaderWithConvert[int, int](s.asReader(), convA) + sta.SetAutomaticClose() + + s.send(1, nil) + s.send(2, nil) + s.closeSend() + + for { + item, err := sta.Recv() + if err != nil { + if err == io.EOF { + break + } + + e = err + continue + } + + cntA += item + } + + assert.NotNil(t, e) + assert.Equal(t, cntA, 2) +} + +func TestStreamReaderWithConvert_ErrWrapperContinue(t *testing.T) { + s := newStream[int](5) + + s.send(1, nil) + s.send(0, fmt.Errorf("transient error 1")) + s.send(2, nil) + s.send(0, fmt.Errorf("transient error 2")) + s.send(3, nil) + s.closeSend() + + wrapperCalls := 0 + sr := StreamReaderWithConvert[int, int](s.asReader(), func(v int) (int, error) { + return v, nil + }, WithErrWrapper(func(err error) error { + wrapperCalls++ + return nil + })) + + var results []int + for { + v, err := sr.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + results = append(results, v) + } + + assert.Equal(t, []int{1, 2, 3}, results) + assert.Equal(t, 2, wrapperCalls) + + s2 := newStream[int](3) + s2.send(0, fmt.Errorf("skip me")) + s2.send(42, nil) + s2.closeSend() + + sr2 := StreamReaderWithConvert[int, int](s2.asReader(), func(v int) (int, error) { + return v, nil + }, WithErrWrapper(func(err error) error { + return nil + })) + + v, err := sr2.Recv() + assert.NoError(t, err) + assert.Equal(t, 42, v) + sr2.Close() +} + +func TestArrayStreamCombined(t *testing.T) { + asr := &StreamReader[int]{ + typ: readerTypeArray, + ar: &arrayReader[int]{ + arr: []int{0, 1, 2}, + index: 0, + }, + } + + s := newStream[int](3) + for i := 3; i < 6; i++ { + s.send(i, nil) + } + s.closeSend() + + nSR := MergeStreamReaders([]*StreamReader[int]{asr, s.asReader()}) + nSR.SetAutomaticClose() + + record := make([]bool, 6) + for i := 0; i < 6; i++ { + chunk, err := nSR.Recv() + if err != nil { + t.Fatal(err) + } + if record[chunk] { + t.Fatal("record duplicated") + } + record[chunk] = true + } + + _, err := nSR.Recv() + if err != io.EOF { + t.Fatal("reader haven't finish correctly") + } + + for i := range record { + if !record[i] { + t.Fatal("record missing") + } + } +} + +func TestMultiStream(t *testing.T) { + var sts []*stream[int] + sum := 0 + for i := 0; i < 10; i++ { + size := rand.Intn(10) + 1 + sum += size + st := newStream[int](size) + for j := 1; j <= size; j++ { + st.send(j&0xffff+i<<16, nil) + } + st.closeSend() + sts = append(sts, st) + } + mst := newMultiStreamReader(sts) + receiveList := make([]int, 10) + for i := 0; i < sum; i++ { + chunk, err := mst.recv() + if err != nil { + t.Fatal(err) + } + if receiveList[chunk>>16] >= chunk&0xffff { + t.Fatal("out of order") + } + receiveList[chunk>>16] = chunk & 0xffff + } + _, err := mst.recv() + if err != io.EOF { + t.Fatal("end stream haven't return EOF") + } +} + +// TestMergeNamedStreamReaders tests the functionality of MergeNamedStreamReaders +// with a focus on SourceEOF error handling. +func TestMergeNamedStreamReaders(t *testing.T) { + t.Run("BasicSourceEOF", func(t *testing.T) { + // Create two named streams + sr1, sw1 := Pipe[string](2) + sr2, sw2 := Pipe[string](2) + + // Merge the streams with names + namedStreams := map[string]*StreamReader[string]{ + "stream1": sr1, + "stream2": sr2, + } + mergedSR := MergeNamedStreamReaders(namedStreams) + mergedSR.SetAutomaticClose() + + // Send data to the first stream and close it immediately + go func() { + defer sw1.Close() + sw1.Send("data1-1", nil) + sw1.Send("data1-2", nil) + // First stream ends + }() + + // Send data to the second stream with a delay before closing + go func() { + defer sw2.Close() + sw2.Send("data2-1", nil) + sw2.Send("data2-2", nil) + sw2.Send("data2-3", nil) + // Second stream ends + }() + + // Track received data and EOF sources + receivedData := make(map[string][]string) + eofSources := make([]string, 0, 2) + + for { + chunk, err := mergedSR.Recv() + if err != nil { + // Check if it's a SourceEOF error + if sourceName, ok := GetSourceName(err); ok { + eofSources = append(eofSources, sourceName) + t.Logf("Received EOF from source: %s", sourceName) + continue // Continue receiving from other streams + } + + // If it's a regular EOF, all streams have ended + if errors.Is(err, io.EOF) { + break + } + + // Handle other errors + t.Errorf("Error receiving data: %v", err) + break + } + + // Categorize data by prefix + if len(chunk) >= 5 { + prefix := chunk[:5] + if prefix == "data1" { + receivedData["stream1"] = append(receivedData["stream1"], chunk) + } else if prefix == "data2" { + receivedData["stream2"] = append(receivedData["stream2"], chunk) + } + } + } + + // Verify we received both SourceEOF errors + if len(eofSources) != 2 { + t.Errorf("Expected 2 SourceEOF errors, got %d", len(eofSources)) + } + + // Verify the source names are correct + expectedSources := map[string]bool{"stream1": false, "stream2": false} + for _, source := range eofSources { + if _, exists := expectedSources[source]; !exists { + t.Errorf("Unexpected source name: %s", source) + } else { + expectedSources[source] = true + } + } + + // Verify all expected sources were seen + for source, seen := range expectedSources { + if !seen { + t.Errorf("Did not receive SourceEOF for %s", source) + } + } + + // Verify we received all expected data + if len(receivedData["stream1"]) != 2 { + t.Errorf("Expected 2 items from stream1, got %d", len(receivedData["stream1"])) + } + + if len(receivedData["stream2"]) != 3 { + t.Errorf("Expected 3 items from stream2, got %d", len(receivedData["stream2"])) + } + }) + + t.Run("EmptyStream", func(t *testing.T) { + // Create two streams, one will be empty + sr1, sw1 := Pipe[string](2) + sr2, sw2 := Pipe[string](2) + + // Close the first stream immediately to make it empty + sw1.Close() + + // Merge the streams with names + namedStreams := map[string]*StreamReader[string]{ + "empty": sr1, + "data": sr2, + } + mergedSR := MergeNamedStreamReaders(namedStreams) + mergedSR.SetAutomaticClose() + + // Send data to the second stream + go func() { + defer sw2.Close() + sw2.Send("test-data", nil) + }() + + // Track received EOFs and data + eofSources := make(map[string]bool, 2) + receivedData := make([]string, 0, 1) + + for { + chunk, err := mergedSR.Recv() + if err != nil { + if sourceName, ok := GetSourceName(err); ok { + eofSources[sourceName] = true + continue + } + + if errors.Is(err, io.EOF) { + break + } + + t.Errorf("Error receiving data: %v", err) + break + } + + receivedData = append(receivedData, chunk) + } + + // Verify we received EOF from the empty stream + if len(eofSources) != 2 { + t.Errorf("Expected 2 SourceEOF errors, got %d", len(eofSources)) + } + + if _, exist := eofSources["empty"]; !exist { + t.Errorf("Expected EOF from 'empty' stream, got '%v'", eofSources) + } + if _, exist := eofSources["data"]; !exist { + t.Errorf("Expected EOF from 'data' stream, got '%v'", eofSources) + } + + // Verify we received the data from the non-empty stream + if len(receivedData) != 1 || receivedData[0] != "test-data" { + t.Errorf("Expected to receive 'test-data', got %v", receivedData) + } + }) + + t.Run("ArraySource", func(t *testing.T) { + // Create three named streams + sr1, sw1 := Pipe[string](2) + sr2, sw2 := Pipe[string](2) + sr3 := StreamReaderFromArray([]string{"data3-1", "data3-2", "data3-3"}) + + // Merge the streams with names + namedStreams := map[string]*StreamReader[string]{ + "stream1": sr1, + "stream2": sr2, + "stream3": sr3, + } + mergedSR := MergeNamedStreamReaders(namedStreams) + mergedSR.SetAutomaticClose() + + // Send data and close streams in sequence + go func() { + // First stream sends one item then closes + sw1.Send("data1", nil) + sw1.Close() + + // Second stream sends two items then closes + sw2.Send("data2-1", nil) + sw2.Send("data2-2", nil) + sw2.Close() + }() + + // Track EOF order and data count + eofOrder := make([]string, 0, 3) + dataCount := 0 + + for { + _, err := mergedSR.Recv() + if err != nil { + if sourceName, ok := GetSourceName(err); ok { + eofOrder = append(eofOrder, sourceName) + continue + } + + if errors.Is(err, io.EOF) { + break + } + + t.Errorf("Error receiving data: %v", err) + break + } + + dataCount++ + } + + // Verify EOF count + if len(eofOrder) != 3 { + t.Errorf("Expected 3 SourceEOF errors, got %d", len(eofOrder)) + } + + // Verify data count + if dataCount != 6 { + t.Errorf("Expected 6 data items, got %d", dataCount) + } + }) + + t.Run("ErrorPropagation", func(t *testing.T) { + // Create two streams + sr1, sw1 := Pipe[string](2) + sr2, sw2 := Pipe[string](2) + + // Merge the streams with names + namedStreams := map[string]*StreamReader[string]{ + "normal": sr1, + "error": sr2, + } + mergedSR := MergeNamedStreamReaders(namedStreams) + defer mergedSR.Close() + + testError := errors.New("test error") + + // Send normal data to first stream + go func() { + defer sw1.Close() + sw1.Send("normal-data", nil) + }() + + // Send error to second stream + go func() { + defer sw2.Close() + sw2.Send("", testError) + }() + + // Track received errors + var receivedError error + + for { + _, err := mergedSR.Recv() + if err != nil { + // Skip SourceEOF errors + if _, ok := GetSourceName(err); ok { + continue + } + + if errors.Is(err, io.EOF) { + break + } + + // Store the first non-EOF error + receivedError = err + break + } + } + + // Verify we received the test error + if receivedError == nil || receivedError.Error() != testError.Error() { + t.Errorf("Expected error '%v', got '%v'", testError, receivedError) + } + }) +} diff --git a/schema/tool.go b/schema/tool.go new file mode 100644 index 0000000..170db11 --- /dev/null +++ b/schema/tool.go @@ -0,0 +1,580 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "bytes" + "encoding/gob" + "encoding/json" + "fmt" + "sort" + "strings" + + "github.com/eino-contrib/jsonschema" + orderedmap "github.com/wk8/go-ordered-map/v2" +) + +// DataType is the type of the parameter. +// It must be one of the following values: "object", "number", "integer", "string", "array", "null", "boolean", which is the same as the type of the parameter in JSONSchema. +type DataType string + +// Supported JSONSchema data types for tool parameters. +const ( + Object DataType = "object" + Number DataType = "number" + Integer DataType = "integer" + String DataType = "string" + Array DataType = "array" + Null DataType = "null" + Boolean DataType = "boolean" +) + +// ToolChoice controls how the model uses the tools provided to it. +// Pass as part of the model option via [model.WithToolChoice]. +type ToolChoice string + +const ( + // ToolChoiceForbidden instructs the model not to call any tools, even if + // tools are bound. The model responds with a plain text message instead. + // Corresponds to "none" in OpenAI Chat Completion. + ToolChoiceForbidden ToolChoice = "forbidden" + + // ToolChoiceAllowed lets the model decide: it may generate a plain message + // or call one or more tools. This is the default when tools are provided. + // Corresponds to "auto" in OpenAI Chat Completion. + ToolChoiceAllowed ToolChoice = "allowed" + + // ToolChoiceForced requires the model to call at least one tool. Use this + // when you want to guarantee structured output via tool calling. + // Corresponds to "required" in OpenAI Chat Completion. + ToolChoiceForced ToolChoice = "forced" +) + +type AgenticToolChoice struct { + // Type is the tool choice mode. + Type ToolChoice + + // Allowed optionally specifies the list of tools that the model is permitted to call. + // Optional. + Allowed *AgenticAllowedToolChoice + + // Forced optionally specifies the list of tools that the model is required to call. + // Optional. + Forced *AgenticForcedToolChoice +} + +// AgenticAllowedToolChoice specifies a list of allowed tools for the model. +type AgenticAllowedToolChoice struct { + // Tools is the list of allowed tools for the model to call. + // Optional. + Tools []*AllowedTool +} + +// AgenticForcedToolChoice specifies a list of tools that the model must call. +type AgenticForcedToolChoice struct { + // Tools is the list of tools that the model must call. + // Optional. + Tools []*AllowedTool +} + +// AllowedTool represents a tool that the model is allowed or forced to call. +// Exactly one of FunctionName, MCPTool, or ServerTool must be specified. +type AllowedTool struct { + // FunctionName specifies a function tool by name. + FunctionName string + + // MCPTool specifies an MCP tool. + MCPTool *AllowedMCPTool + + // ServerTool specifies a server tool. + ServerTool *AllowedServerTool +} + +// AllowedMCPTool contains the information for identifying an MCP tool. +type AllowedMCPTool struct { + // ServerLabel is the label of the MCP server. + ServerLabel string + // Name is the name of the MCP tool. + Name string +} + +// AllowedServerTool contains the information for identifying a server tool. +type AllowedServerTool struct { + // Name is the name of the server tool. + Name string +} + +// ToolInfo is the information of a tool. +// ToolInfo describes a tool that can be passed to a ChatModel via +// [ToolCallingChatModel.WithTools] or [ChatModel.BindTools]. +// +// Name should be concise and unique within the tool set. Desc should explain +// when and why to use the tool; few-shot examples in Desc significantly improve +// model accuracy. ParamsOneOf may be nil for tools that take no arguments. +type ToolInfo struct { + // The unique name of the tool that clearly communicates its purpose. + Name string + // Used to tell the model how/when/why to use the tool. + // You can provide few-shot examples as a part of the description. + Desc string + // Extra is the extra information for the tool. + Extra map[string]any + + // The parameters the functions accepts (different models may require different parameter types). + // can be described in two ways: + // - use params: schema.NewParamsOneOfByParams(params) + // - use jsonschema: schema.NewParamsOneOfByJSONSchema(jsonschema) + // If is nil, signals that the tool does not need any input parameter + *ParamsOneOf +} + +type toolInfoForJSON struct { + Name string `json:"name,omitempty"` + Desc string `json:"desc,omitempty"` + Extra map[string]any `json:"extra,omitempty"` + HasParamsOneOf bool `json:"has_params_one_of,omitempty"` + Params map[string]*ParameterInfo `json:"params,omitempty"` + JSONSchema *jsonschema.Schema `json:"json_schema,omitempty"` +} + +type toolInfoForGob struct { + Name string + Desc string + Extra map[string]any + HasParamsOneOf bool + Params map[string]*ParameterInfo + JSONSchema *string +} + +func (t *ToolInfo) MarshalJSON() ([]byte, error) { + tmp := &toolInfoForJSON{ + Name: t.Name, + Desc: t.Desc, + Extra: t.Extra, + } + if t.ParamsOneOf != nil { + tmp.HasParamsOneOf = true + tmp.Params = t.ParamsOneOf.params + tmp.JSONSchema = t.ParamsOneOf.jsonschema + } + return json.Marshal(tmp) +} + +func (t *ToolInfo) UnmarshalJSON(data []byte) error { + tmp := &toolInfoForJSON{} + if err := json.Unmarshal(data, tmp); err != nil { + return err + } + t.Name = tmp.Name + t.Desc = tmp.Desc + t.Extra = tmp.Extra + if tmp.HasParamsOneOf { + t.ParamsOneOf = &ParamsOneOf{ + params: tmp.Params, + jsonschema: tmp.JSONSchema, + } + // An empty-but-non-nil params map is dropped by `omitempty` during + // marshaling. When jsonschema is also absent, the params form was the + // chosen representation, so restore the empty map to preserve the + // roundtrip invariant. + if t.ParamsOneOf.params == nil && t.ParamsOneOf.jsonschema == nil { + t.ParamsOneOf.params = map[string]*ParameterInfo{} + } + } + return nil +} + +func (t *ToolInfo) GobEncode() ([]byte, error) { + tmp := &toolInfoForGob{ + Name: t.Name, + Desc: t.Desc, + Extra: t.Extra, + } + if t.ParamsOneOf != nil { + tmp.HasParamsOneOf = true + tmp.Params = t.ParamsOneOf.params + if t.ParamsOneOf.jsonschema != nil { + b, err := json.Marshal(t.ParamsOneOf.jsonschema) + if err != nil { + return nil, err + } + str := string(b) + tmp.JSONSchema = &str + } + } + buf := new(bytes.Buffer) + if err := gob.NewEncoder(buf).Encode(tmp); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +func (t *ToolInfo) GobDecode(b []byte) error { + tmp := &toolInfoForGob{} + if err := gob.NewDecoder(bytes.NewBuffer(b)).Decode(tmp); err != nil { + return err + } + t.Name = tmp.Name + t.Desc = tmp.Desc + t.Extra = tmp.Extra + if !tmp.HasParamsOneOf { + return nil + } + t.ParamsOneOf = &ParamsOneOf{ + params: tmp.Params, + } + if tmp.JSONSchema != nil { + s := &jsonschema.Schema{} + if err := json.Unmarshal([]byte(*tmp.JSONSchema), s); err != nil { + return err + } + t.ParamsOneOf.jsonschema = s + } + return nil +} + +// ParameterInfo is the information of a parameter. +// It is used to describe the parameters of a tool. +type ParameterInfo struct { + // The type of the parameter. + Type DataType + // The element type of the parameter, only for array. + ElemInfo *ParameterInfo + // The sub parameters of the parameter, only for object. + SubParams map[string]*ParameterInfo + // The description of the parameter. + Desc string + // The enum values of the parameter, only for string. + Enum []string + // Whether the parameter is required. + Required bool +} + +// ParamsOneOf holds a tool's parameter schema using exactly one of two +// representations. Choose the one that best fits your needs: +// +// 1. [NewParamsOneOfByParams] — lightweight: describe parameters as a +// map[string]*[ParameterInfo]. Covers the most common cases (scalars, +// arrays, nested objects, enums, required flags). +// +// 2. [NewParamsOneOfByJSONSchema] — powerful: supply a full +// *jsonschema.Schema (JSON Schema 2020-12). Required when you need +// features not expressible via ParameterInfo, such as anyOf, oneOf, or +// $defs references. [utils.InferTool] generates this form automatically +// from Go struct tags. +// +// You must use exactly one constructor — setting both fields is invalid. +// If ParamsOneOf is nil, the tool takes no input parameters. +type ParamsOneOf struct { + // use NewParamsOneOfByParams to set this field + params map[string]*ParameterInfo + + jsonschema *jsonschema.Schema +} + +// NewParamsOneOfByParams creates a ParamsOneOf with map[string]*ParameterInfo. +func NewParamsOneOfByParams(params map[string]*ParameterInfo) *ParamsOneOf { + return &ParamsOneOf{ + params: params, + } +} + +// NewParamsOneOfByJSONSchema creates a ParamsOneOf with *jsonschema.Schema. +func NewParamsOneOfByJSONSchema(s *jsonschema.Schema) *ParamsOneOf { + return &ParamsOneOf{ + jsonschema: s, + } +} + +// ToJSONSchema parses ParamsOneOf, converts the parameter description that user actually provides, into the format ready to be passed to Model. +func (p *ParamsOneOf) ToJSONSchema() (*jsonschema.Schema, error) { + if p == nil { + return nil, nil + } + + if p.params != nil { + sc := &jsonschema.Schema{ + Properties: orderedmap.New[string, *jsonschema.Schema](), + Type: string(Object), + Required: make([]string, 0, len(p.params)), + } + + keys := make([]string, 0, len(p.params)) + for k := range p.params { + keys = append(keys, k) + } + sort.Strings(keys) + + for _, k := range keys { + v := p.params[k] + sc.Properties.Set(k, paramInfoToJSONSchema(v)) + if v.Required { + sc.Required = append(sc.Required, k) + } + } + + return sc, nil + } + + return p.jsonschema, nil +} + +func paramInfoToJSONSchema(paramInfo *ParameterInfo) *jsonschema.Schema { + js := &jsonschema.Schema{ + Type: string(paramInfo.Type), + Description: paramInfo.Desc, + } + + if len(paramInfo.Enum) > 0 { + js.Enum = make([]any, len(paramInfo.Enum)) + for i, enum := range paramInfo.Enum { + js.Enum[i] = enum + } + } + + if paramInfo.ElemInfo != nil { + js.Items = paramInfoToJSONSchema(paramInfo.ElemInfo) + } + + if len(paramInfo.SubParams) > 0 { + required := make([]string, 0, len(paramInfo.SubParams)) + js.Properties = orderedmap.New[string, *jsonschema.Schema]() + keys := make([]string, 0, len(paramInfo.SubParams)) + for k := range paramInfo.SubParams { + keys = append(keys, k) + } + sort.Strings(keys) + + for _, k := range keys { + v := paramInfo.SubParams[k] + item := paramInfoToJSONSchema(v) + js.Properties.Set(k, item) + if v.Required { + required = append(required, k) + } + } + + js.Required = required + } + + return js +} + +// ToolPartType defines the type of content in a tool output part. +// It is used to distinguish between different types of multimodal content returned by tools. +type ToolPartType string + +const ( + // ToolPartTypeText means the part is a text. + ToolPartTypeText ToolPartType = "text" + + // ToolPartTypeImage means the part is an image url. + ToolPartTypeImage ToolPartType = "image" + + // ToolPartTypeAudio means the part is an audio url. + ToolPartTypeAudio ToolPartType = "audio" + + // ToolPartTypeVideo means the part is a video url. + ToolPartTypeVideo ToolPartType = "video" + + // ToolPartTypeFile means the part is a file url. + ToolPartTypeFile ToolPartType = "file" + + // ToolPartTypeToolSearchResult means the part contains tool search results. + ToolPartTypeToolSearchResult ToolPartType = "tool_search_result" +) + +// ToolOutputImage represents an image in tool output. +// It contains URL or Base64-encoded data along with MIME type information. +type ToolOutputImage struct { + MessagePartCommon +} + +// ToolOutputAudio represents an audio file in tool output. +// It contains URL or Base64-encoded data along with MIME type information. +type ToolOutputAudio struct { + MessagePartCommon +} + +// ToolOutputVideo represents a video file in tool output. +// It contains URL or Base64-encoded data along with MIME type information. +type ToolOutputVideo struct { + MessagePartCommon +} + +// ToolOutputFile represents a generic file in tool output. +// It contains URL or Base64-encoded data along with MIME type information. +type ToolOutputFile struct { + MessagePartCommon +} + +// ToolSearchResult represents the result of a tool search operation. +// When a model issues a tool search call, the framework searches for matching tools +// and returns the results via this struct. +type ToolSearchResult struct { + // Tools contains the full definitions of matched tools that were not previously + // registered. Their complete definitions are required so that the model can + // understand their parameters and usage. + Tools []*ToolInfo +} + +func (t *ToolSearchResult) String() string { + sb := new(strings.Builder) + sb.WriteString("ToolSearchResult[") + for _, tool := range t.Tools { + sb.WriteString(tool.Name) + sb.WriteString(",") + } + sb.WriteString("]") + return sb.String() +} + +// ToolOutputPart represents a part of tool execution output. +// It supports streaming scenarios through the Index field for chunk merging. +type ToolOutputPart struct { + + // Type is the type of the part, e.g., "text", "image_url", "audio_url", "video_url". + Type ToolPartType `json:"type"` + + // Text is the text content, used when Type is "text". + Text string `json:"text,omitempty"` + + // Image is the image content, used when Type is ToolPartTypeImage. + Image *ToolOutputImage `json:"image,omitempty"` + + // Audio is the audio content, used when Type is ToolPartTypeAudio. + Audio *ToolOutputAudio `json:"audio,omitempty"` + + // Video is the video content, used when Type is ToolPartTypeVideo. + Video *ToolOutputVideo `json:"video,omitempty"` + + // File is the file content, used when Type is ToolPartTypeFile. + File *ToolOutputFile `json:"file,omitempty"` + + // ToolSearchResult holds the tool search results, used when Type is ToolPartTypeToolSearchResult. + ToolSearchResult *ToolSearchResult `json:"tool_search_result,omitempty"` + + // Extra is used to store extra information. + Extra map[string]any `json:"extra,omitempty"` +} + +// ToolArgument contains the input information for a tool call. +// It is used to pass tool call arguments to enhanced tools. +type ToolArgument struct { + // Text contains the arguments for the tool call in JSON format. + Text string `json:"text,omitempty"` +} + +// ToolResult represents the structured multimodal output from a tool execution. +// It is used when a tool needs to return more than just a simple string, +// such as images, files, or other structured data. +type ToolResult struct { + // Parts contains the multimodal output parts. Each part can be a different + // type of content, like text, an image, or a file. + Parts []ToolOutputPart `json:"parts,omitempty"` +} + +func convToolOutputPartToMessageInputPart(toolPart ToolOutputPart) (MessageInputPart, error) { + switch toolPart.Type { + case ToolPartTypeText: + return MessageInputPart{ + Type: ChatMessagePartTypeText, + Text: toolPart.Text, + Extra: toolPart.Extra, + }, nil + case ToolPartTypeImage: + if toolPart.Image == nil { + return MessageInputPart{}, fmt.Errorf("image content is nil for tool part type %v", toolPart.Type) + } + return MessageInputPart{ + Type: ChatMessagePartTypeImageURL, + Image: &MessageInputImage{MessagePartCommon: toolPart.Image.MessagePartCommon}, + Extra: toolPart.Extra, + }, nil + case ToolPartTypeAudio: + if toolPart.Audio == nil { + return MessageInputPart{}, fmt.Errorf("audio content is nil for tool part type %v", toolPart.Type) + } + return MessageInputPart{ + Type: ChatMessagePartTypeAudioURL, + Audio: &MessageInputAudio{MessagePartCommon: toolPart.Audio.MessagePartCommon}, + Extra: toolPart.Extra, + }, nil + case ToolPartTypeVideo: + if toolPart.Video == nil { + return MessageInputPart{}, fmt.Errorf("video content is nil for tool part type %v", toolPart.Type) + } + return MessageInputPart{ + Type: ChatMessagePartTypeVideoURL, + Video: &MessageInputVideo{MessagePartCommon: toolPart.Video.MessagePartCommon}, + Extra: toolPart.Extra, + }, nil + case ToolPartTypeFile: + if toolPart.File == nil { + return MessageInputPart{}, fmt.Errorf("file content is nil for tool part type %v", toolPart.Type) + } + return MessageInputPart{ + Type: ChatMessagePartTypeFileURL, + File: &MessageInputFile{MessagePartCommon: toolPart.File.MessagePartCommon}, + Extra: toolPart.Extra, + }, nil + case ToolPartTypeToolSearchResult: + if toolPart.ToolSearchResult == nil { + return MessageInputPart{}, fmt.Errorf("tool search result is nil for tool part type %v", toolPart.Type) + } + return MessageInputPart{ + Type: ChatMessagePartTypeToolSearchResult, + ToolSearchResult: toolPart.ToolSearchResult, + }, nil + default: + return MessageInputPart{}, fmt.Errorf("unknown tool part type: %v", toolPart.Type) + } +} + +// ToMessageInputParts converts ToolOutputPart slice to MessageInputPart slice. +// This is used when passing tool results as input to the model. +// +// Parameters: +// - None (method receiver is *ToolResult) +// +// Returns: +// - []MessageInputPart: The converted message input parts that can be used in a Message. +// - error: An error if conversion fails due to unknown part types or nil content fields. +// +// Example: +// +// toolResult := &schema.ToolResult{ +// Parts: []schema.ToolOutputPart{ +// {Type: schema.ToolPartTypeText, Text: "Result text"}, +// {Type: schema.ToolPartTypeImage, Image: &schema.ToolOutputImage{...}}, +// }, +// } +// inputParts, err := toolResult.ToMessageInputParts() +func (tr *ToolResult) ToMessageInputParts() ([]MessageInputPart, error) { + if tr == nil || len(tr.Parts) == 0 { + return nil, nil + } + result := make([]MessageInputPart, len(tr.Parts)) + for i, part := range tr.Parts { + var err error + result[i], err = convToolOutputPartToMessageInputPart(part) + if err != nil { + return nil, err + } + } + return result, nil +} diff --git a/schema/tool_test.go b/schema/tool_test.go new file mode 100644 index 0000000..76a6a25 --- /dev/null +++ b/schema/tool_test.go @@ -0,0 +1,235 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package schema + +import ( + "bytes" + "encoding/gob" + "encoding/json" + "testing" + + "github.com/eino-contrib/jsonschema" + "github.com/smartystreets/goconvey/convey" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParamsOneOfToJSONSchema(t *testing.T) { + convey.Convey("ParamsOneOfToJSONSchema", t, func() { + var ( + oneOf ParamsOneOf + converted any + err error + ) + + convey.Convey("user provides JSON schema directly, use what the user provides", func() { + oneOf.jsonschema = &jsonschema.Schema{ + Type: "string", + Description: "this is the only argument", + } + converted, err = oneOf.ToJSONSchema() + convey.So(err, convey.ShouldBeNil) + convey.So(converted, convey.ShouldResemble, oneOf.jsonschema) + }) + + convey.Convey("user provides map[string]ParameterInfo, converts to json schema", func() { + oneOf.params = map[string]*ParameterInfo{ + "arg1": { + Type: String, + Desc: "this is the first argument", + Required: true, + Enum: []string{"1", "2"}, + }, + "arg2": { + Type: Object, + Desc: "this is the second argument", + SubParams: map[string]*ParameterInfo{ + "sub_arg1": { + Type: String, + Desc: "this is the sub argument", + Required: true, + Enum: []string{"1", "2"}, + }, + "sub_arg2": { + Type: String, + Desc: "this is the sub argument 2", + }, + }, + Required: true, + }, + "arg3": { + Type: Array, + Desc: "this is the third argument", + ElemInfo: &ParameterInfo{ + Type: String, + Desc: "this is the element of the third argument", + Required: true, + Enum: []string{"1", "2"}, + }, + Required: true, + }, + } + converted, err = oneOf.ToJSONSchema() + convey.So(err, convey.ShouldBeNil) + }) + + convey.Convey("user provides map[string]ParameterInfo, converts to json schema in order", func() { + params := &ParamsOneOf{ + params: map[string]*ParameterInfo{ + "c": { + Type: "string", + }, + "a": { + Type: "object", + SubParams: map[string]*ParameterInfo{ + "z": { + Type: "number", + }, + "y": { + Type: "string", + }, + }, + }, + "b": { + Type: "array", + ElemInfo: &ParameterInfo{ + Type: "object", + SubParams: map[string]*ParameterInfo{ + "p": { + Type: "integer", + }, + "o": { + Type: "boolean", + }, + }, + }, + }, + }, + } + + schema1, err := params.ToJSONSchema() + assert.NoError(t, err) + json1, err := json.Marshal(schema1) + assert.NoError(t, err) + + schema2, err := params.ToJSONSchema() + assert.NoError(t, err) + json2, err := json.Marshal(schema2) + assert.NoError(t, err) + + assert.Equal(t, string(json1), string(json2)) + }) + + }) +} + +func TestToolInfoSerialization(t *testing.T) { + ti1 := &ToolInfo{ + ParamsOneOf: NewParamsOneOfByParams(map[string]*ParameterInfo{ + "a": { + Type: String, + Desc: "desc", + }, + }), + } + ti2 := &ToolInfo{ + ParamsOneOf: NewParamsOneOfByJSONSchema(&jsonschema.Schema{ + Type: "string", + }), + } + + // json + b, err := json.Marshal(ti1) + assert.NoError(t, err) + result := &ToolInfo{} + err = json.Unmarshal(b, result) + assert.NoError(t, err) + assert.Equal(t, ti1, result) + b, err = json.Marshal(ti2) + assert.NoError(t, err) + result = &ToolInfo{} + err = json.Unmarshal(b, result) + assert.NoError(t, err) + assert.Equal(t, ti2, result) + + // gob + buf := new(bytes.Buffer) + err = gob.NewEncoder(buf).Encode(ti1) + assert.NoError(t, err) + result = &ToolInfo{} + err = gob.NewDecoder(buf).Decode(result) + assert.NoError(t, err) + assert.Equal(t, ti1, result) + buf = new(bytes.Buffer) + err = gob.NewEncoder(buf).Encode(ti2) + assert.NoError(t, err) + result = &ToolInfo{} + err = gob.NewDecoder(buf).Decode(result) + assert.NoError(t, err) + assert.Equal(t, ti2, result) + + // json roundtrip with empty-but-non-nil params map: must not collapse to nil, + // otherwise the params form is silently dropped. + tiEmpty := &ToolInfo{ + ParamsOneOf: NewParamsOneOfByParams(map[string]*ParameterInfo{}), + } + b, err = json.Marshal(tiEmpty) + assert.NoError(t, err) + result = &ToolInfo{} + err = json.Unmarshal(b, result) + assert.NoError(t, err) + assert.NotNil(t, result.ParamsOneOf) + assert.NotNil(t, result.ParamsOneOf.params) + assert.Equal(t, tiEmpty, result) +} + +func TestMCPToolResult_NilErrorCode(t *testing.T) { + result := &MCPToolResult{ + CallID: "test-call", + Name: "test-tool", + Content: "some result", + Error: &MCPToolCallError{ + Code: nil, + Message: "something went wrong", + }, + } + + require.NotPanics(t, func() { + s := result.String() + t.Logf("String output: %s", s) + assert.Contains(t, s, "something went wrong") + }, "BUG: MCPToolResult.String() should not panic when Error.Code is nil") +} + +func TestMCPToolResult_WithErrorCode(t *testing.T) { + code := int64(500) + result := &MCPToolResult{ + CallID: "test-call", + Name: "test-tool", + Content: "", + Error: &MCPToolCallError{ + Code: &code, + Message: "internal server error", + }, + } + + require.NotPanics(t, func() { + s := result.String() + assert.Contains(t, s, "500") + assert.Contains(t, s, "internal server error") + }) +} diff --git a/scripts/dev_setup.sh b/scripts/dev_setup.sh new file mode 100755 index 0000000..0be6017 --- /dev/null +++ b/scripts/dev_setup.sh @@ -0,0 +1,122 @@ +#!/usr/bin/env bash +# dev_setup.sh — set up a local multi-module workspace for eino development. +# +# BACKGROUND +# eino, eino-ext, and eino-examples live in separate GitHub repositories to +# keep their Go modules, versioning, and maintenance independent. However, +# working across them is inconvenient: editors and AI coding tools lack +# cross-repo type information and can't navigate between them. +# +# This script brings all three repos together locally so that a single +# go.work file provides full cross-module LSP (go-to-definition, type +# inference, autocomplete) across all ~83 modules — without touching any +# remote repository. +# +# WHAT IT DOES +# 1. Clones eino-ext → ext/ +# 2. Clones eino-examples → examples/ +# 3. Registers ext/ and examples/ in .git/info/exclude so eino's git +# never sees them (local-only, never committed) +# 4. Creates go.work at the repo root covering eino + all modules in +# ext/ and examples/ (go.work is already in .gitignore) +# +# RESULTING LAYOUT +# eino/ ← you are here (github.com/cloudwego/eino) +# eino/ext/ ← github.com/cloudwego/eino-ext (full git repo) +# eino/examples/ ← github.com/cloudwego/eino-examples (full git repo) +# eino/go.work ← wires all modules together (gitignored) +# +# WORKING ACROSS REPOS +# Each subdirectory is a full independent git repo tracking its own remote. +# To contribute to eino-ext or eino-examples, work inside that directory: +# +# cd ext +# git checkout -b feat/my-feature +# # make changes — editor has full cross-repo type info via go.work +# git commit -m "feat: ..." +# git push origin feat/my-feature # pushes to cloudwego/eino-ext +# +# KEEPING REPOS UP TO DATE +# git -C ext pull +# git -C examples pull +# +# USAGE +# bash scripts/dev_setup.sh # first-time setup +# bash scripts/dev_setup.sh --reset # re-clone everything from scratch + +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$REPO_ROOT" + +EXT_DIR="ext" +EXAMPLES_DIR="examples" +EINO_EXT_REPO="https://github.com/cloudwego/eino-ext" +EINO_EXAMPLES_REPO="https://github.com/cloudwego/eino-examples" + +# Parse flags +RESET=false +for arg in "$@"; do + case $arg in + --reset) RESET=true ;; + esac +done + +echo "==> Setting up eino dev workspace in: $REPO_ROOT" + +# --reset: remove existing dirs +if [ "$RESET" = true ]; then + echo "==> --reset: removing $EXT_DIR/ and $EXAMPLES_DIR/" + rm -rf "$EXT_DIR" "$EXAMPLES_DIR" go.work go.work.sum +fi + +# Clone repos if not already present +if [ ! -d "$EXT_DIR/.git" ]; then + echo "==> Cloning eino-ext into $EXT_DIR/" + git clone "$EINO_EXT_REPO" "$EXT_DIR" +else + echo "==> $EXT_DIR/ already exists, skipping clone" +fi + +if [ ! -d "$EXAMPLES_DIR/.git" ]; then + echo "==> Cloning eino-examples into $EXAMPLES_DIR/" + git clone "$EINO_EXAMPLES_REPO" "$EXAMPLES_DIR" +else + echo "==> $EXAMPLES_DIR/ already exists, skipping clone" +fi + +# Exclude dirs from eino's git tracking (local only, not committed) +EXCLUDE_FILE=".git/info/exclude" +add_exclude() { + local entry="$1" + if ! grep -qxF "$entry" "$EXCLUDE_FILE" 2>/dev/null; then + echo "$entry" >> "$EXCLUDE_FILE" + echo "==> Added '$entry' to $EXCLUDE_FILE" + fi +} +add_exclude "$EXT_DIR/" +add_exclude "$EXAMPLES_DIR/" + +# Build go.work covering eino root + every go.mod found in ext/ and examples/ +if [ ! -f "go.work" ]; then + echo "==> Creating go.work" + go work init . + + # Collect all module directories (directories containing a go.mod) + while IFS= read -r modfile; do + dir="$(dirname "$modfile")" + go work use "$dir" + done < <(find "$EXT_DIR" "$EXAMPLES_DIR" -name "go.mod" | sort) + + echo "==> go.work created with $(grep -c '^\s\+\.' go.work || true) module(s)" +else + echo "==> go.work already exists, skipping (use --reset to recreate)" +fi + +echo "" +echo "Done. Your workspace includes:" +echo " . — github.com/cloudwego/eino" +echo " $EXT_DIR/ — github.com/cloudwego/eino-ext ($(find "$EXT_DIR" -name "go.mod" | wc -l | tr -d ' ') modules)" +echo " $EXAMPLES_DIR/ — github.com/cloudwego/eino-examples ($(find "$EXAMPLES_DIR" -name "go.mod" | wc -l | tr -d ' ') modules)" +echo "" +echo "Run 'go build ./...' or open this directory in your editor." diff --git a/scripts/eino_setup.sh b/scripts/eino_setup.sh new file mode 100755 index 0000000..5b99bed --- /dev/null +++ b/scripts/eino_setup.sh @@ -0,0 +1,180 @@ +#!/usr/bin/env bash +# eino_setup.sh — fetch eino framework source into your project for AI-assisted development. +# +# BACKGROUND +# When building applications with eino, your AI coding assistant (Claude Code, +# Cursor, Copilot, etc.) only sees your code. It cannot navigate into eino's +# source to understand how components work, what patterns are idiomatic, or +# how to wire things together correctly. +# +# This script clones eino, eino-ext, and eino-examples into a _eino/ directory +# inside your project. Your AI assistant can then browse the actual source, +# examples, and extensions — giving it full context to help you build correctly. +# +# WHAT IT DOES +# 1. Clones eino → _eino/eino/ +# 2. Clones eino-ext → _eino/eino-ext/ +# 3. Clones eino-examples → _eino/eino-examples/ +# 4. Adds _eino/ to .gitignore (read-only reference, never committed) +# 5. Writes a _eino/README.md explaining the directory to future readers +# +# RESULTING LAYOUT +# your-project/ +# ├── _eino/ +# │ ├── eino/ ← github.com/cloudwego/eino (core framework) +# │ ├── eino-ext/ ← github.com/cloudwego/eino-ext (components & integrations) +# │ └── eino-examples/ ← github.com/cloudwego/eino-examples (patterns & recipes) +# └── ... your code +# +# NOTE: _eino/ is read-only reference material. Do not edit files inside it. +# Your go.mod is unchanged — eino remains a normal dependency. +# +# KEEPING UP TO DATE +# bash eino_setup.sh --update # pull latest on all three repos +# +# USAGE +# bash eino_setup.sh # first-time setup +# bash eino_setup.sh --reset # re-clone everything from scratch +# bash eino_setup.sh --update # pull latest without re-cloning +# +# SYSTEM PROMPT +# After running this script, add the following to your AI assistant's project +# instructions (CLAUDE.md, .cursorrules, .github/copilot-instructions.md, etc.): +# +# --- +# ## eino Framework Reference +# +# This project uses the eino framework (github.com/cloudwego/eino). +# The full framework source is available locally in `_eino/`: +# +# - `_eino/eino/` — core framework (components, graph, compose, callbacks) +# - `_eino/eino-ext/` — official components and integrations (models, tools, retrievers, etc.) +# - `_eino/eino-examples/` — working examples and patterns +# +# When answering questions about eino APIs, component wiring, graph construction, +# callbacks, or any eino-specific patterns: explore `_eino/` first. +# Prefer examples from `_eino/eino-examples/` as the canonical reference for +# idiomatic usage. +# --- + +set -euo pipefail + +PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$PROJECT_ROOT" + +EINO_DIR="_eino" +EINO_REPO="https://github.com/cloudwego/eino" +EINO_EXT_REPO="https://github.com/cloudwego/eino-ext" +EINO_EXAMPLES_REPO="https://github.com/cloudwego/eino-examples" + +# Parse flags +RESET=false +UPDATE=false +for arg in "$@"; do + case $arg in + --reset) RESET=true ;; + --update) UPDATE=true ;; + esac +done + +echo "==> eino setup in: $PROJECT_ROOT" + +# --reset: remove and re-clone +if [ "$RESET" = true ]; then + echo "==> --reset: removing $EINO_DIR/" + rm -rf "$EINO_DIR" +fi + +# --update: pull latest on existing clones +if [ "$UPDATE" = true ]; then + for repo in eino eino-ext eino-examples; do + dir="$EINO_DIR/$repo" + if [ -d "$dir/.git" ]; then + echo "==> Updating $dir/" + git -C "$dir" pull --ff-only + else + echo "==> $dir/ not found, skipping update (run without --update to clone)" + fi + done + echo "" + echo "Done. Run 'bash eino_setup.sh' to clone any missing repos." + exit 0 +fi + +mkdir -p "$EINO_DIR" + +# Clone repos (shallow — we only need source to read, not full history) +clone_if_missing() { + local repo_url="$1" + local dest="$2" + if [ ! -d "$dest/.git" ]; then + echo "==> Cloning $(basename "$dest")/" + git clone --depth=1 "$repo_url" "$dest" + else + echo "==> $dest/ already exists, skipping clone" + fi +} + +clone_if_missing "$EINO_REPO" "$EINO_DIR/eino" +clone_if_missing "$EINO_EXT_REPO" "$EINO_DIR/eino-ext" +clone_if_missing "$EINO_EXAMPLES_REPO" "$EINO_DIR/eino-examples" + +# Add _eino/ to .gitignore +GITIGNORE=".gitignore" +if ! grep -qxF "$EINO_DIR/" "$GITIGNORE" 2>/dev/null; then + echo "" >> "$GITIGNORE" + echo "# eino framework source (AI coding reference — see eino_setup.sh)" >> "$GITIGNORE" + echo "$EINO_DIR/" >> "$GITIGNORE" + echo "==> Added '$EINO_DIR/' to $GITIGNORE" +fi + +# Write a README so the directory is self-explanatory +cat > "$EINO_DIR/README.md" <<'EOF' +# _eino — eino framework source reference + +This directory contains read-only clones of the eino framework repositories, +checked out for use by AI coding assistants (Claude Code, Cursor, Copilot, etc.). + +| Directory | Repository | Purpose | +|----------------|-----------------------------------------|--------------------------------| +| `eino/` | github.com/cloudwego/eino | Core framework source | +| `eino-ext/` | github.com/cloudwego/eino-ext | Components and integrations | +| `eino-examples/` | github.com/cloudwego/eino-examples | Patterns, recipes, and samples | + +**Do not edit files here.** This directory is in `.gitignore` and is never committed. + +To update to the latest: + + bash eino_setup.sh --update + +To re-clone from scratch: + + bash eino_setup.sh --reset +EOF + +echo "" +echo "Done. Your AI assistant now has full eino context in $EINO_DIR/:" +echo " $EINO_DIR/eino/ — core framework ($(find "$EINO_DIR/eino" -name "*.go" | wc -l | tr -d ' ') .go files)" +echo " $EINO_DIR/eino-ext/ — components & integrations ($(find "$EINO_DIR/eino-ext" -name "*.go" | wc -l | tr -d ' ') .go files)" +echo " $EINO_DIR/eino-examples/ — patterns & recipes ($(find "$EINO_DIR/eino-examples" -name "*.go" | wc -l | tr -d ' ') .go files)" +echo "" +echo "Add the following to your AI assistant's system prompt or project instructions" +echo "(e.g. CLAUDE.md, .cursorrules, .github/copilot-instructions.md):" +echo "" +echo "---" +cat <<'PROMPT' +## eino Framework Reference + +This project uses the eino framework (github.com/cloudwego/eino). +The full framework source is available locally in `_eino/`: + +- `_eino/eino/` — core framework (components, graph, compose, callbacks) +- `_eino/eino-ext/` — official components and integrations (models, tools, retrievers, etc.) +- `_eino/eino-examples/` — working examples and patterns + +When answering questions about eino APIs, component wiring, graph construction, +callbacks, or any eino-specific patterns: explore `_eino/` first. +Prefer examples from `_eino/eino-examples/` as the canonical reference for +idiomatic usage. +PROMPT +echo "---" diff --git a/utils/callbacks/template.go b/utils/callbacks/template.go new file mode 100644 index 0000000..850e301 --- /dev/null +++ b/utils/callbacks/template.go @@ -0,0 +1,793 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Package callbacks provides ready-to-use callback handler templates for components. +package callbacks + +import ( + "context" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +// NewHandlerHelper creates a new component template handler builder. +// This builder can be used to configure and build a component template handler, +// which can handle callback events for different components with its own struct definition, +// and fallbackTemplate can be used to handle scenarios where none of the cases are hit as a fallback. +func NewHandlerHelper() *HandlerHelper { + return &HandlerHelper{ + composeTemplates: map[components.Component]callbacks.Handler{}, + } +} + +// HandlerHelper is a builder for creating a callbacks.Handler with specific handlers for different component types. +// create a handler with callbacks.NewHandlerHelper(). +// eg. +// +// helper := template.NewHandlerHelper(). +// ChatModel(&model.IndexerCallbackHandler{}). +// Prompt(&prompt.IndexerCallbackHandler{}). +// Handler() +// +// then use the handler with runnable.Invoke(ctx, input, compose.WithCallbacks(handler)) +type HandlerHelper struct { + promptHandler *PromptCallbackHandler + chatModelHandler *ModelCallbackHandler + embeddingHandler *EmbeddingCallbackHandler + indexerHandler *IndexerCallbackHandler + retrieverHandler *RetrieverCallbackHandler + loaderHandler *LoaderCallbackHandler + transformerHandler *TransformerCallbackHandler + toolHandler *ToolCallbackHandler + toolsNodeHandler *ToolsNodeCallbackHandlers + agentHandler *AgentCallbackHandler + agenticAgentHandler *AgenticAgentCallbackHandler + agenticPromptHandler *AgenticPromptCallbackHandler + agenticModelHandler *AgenticModelCallbackHandler + agenticToolsNodeHandler *AgenticToolsNodeCallbackHandlers + composeTemplates map[components.Component]callbacks.Handler +} + +// Handler returns the callbacks.Handler created by HandlerHelper. +func (c *HandlerHelper) Handler() callbacks.Handler { + return &handlerTemplate{c} +} + +// Prompt sets the prompt handler for the handler helper, which will be called when the prompt component is executed. +func (c *HandlerHelper) Prompt(handler *PromptCallbackHandler) *HandlerHelper { + c.promptHandler = handler + return c +} + +// ChatModel sets the chat model handler for the handler helper, which will be called when the chat model component is executed. +func (c *HandlerHelper) ChatModel(handler *ModelCallbackHandler) *HandlerHelper { + c.chatModelHandler = handler + return c +} + +// Embedding sets the embedding handler for the handler helper, which will be called when the embedding component is executed. +func (c *HandlerHelper) Embedding(handler *EmbeddingCallbackHandler) *HandlerHelper { + c.embeddingHandler = handler + return c +} + +// Indexer sets the indexer handler for the handler helper, which will be called when the indexer component is executed. +func (c *HandlerHelper) Indexer(handler *IndexerCallbackHandler) *HandlerHelper { + c.indexerHandler = handler + return c +} + +// Retriever sets the retriever handler for the handler helper, which will be called when the retriever component is executed. +func (c *HandlerHelper) Retriever(handler *RetrieverCallbackHandler) *HandlerHelper { + c.retrieverHandler = handler + return c +} + +// Loader sets the loader handler for the handler helper, which will be called when the loader component is executed. +func (c *HandlerHelper) Loader(handler *LoaderCallbackHandler) *HandlerHelper { + c.loaderHandler = handler + return c +} + +// Transformer sets the transformer handler for the handler helper, which will be called when the transformer component is executed. +func (c *HandlerHelper) Transformer(handler *TransformerCallbackHandler) *HandlerHelper { + c.transformerHandler = handler + return c +} + +// Tool sets the tool handler for the handler helper, which will be called when the tool component is executed. +func (c *HandlerHelper) Tool(handler *ToolCallbackHandler) *HandlerHelper { + c.toolHandler = handler + return c +} + +// ToolsNode sets the tools node handler for the handler helper, which will be called when the tools node is executed. +func (c *HandlerHelper) ToolsNode(handler *ToolsNodeCallbackHandlers) *HandlerHelper { + c.toolsNodeHandler = handler + return c +} + +// AgenticPrompt sets the agentic prompt handler for the handler helper, which will be called when the agentic prompt component is executed. +func (c *HandlerHelper) AgenticPrompt(handler *AgenticPromptCallbackHandler) *HandlerHelper { + c.agenticPromptHandler = handler + return c +} + +// AgenticModel sets the agentic chat model handler for the handler helper, which will be called when the agentic chat model component is executed. +func (c *HandlerHelper) AgenticModel(handler *AgenticModelCallbackHandler) *HandlerHelper { + c.agenticModelHandler = handler + return c +} + +// AgenticToolsNode sets the agentic tools node handler for the handler helper, which will be called when the agentic tools node is executed. +func (c *HandlerHelper) AgenticToolsNode(handler *AgenticToolsNodeCallbackHandlers) *HandlerHelper { + c.agenticToolsNodeHandler = handler + return c +} + +// Agent sets the agent handler for the handler helper, which will be called when the agent is executed. +func (c *HandlerHelper) Agent(handler *AgentCallbackHandler) *HandlerHelper { + c.agentHandler = handler + return c +} + +// AgenticAgent sets the agentic agent callback handler for the handler helper, which will be called when an agentic agent is executed. +func (c *HandlerHelper) AgenticAgent(handler *AgenticAgentCallbackHandler) *HandlerHelper { + c.agenticAgentHandler = handler + return c +} + +// Graph sets the graph handler for the handler helper, which will be called when the graph is executed. +func (c *HandlerHelper) Graph(handler callbacks.Handler) *HandlerHelper { + c.composeTemplates[compose.ComponentOfGraph] = handler + return c +} + +// Chain sets the chain handler for the handler helper, which will be called when the chain is executed. +func (c *HandlerHelper) Chain(handler callbacks.Handler) *HandlerHelper { + c.composeTemplates[compose.ComponentOfChain] = handler + return c +} + +// Lambda sets the lambda handler for the handler helper, which will be called when the lambda is executed. +func (c *HandlerHelper) Lambda(handler callbacks.Handler) *HandlerHelper { + c.composeTemplates[compose.ComponentOfLambda] = handler + return c +} + +type handlerTemplate struct { + *HandlerHelper +} + +// OnStart is the callback function for the start event of a component. +// implement the callbacks Handler interface. +func (c *handlerTemplate) OnStart(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + switch info.Component { + case components.ComponentOfPrompt: + return c.promptHandler.OnStart(ctx, info, prompt.ConvCallbackInput(input)) + case components.ComponentOfAgenticPrompt: + return c.agenticPromptHandler.OnStart(ctx, info, prompt.ConvCallbackInput(input)) + case components.ComponentOfChatModel: + return c.chatModelHandler.OnStart(ctx, info, model.ConvCallbackInput(input)) + case components.ComponentOfAgenticModel: + return c.agenticModelHandler.OnStart(ctx, info, model.ConvAgenticCallbackInput(input)) + case components.ComponentOfEmbedding: + return c.embeddingHandler.OnStart(ctx, info, embedding.ConvCallbackInput(input)) + case components.ComponentOfIndexer: + return c.indexerHandler.OnStart(ctx, info, indexer.ConvCallbackInput(input)) + case components.ComponentOfRetriever: + return c.retrieverHandler.OnStart(ctx, info, retriever.ConvCallbackInput(input)) + case components.ComponentOfLoader: + return c.loaderHandler.OnStart(ctx, info, document.ConvLoaderCallbackInput(input)) + case components.ComponentOfTransformer: + return c.transformerHandler.OnStart(ctx, info, document.ConvTransformerCallbackInput(input)) + case components.ComponentOfTool: + return c.toolHandler.OnStart(ctx, info, tool.ConvCallbackInput(input)) + case compose.ComponentOfToolsNode: + return c.toolsNodeHandler.OnStart(ctx, info, convToolsNodeCallbackInput(input)) + case compose.ComponentOfAgenticToolsNode: + return c.agenticToolsNodeHandler.OnStart(ctx, info, convAgenticToolsNodeCallbackInput(input)) + case adk.ComponentOfAgent: + return c.agentHandler.OnStart(ctx, info, adk.ConvAgentCallbackInput(input)) + case adk.ComponentOfAgenticAgent: + return c.agenticAgentHandler.OnStart(ctx, info, adk.ConvTypedCallbackInput[*schema.AgenticMessage](input)) + case compose.ComponentOfGraph, + compose.ComponentOfChain, + compose.ComponentOfLambda: + return c.composeTemplates[info.Component].OnStart(ctx, info, input) + default: + return ctx + } +} + +// OnEnd is the callback function for the end event of a component. +// implement the callbacks Handler interface. +func (c *handlerTemplate) OnEnd(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + switch info.Component { + case components.ComponentOfPrompt: + return c.promptHandler.OnEnd(ctx, info, prompt.ConvCallbackOutput(output)) + case components.ComponentOfAgenticPrompt: + return c.agenticPromptHandler.OnEnd(ctx, info, prompt.ConvCallbackOutput(output)) + case components.ComponentOfChatModel: + return c.chatModelHandler.OnEnd(ctx, info, model.ConvCallbackOutput(output)) + case components.ComponentOfAgenticModel: + return c.agenticModelHandler.OnEnd(ctx, info, model.ConvAgenticCallbackOutput(output)) + case components.ComponentOfEmbedding: + return c.embeddingHandler.OnEnd(ctx, info, embedding.ConvCallbackOutput(output)) + case components.ComponentOfIndexer: + return c.indexerHandler.OnEnd(ctx, info, indexer.ConvCallbackOutput(output)) + case components.ComponentOfRetriever: + return c.retrieverHandler.OnEnd(ctx, info, retriever.ConvCallbackOutput(output)) + case components.ComponentOfLoader: + return c.loaderHandler.OnEnd(ctx, info, document.ConvLoaderCallbackOutput(output)) + case components.ComponentOfTransformer: + return c.transformerHandler.OnEnd(ctx, info, document.ConvTransformerCallbackOutput(output)) + case components.ComponentOfTool: + return c.toolHandler.OnEnd(ctx, info, tool.ConvCallbackOutput(output)) + case compose.ComponentOfToolsNode: + return c.toolsNodeHandler.OnEnd(ctx, info, convToolsNodeCallbackOutput(output)) + case compose.ComponentOfAgenticToolsNode: + return c.agenticToolsNodeHandler.OnEnd(ctx, info, convAgenticToolsNodeCallbackOutput(output)) + case adk.ComponentOfAgent: + return c.agentHandler.OnEnd(ctx, info, adk.ConvAgentCallbackOutput(output)) + case adk.ComponentOfAgenticAgent: + return c.agenticAgentHandler.OnEnd(ctx, info, adk.ConvTypedCallbackOutput[*schema.AgenticMessage](output)) + case compose.ComponentOfGraph, + compose.ComponentOfChain, + compose.ComponentOfLambda: + return c.composeTemplates[info.Component].OnEnd(ctx, info, output) + default: + return ctx + } +} + +// OnError is the callback function for the error event of a component. +// implement the callbacks Handler interface. +func (c *handlerTemplate) OnError(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + switch info.Component { + case components.ComponentOfPrompt: + return c.promptHandler.OnError(ctx, info, err) + case components.ComponentOfAgenticPrompt: + return c.agenticPromptHandler.OnError(ctx, info, err) + case components.ComponentOfChatModel: + return c.chatModelHandler.OnError(ctx, info, err) + case components.ComponentOfAgenticModel: + return c.agenticModelHandler.OnError(ctx, info, err) + case components.ComponentOfEmbedding: + return c.embeddingHandler.OnError(ctx, info, err) + case components.ComponentOfIndexer: + return c.indexerHandler.OnError(ctx, info, err) + case components.ComponentOfRetriever: + return c.retrieverHandler.OnError(ctx, info, err) + case components.ComponentOfLoader: + return c.loaderHandler.OnError(ctx, info, err) + case components.ComponentOfTransformer: + return c.transformerHandler.OnError(ctx, info, err) + case components.ComponentOfTool: + return c.toolHandler.OnError(ctx, info, err) + case compose.ComponentOfToolsNode: + return c.toolsNodeHandler.OnError(ctx, info, err) + case compose.ComponentOfAgenticToolsNode: + return c.agenticToolsNodeHandler.OnError(ctx, info, err) + case compose.ComponentOfGraph, + compose.ComponentOfChain, + compose.ComponentOfLambda: + return c.composeTemplates[info.Component].OnError(ctx, info, err) + default: + return ctx + } +} + +// OnStartWithStreamInput is the callback function for the start event of a component with stream input. +// implement the callbacks Handler interface. +func (c *handlerTemplate) OnStartWithStreamInput(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + switch info.Component { + // currently no components.Component receive stream as input + case compose.ComponentOfGraph, + compose.ComponentOfChain, + compose.ComponentOfLambda: + return c.composeTemplates[info.Component].OnStartWithStreamInput(ctx, info, input) + default: + return ctx + } +} + +// OnEndWithStreamOutput is the callback function for the end event of a component with stream output. +// implement the callbacks Handler interface. +func (c *handlerTemplate) OnEndWithStreamOutput(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + switch info.Component { + case components.ComponentOfChatModel: + return c.chatModelHandler.OnEndWithStreamOutput(ctx, info, + schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*model.CallbackOutput, error) { + return model.ConvCallbackOutput(item), nil + })) + case components.ComponentOfAgenticModel: + return c.agenticModelHandler.OnEndWithStreamOutput(ctx, info, + schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*model.AgenticCallbackOutput, error) { + return model.ConvAgenticCallbackOutput(item), nil + })) + case components.ComponentOfTool: + return c.toolHandler.OnEndWithStreamOutput(ctx, info, + schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) (*tool.CallbackOutput, error) { + return tool.ConvCallbackOutput(item), nil + })) + case compose.ComponentOfToolsNode: + return c.toolsNodeHandler.OnEndWithStreamOutput(ctx, info, + schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) ([]*schema.Message, error) { + return convToolsNodeCallbackOutput(item), nil + })) + case compose.ComponentOfAgenticToolsNode: + return c.agenticToolsNodeHandler.OnEndWithStreamOutput(ctx, info, + schema.StreamReaderWithConvert(output, func(item callbacks.CallbackOutput) ([]*schema.AgenticMessage, error) { + return convAgenticToolsNodeCallbackOutput(item), nil + })) + case compose.ComponentOfGraph, + compose.ComponentOfChain, + compose.ComponentOfLambda: + return c.composeTemplates[info.Component].OnEndWithStreamOutput(ctx, info, output) + default: + return ctx + } +} + +// Needed checks if the callback handler is needed for the given timing. +// +//nolint:cyclop +func (c *handlerTemplate) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + if info == nil { + return false + } + + switch info.Component { + case components.ComponentOfChatModel: + if c.chatModelHandler != nil && c.chatModelHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfAgenticModel: + if c.agenticModelHandler != nil && c.agenticModelHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfEmbedding: + if c.embeddingHandler != nil && c.embeddingHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfIndexer: + if c.indexerHandler != nil && c.indexerHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfLoader: + if c.loaderHandler != nil && c.loaderHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfPrompt: + if c.promptHandler != nil && c.promptHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfAgenticPrompt: + if c.agenticPromptHandler != nil && c.agenticPromptHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfRetriever: + if c.retrieverHandler != nil && c.retrieverHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfTool: + if c.toolHandler != nil && c.toolHandler.Needed(ctx, info, timing) { + return true + } + case components.ComponentOfTransformer: + if c.transformerHandler != nil && c.transformerHandler.Needed(ctx, info, timing) { + return true + } + case compose.ComponentOfToolsNode: + if c.toolsNodeHandler != nil && c.toolsNodeHandler.Needed(ctx, info, timing) { + return true + } + case compose.ComponentOfAgenticToolsNode: + if c.agenticToolsNodeHandler != nil && c.agenticToolsNodeHandler.Needed(ctx, info, timing) { + return true + } + case adk.ComponentOfAgent: + if c.agentHandler != nil && c.agentHandler.Needed(ctx, info, timing) { + return true + } + case adk.ComponentOfAgenticAgent: + if c.agenticAgentHandler != nil && c.agenticAgentHandler.Needed(ctx, info, timing) { + return true + } + case compose.ComponentOfGraph, + compose.ComponentOfChain, + compose.ComponentOfLambda: + handler := c.composeTemplates[info.Component] + if handler != nil { + checker, ok := handler.(callbacks.TimingChecker) + if !ok || checker.Needed(ctx, info, timing) { + return true + } + } + default: + return false + } + + return false +} + +// LoaderCallbackHandler is the handler for the loader callback. +type LoaderCallbackHandler struct { + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.LoaderCallbackOutput) context.Context + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *LoaderCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// TransformerCallbackHandler is the handler for the transformer callback. +type TransformerCallbackHandler struct { + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.TransformerCallbackOutput) context.Context + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *TransformerCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// EmbeddingCallbackHandler is the handler for the embedding callback. +type EmbeddingCallbackHandler struct { + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *embedding.CallbackInput) context.Context + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *embedding.CallbackOutput) context.Context + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *EmbeddingCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// IndexerCallbackHandler is the handler for the indexer callback. +type IndexerCallbackHandler struct { + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *indexer.CallbackInput) context.Context + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *indexer.CallbackOutput) context.Context + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *IndexerCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// ModelCallbackHandler is the handler for the model callback. +type ModelCallbackHandler struct { + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context + OnEndWithStreamOutput func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *ModelCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + case callbacks.TimingOnEndWithStreamOutput: + return ch.OnEndWithStreamOutput != nil + default: + return false + } +} + +// PromptCallbackHandler is the handler for the callback. +type PromptCallbackHandler struct { + // OnStart is the callback function for the start of the callback. + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context + // OnEnd is the callback function for the end of the callback. + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context + // OnError is the callback function for the error of the callback. + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *PromptCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// RetrieverCallbackHandler is the handler for the retriever callback. +type RetrieverCallbackHandler struct { + // OnStart is the callback function for the start of the retriever. + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *retriever.CallbackInput) context.Context + // OnEnd is the callback function for the end of the retriever. + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *retriever.CallbackOutput) context.Context + // OnError is the callback function for the error of the retriever. + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *RetrieverCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// ToolCallbackHandler is the handler for the tool callback. +type ToolCallbackHandler struct { + OnStart func(ctx context.Context, info *callbacks.RunInfo, input *tool.CallbackInput) context.Context + OnEnd func(ctx context.Context, info *callbacks.RunInfo, output *tool.CallbackOutput) context.Context + OnEndWithStreamOutput func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[*tool.CallbackOutput]) context.Context + OnError func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *ToolCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnEndWithStreamOutput: + return ch.OnEndWithStreamOutput != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// ToolsNodeCallbackHandlers defines optional callbacks for the Tools node +// lifecycle events. +type ToolsNodeCallbackHandlers struct { + OnStart func(ctx context.Context, info *callbacks.RunInfo, input *schema.Message) context.Context + OnEnd func(ctx context.Context, info *callbacks.RunInfo, input []*schema.Message) context.Context + OnEndWithStreamOutput func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.Message]) context.Context + OnError func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context +} + +// Needed reports whether a handler is registered for the given timing. +func (ch *ToolsNodeCallbackHandlers) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnEndWithStreamOutput: + return ch.OnEndWithStreamOutput != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +func convToolsNodeCallbackInput(src callbacks.CallbackInput) *schema.Message { + switch t := src.(type) { + case *schema.Message: + return t + default: + return nil + } +} + +func convToolsNodeCallbackOutput(src callbacks.CallbackInput) []*schema.Message { + switch t := src.(type) { + case []*schema.Message: + return t + default: + return nil + } +} + +// AgentCallbackHandler handles callbacks for agents using *schema.Message. +// Use ComponentOfAgent to filter callback events to agent-related events. +type AgentCallbackHandler struct { + // OnStart is called when an agent run begins. Return a modified context to propagate values. + OnStart func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context + // OnEnd is called when an agent run completes. The output's Events iterator should be + // consumed asynchronously to avoid blocking. + OnEnd func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context +} + +func (ch *AgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + default: + return false + } +} + +// AgenticAgentCallbackHandler handles callbacks for agentic agents using *schema.AgenticMessage. +// Use ComponentOfAgenticAgent to filter callback events to agentic-agent-related events. +type AgenticAgentCallbackHandler struct { + // OnStart is called when an agentic agent run begins. Return a modified context to propagate values. + OnStart func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context + // OnEnd is called when an agentic agent run completes. The output's Events iterator should be + // consumed asynchronously to avoid blocking. + OnEnd func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context +} + +func (ch *AgenticAgentCallbackHandler) Needed(ctx context.Context, info *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + default: + return false + } +} + +// AgenticPromptCallbackHandler is the handler for the agentic prompt callback. +type AgenticPromptCallbackHandler struct { + // OnStart is the callback function for the start of the agentic prompt. + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context + // OnEnd is the callback function for the end of the agentic prompt. + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context + // OnError is the callback function for the error of the agentic prompt. + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *AgenticPromptCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +// AgenticModelCallbackHandler is the handler for the agentic chat model callback. +type AgenticModelCallbackHandler struct { + OnStart func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.AgenticCallbackInput) context.Context + OnEnd func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.AgenticCallbackOutput) context.Context + OnEndWithStreamOutput func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.AgenticCallbackOutput]) context.Context + OnError func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context +} + +// Needed checks if the callback handler is needed for the given timing. +func (ch *AgenticModelCallbackHandler) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnError: + return ch.OnError != nil + case callbacks.TimingOnEndWithStreamOutput: + return ch.OnEndWithStreamOutput != nil + default: + return false + } +} + +// AgenticToolsNodeCallbackHandlers defines optional callbacks for the Agentic Tools node +// lifecycle events. +type AgenticToolsNodeCallbackHandlers struct { + OnStart func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context + OnEnd func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context + OnEndWithStreamOutput func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context + OnError func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context +} + +// Needed reports whether a handler is registered for the given timing. +func (ch *AgenticToolsNodeCallbackHandlers) Needed(ctx context.Context, runInfo *callbacks.RunInfo, timing callbacks.CallbackTiming) bool { + switch timing { + case callbacks.TimingOnStart: + return ch.OnStart != nil + case callbacks.TimingOnEnd: + return ch.OnEnd != nil + case callbacks.TimingOnEndWithStreamOutput: + return ch.OnEndWithStreamOutput != nil + case callbacks.TimingOnError: + return ch.OnError != nil + default: + return false + } +} + +func convAgenticToolsNodeCallbackInput(src callbacks.CallbackInput) *schema.AgenticMessage { + switch t := src.(type) { + case *schema.AgenticMessage: + return t + default: + return nil + } +} + +func convAgenticToolsNodeCallbackOutput(src callbacks.CallbackInput) []*schema.AgenticMessage { + switch t := src.(type) { + case []*schema.AgenticMessage: + return t + default: + return nil + } +} diff --git a/utils/callbacks/template_test.go b/utils/callbacks/template_test.go new file mode 100644 index 0000000..79be157 --- /dev/null +++ b/utils/callbacks/template_test.go @@ -0,0 +1,807 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package callbacks + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/cloudwego/eino/adk" + "github.com/cloudwego/eino/callbacks" + "github.com/cloudwego/eino/components" + "github.com/cloudwego/eino/components/document" + "github.com/cloudwego/eino/components/embedding" + "github.com/cloudwego/eino/components/indexer" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/components/prompt" + "github.com/cloudwego/eino/components/retriever" + "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" +) + +func TestNewComponentTemplate(t *testing.T) { + t.Run("TestNewComponentTemplate", func(t *testing.T) { + cnt := 0 + tpl := NewHandlerHelper() + tpl.ChatModel(&ModelCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.CallbackOutput]) context.Context { + output.Close() + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }}). + Embedding(&EmbeddingCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *embedding.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *embedding.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + Prompt(&PromptCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + Retriever(&RetrieverCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *retriever.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *retriever.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + Tool(&ToolCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *tool.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *tool.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*tool.CallbackOutput]) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + Lambda(callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + cnt++ + return ctx + }). + OnStartWithStreamInputFn(func(ctx context.Context, info *callbacks.RunInfo, input *schema.StreamReader[callbacks.CallbackInput]) context.Context { + input.Close() + cnt++ + return ctx + }). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + cnt++ + return ctx + }). + OnEndWithStreamOutputFn(func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[callbacks.CallbackOutput]) context.Context { + output.Close() + cnt++ + return ctx + }). + OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }).Build()). + AgenticModel(&AgenticModelCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.AgenticCallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *model.AgenticCallbackOutput) context.Context { + cnt++ + return ctx + }, + OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[*model.AgenticCallbackOutput]) context.Context { + output.Close() + cnt++ + return ctx + }, + OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + AgenticPrompt(&AgenticPromptCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context { + cnt++ + return ctx + }, + OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context { + output.Close() + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }). + Handler() + + types := []components.Component{ + components.ComponentOfPrompt, + components.ComponentOfChatModel, + components.ComponentOfEmbedding, + components.ComponentOfRetriever, + components.ComponentOfTool, + compose.ComponentOfLambda, + components.ComponentOfAgenticModel, + components.ComponentOfAgenticPrompt, + compose.ComponentOfAgenticToolsNode, + } + + handler := tpl.Handler() + ctx := context.Background() + for _, typ := range types { + handler.OnStart(ctx, &callbacks.RunInfo{Component: typ}, nil) + handler.OnEnd(ctx, &callbacks.RunInfo{Component: typ}, nil) + handler.OnError(ctx, &callbacks.RunInfo{Component: typ}, fmt.Errorf("mock err")) + + sir, siw := schema.Pipe[callbacks.CallbackInput](1) + siw.Close() + handler.OnStartWithStreamInput(ctx, &callbacks.RunInfo{Component: typ}, sir) + + sor, sow := schema.Pipe[callbacks.CallbackOutput](1) + sow.Close() + handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: typ}, sor) + } + + assert.Equal(t, 33, cnt) + + ctx = context.Background() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler) + callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 33, cnt) + + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt}) + ctx = callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 34, cnt) + + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}) + callbacks.OnEnd[any](ctx, nil) + assert.Equal(t, 34, cnt) + + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}) + callbacks.OnError(ctx, nil) + assert.Equal(t, 35, cnt) + + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}) + callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 35, cnt) + + tpl.Transformer(&TransformerCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.TransformerCallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }).Indexer(&IndexerCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *indexer.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *indexer.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }).Loader(&LoaderCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *document.LoaderCallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }).ToolsNode(&ToolsNodeCallbackHandlers{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *schema.Message) context.Context { + cnt++ + return ctx + }, + OnEndWithStreamOutput: func(ctx context.Context, runInfo *callbacks.RunInfo, output *schema.StreamReader[[]*schema.Message]) context.Context { + cnt++ + + if output == nil { + return ctx + } + + for { + _, err := output.Recv() + if err != nil { + return ctx + } + } + }, + }).AgenticPrompt(&AgenticPromptCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, runInfo *callbacks.RunInfo, output *prompt.CallbackOutput) context.Context { + cnt++ + return ctx + }, + OnError: func(ctx context.Context, runInfo *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }).AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context { + cnt++ + return ctx + }, + OnEndWithStreamOutput: func(ctx context.Context, info *callbacks.RunInfo, output *schema.StreamReader[[]*schema.AgenticMessage]) context.Context { + output.Close() + cnt++ + return ctx + }, + OnError: func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context { + cnt++ + return ctx + }, + }) + + handler = tpl.Handler() + ctx = context.Background() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, handler) + + ctx = callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 36, cnt) + + callbacks.OnEnd[any](ctx, nil) + assert.Equal(t, 37, cnt) + + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}) + callbacks.OnEnd[any](ctx, nil) + assert.Equal(t, 38, cnt) + + ctx = callbacks.ReuseHandlers(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}) + callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 39, cnt) + + sr, sw := schema.Pipe[any](0) + sw.Close() + callbacks.OnEndWithStreamOutput[any](ctx, sr) + assert.Equal(t, 40, cnt) + + sr1, sw1 := schema.Pipe[[]*schema.Message](1) + sw1.Send([]*schema.Message{{}}, nil) + sw1.Close() + callbacks.OnEndWithStreamOutput[[]*schema.Message](ctx, sr1) + // Check AgenticModel stream + sir2, siw2 := schema.Pipe[callbacks.CallbackOutput](1) + siw2.Close() + handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, sir2) + assert.Equal(t, 42, cnt) + + // Check AgenticToolsNode stream + sir3, siw3 := schema.Pipe[callbacks.CallbackOutput](1) + siw3.Close() + handler.OnEndWithStreamOutput(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, sir3) + assert.Equal(t, 43, cnt) + + ctx = callbacks.ReuseHandlers(ctx, nil) + callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 43, cnt) + }) + + t.Run("EdgeCases", func(t *testing.T) { + ctx := context.Background() + cnt := 0 + + // 1. Test Graph and Chain Setters and Execution + tpl := NewHandlerHelper(). + Graph(callbacks.NewHandlerBuilder(). + OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context { + cnt++ + return ctx + }).Build()). + Chain(callbacks.NewHandlerBuilder(). + OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context { + cnt++ + return ctx + }).Build()) + + h := tpl.Handler() + + // Trigger Graph OnStart + h.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfGraph}, nil) + assert.Equal(t, 1, cnt) + + // Trigger Chain OnEnd + h.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfChain}, nil) + assert.Equal(t, 2, cnt) + + // 2. Test Needed logic for Graph/Chain when handler is present/absent + // Graph is present (OnStart) + needed := h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfGraph}, callbacks.TimingOnStart) + assert.True(t, needed) + + // Chain is present (OnEnd) - but we check OnStart which is not defined in the builder above? + // NewHandlerBuilder returns a handler that usually returns true for Needed if the specific func is not nil. + // Let's verify Chain OnStart is NOT needed because we only set OnEndFn. + needed = h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfChain}, callbacks.TimingOnStart) + assert.False(t, needed) // Should be false because OnStartFn wasn't set for Chain + + // Lambda is NOT present + needed = h.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfLambda}, callbacks.TimingOnStart) + assert.False(t, needed) + + // 3. Test Conversion Fallbacks (Default cases) + // We need a handler with ToolsNode and AgenticToolsNode to test their conversion fallbacks + tpl2 := NewHandlerHelper(). + ToolsNode(&ToolsNodeCallbackHandlers{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.Message) context.Context { + if input == nil { + cnt++ + } + return ctx + }, + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.Message) context.Context { + if input == nil { + cnt++ + } + return ctx + }, + }). + AgenticToolsNode(&AgenticToolsNodeCallbackHandlers{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *schema.AgenticMessage) context.Context { + if input == nil { + cnt++ + } + return ctx + }, + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, input []*schema.AgenticMessage) context.Context { + if input == nil { + cnt++ + } + return ctx + }, + }) + + h2 := tpl2.Handler() + + // Pass wrong type (string) to trigger default case in convToolsNodeCallbackInput -> returns nil + h2.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}, "wrong-input-type") + assert.Equal(t, 3, cnt) // +1 + + // Pass wrong type to trigger default case in convToolsNodeCallbackOutput -> returns nil + h2.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfToolsNode}, "wrong-output-type") + assert.Equal(t, 4, cnt) // +1 + + // Pass wrong type to trigger default case in convAgenticToolsNodeCallbackInput -> returns nil + h2.OnStart(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, "wrong-input-type") + assert.Equal(t, 5, cnt) // +1 + + // Pass wrong type to trigger default case in convAgenticToolsNodeCallbackOutput -> returns nil + h2.OnEnd(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, "wrong-output-type") + assert.Equal(t, 6, cnt) // +1 + + // 4. Test Needed for Agentic components when handlers are Set vs Unset + // tpl2 has AgenticToolsNode set + needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: compose.ComponentOfAgenticToolsNode}, callbacks.TimingOnStart) + assert.True(t, needed) + + // tpl2 does NOT have AgenticModel set + needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, callbacks.TimingOnStart) + assert.False(t, needed) + + // Set it now + tpl2.AgenticModel(&AgenticModelCallbackHandler{ + OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.AgenticCallbackInput) context.Context { + return ctx + }, + }) + + needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticModel}, callbacks.TimingOnStart) + assert.True(t, needed) + + // Check invalid component + needed = h2.(callbacks.TimingChecker).Needed(ctx, &callbacks.RunInfo{Component: "UnknownComponent"}, callbacks.TimingOnStart) + assert.False(t, needed) + + // Check RunInfo nil + needed = h2.(callbacks.TimingChecker).Needed(ctx, nil, callbacks.TimingOnStart) + assert.False(t, needed) + + // 5. Test Needed for Transformer, Loader, Indexer, etc to ensure switch coverage + tpl3 := NewHandlerHelper(). + Transformer(&TransformerCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *document.TransformerCallbackInput) context.Context { + return ctx + }}). + Loader(&LoaderCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *document.LoaderCallbackInput) context.Context { + return ctx + }}). + Indexer(&IndexerCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *indexer.CallbackInput) context.Context { + return ctx + }}). + Retriever(&RetrieverCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *retriever.CallbackInput) context.Context { + return ctx + }}). + Embedding(&EmbeddingCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *embedding.CallbackInput) context.Context { + return ctx + }}). + Tool(&ToolCallbackHandler{OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *tool.CallbackInput) context.Context { + return ctx + }}) + + h3 := tpl3.Handler() + checker := h3.(callbacks.TimingChecker) + + assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, callbacks.TimingOnStart)) + assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}, callbacks.TimingOnStart)) + assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}, callbacks.TimingOnStart)) + assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfRetriever}, callbacks.TimingOnStart)) + assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}, callbacks.TimingOnStart)) + assert.True(t, checker.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTool}, callbacks.TimingOnStart)) + + // Verify False paths (by using a helper without them) + emptyH := NewHandlerHelper().Handler().(callbacks.TimingChecker) + assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTransformer}, callbacks.TimingOnStart)) + assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfLoader}, callbacks.TimingOnStart)) + assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfIndexer}, callbacks.TimingOnStart)) + assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfRetriever}, callbacks.TimingOnStart)) + assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfEmbedding}, callbacks.TimingOnStart)) + assert.False(t, emptyH.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfTool}, callbacks.TimingOnStart)) + + // 6. Test Needed for remaining components (ChatModel, Prompt, AgenticPrompt) + tpl4 := NewHandlerHelper(). + ChatModel(&ModelCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *model.CallbackInput) context.Context { + return ctx + }}). + Prompt(&PromptCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context { + return ctx + }}). + AgenticPrompt(&AgenticPromptCallbackHandler{OnStart: func(ctx context.Context, runInfo *callbacks.RunInfo, input *prompt.CallbackInput) context.Context { + return ctx + }}) + + h4 := tpl4.Handler() + checker4 := h4.(callbacks.TimingChecker) + + assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfChatModel}, callbacks.TimingOnStart)) + assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfPrompt}, callbacks.TimingOnStart)) + assert.True(t, checker4.Needed(ctx, &callbacks.RunInfo{Component: components.ComponentOfAgenticPrompt}, callbacks.TimingOnStart)) + }) +} + +func TestAgentCallbackHandler(t *testing.T) { + t.Run("Needed returns correct values", func(t *testing.T) { + handler := &AgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context { + return ctx + }, + } + + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgent} + + assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnStart)) + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd)) + }) + + t.Run("Needed with OnEnd set", func(t *testing.T) { + handler := &AgentCallbackHandler{ + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context { + return ctx + }, + } + + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgent} + + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart)) + assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnEnd)) + }) + + t.Run("Needed with nil handlers", func(t *testing.T) { + handler := &AgentCallbackHandler{} + + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgent} + + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart)) + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd)) + }) +} + +func TestHandlerHelperWithAgent(t *testing.T) { + t.Run("Agent method sets handler correctly", func(t *testing.T) { + cnt := 0 + tpl := NewHandlerHelper() + tpl.Agent(&AgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context { + cnt++ + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: adk.ComponentOfAgent}, handler) + + ctx = callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 1, cnt) + + callbacks.OnEnd[any](ctx, nil) + assert.Equal(t, 2, cnt) + }) +} + +func TestHandlerTemplateWithAgentComponent(t *testing.T) { + t.Run("OnStart routes to agent handler", func(t *testing.T) { + called := false + tpl := NewHandlerHelper() + tpl.Agent(&AgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context { + called = true + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgent, Name: "TestAgent"} + + handler.OnStart(ctx, info, &adk.AgentCallbackInput{}) + assert.True(t, called) + }) + + t.Run("OnEnd routes to agent handler", func(t *testing.T) { + called := false + tpl := NewHandlerHelper() + tpl.Agent(&AgentCallbackHandler{ + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.AgentCallbackOutput) context.Context { + called = true + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgent, Name: "TestAgent"} + + handler.OnEnd(ctx, info, &adk.AgentCallbackOutput{}) + assert.True(t, called) + }) + + t.Run("Needed returns true for agent component", func(t *testing.T) { + tpl := NewHandlerHelper() + tpl.Agent(&AgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.AgentCallbackInput) context.Context { + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgent} + + checker, ok := handler.(callbacks.TimingChecker) + assert.True(t, ok, "handler should implement TimingChecker") + assert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart)) + }) +} + +func TestAgenticAgentCallbackHandler(t *testing.T) { + t.Run("Needed returns correct values", func(t *testing.T) { + handler := &AgenticAgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context { + return ctx + }, + } + + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent} + + assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnStart)) + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd)) + }) + + t.Run("Needed with OnEnd set", func(t *testing.T) { + handler := &AgenticAgentCallbackHandler{ + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context { + return ctx + }, + } + + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent} + + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart)) + assert.True(t, handler.Needed(ctx, info, callbacks.TimingOnEnd)) + }) + + t.Run("Needed with nil handlers", func(t *testing.T) { + handler := &AgenticAgentCallbackHandler{} + + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent} + + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnStart)) + assert.False(t, handler.Needed(ctx, info, callbacks.TimingOnEnd)) + }) +} + +func TestHandlerHelperWithAgenticAgent(t *testing.T) { + t.Run("AgenticAgent method sets handler correctly", func(t *testing.T) { + cnt := 0 + tpl := NewHandlerHelper() + tpl.AgenticAgent(&AgenticAgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context { + cnt++ + return ctx + }, + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context { + cnt++ + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + ctx = callbacks.InitCallbacks(ctx, &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent}, handler) + + ctx = callbacks.OnStart[any](ctx, nil) + assert.Equal(t, 1, cnt) + + callbacks.OnEnd[any](ctx, nil) + assert.Equal(t, 2, cnt) + }) +} + +func TestHandlerTemplateWithAgenticAgentComponent(t *testing.T) { + t.Run("OnStart routes to agentic agent handler", func(t *testing.T) { + called := false + tpl := NewHandlerHelper() + tpl.AgenticAgent(&AgenticAgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context { + called = true + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent, Name: "TestAgenticAgent"} + + handler.OnStart(ctx, info, &adk.TypedAgentCallbackInput[*schema.AgenticMessage]{}) + assert.True(t, called) + }) + + t.Run("OnEnd routes to agentic agent handler", func(t *testing.T) { + called := false + tpl := NewHandlerHelper() + tpl.AgenticAgent(&AgenticAgentCallbackHandler{ + OnEnd: func(ctx context.Context, info *callbacks.RunInfo, output *adk.TypedAgentCallbackOutput[*schema.AgenticMessage]) context.Context { + called = true + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent, Name: "TestAgenticAgent"} + + handler.OnEnd(ctx, info, &adk.TypedAgentCallbackOutput[*schema.AgenticMessage]{}) + assert.True(t, called) + }) + + t.Run("Needed returns true for agentic agent component", func(t *testing.T) { + tpl := NewHandlerHelper() + tpl.AgenticAgent(&AgenticAgentCallbackHandler{ + OnStart: func(ctx context.Context, info *callbacks.RunInfo, input *adk.TypedAgentCallbackInput[*schema.AgenticMessage]) context.Context { + return ctx + }, + }) + + handler := tpl.Handler() + ctx := context.Background() + info := &callbacks.RunInfo{Component: adk.ComponentOfAgenticAgent} + + checker, ok := handler.(callbacks.TimingChecker) + assert.True(t, ok, "handler should implement TimingChecker") + assert.True(t, checker.Needed(ctx, info, callbacks.TimingOnStart)) + }) +}