From f2c3b2d834235e085d28e8726b81fd089028f5fd Mon Sep 17 00:00:00 2001 From: Pin Hsu Chen <96507253+PaulChen79@users.noreply.github.com> Date: Fri, 8 May 2026 10:26:55 +0800 Subject: [PATCH] feat(sdk): add wren-langchain package for langchain/langgraph integration (#2247) --- .github/workflows/publish-wren-langchain.yml | 120 +++++++++ .github/workflows/release-please.yml | 17 ++ .github/workflows/sdk-langchain-ci.yml | 96 +++++++ .release-please-manifest.json | 3 +- LICENSE | 2 +- README.md | 4 +- release-please-config.json | 8 + sdk/wren-langchain/.gitignore | 14 + sdk/wren-langchain/LICENSE | 201 ++++++++++++++ sdk/wren-langchain/README.md | 247 ++++++++++++++++++ sdk/wren-langchain/examples/langchain_demo.py | 138 ++++++++++ sdk/wren-langchain/examples/langgraph_demo.py | 187 +++++++++++++ sdk/wren-langchain/pyproject.toml | 125 +++++++++ .../src/wren_langchain/__init__.py | 15 ++ .../src/wren_langchain/_envelope.py | 118 +++++++++ .../src/wren_langchain/_format.py | 149 +++++++++++ .../src/wren_langchain/_memory_api.py | 98 +++++++ .../src/wren_langchain/_prompt.py | 224 ++++++++++++++++ .../src/wren_langchain/_providers/__init__.py | 0 .../wren_langchain/_providers/connection.py | 80 ++++++ .../wren_langchain/_providers/mdl_source.py | 42 +++ .../src/wren_langchain/_providers/memory.py | 41 +++ .../src/wren_langchain/_toolkit.py | 229 ++++++++++++++++ .../src/wren_langchain/_tools.py | 120 +++++++++ .../src/wren_langchain/_tools_memory.py | 123 +++++++++ .../src/wren_langchain/exceptions.py | 18 ++ sdk/wren-langchain/tests/__init__.py | 0 .../tests/conformance/__init__.py | 0 .../conformance/test_langchain_contract.py | 122 +++++++++ sdk/wren-langchain/tests/conftest.py | 34 +++ .../tests/integration/__init__.py | 0 .../tests/integration/conftest.py | 78 ++++++ .../integration/test_langgraph_toolnode.py | 76 ++++++ .../tests/integration/test_memory_tools.py | 68 +++++ .../tests/integration/test_runtime_tools.py | 53 ++++ sdk/wren-langchain/tests/unit/__init__.py | 0 .../tests/unit/test_envelope.py | 187 +++++++++++++ .../tests/unit/test_exceptions.py | 21 ++ .../tests/unit/test_memory_api.py | 112 ++++++++ sdk/wren-langchain/tests/unit/test_prompt.py | 154 +++++++++++ .../tests/unit/test_providers_connection.py | 84 ++++++ .../tests/unit/test_providers_mdl.py | 58 ++++ .../tests/unit/test_providers_memory.py | 43 +++ .../tests/unit/test_toolkit_init.py | 104 ++++++++ .../tests/unit/test_toolkit_runtime.py | 121 +++++++++ .../tests/unit/test_tools_memory.py | 159 +++++++++++ .../tests/unit/test_tools_runtime.py | 141 ++++++++++ 47 files changed, 4030 insertions(+), 4 deletions(-) create mode 100644 .github/workflows/publish-wren-langchain.yml create mode 100644 .github/workflows/sdk-langchain-ci.yml create mode 100644 sdk/wren-langchain/.gitignore create mode 100644 sdk/wren-langchain/LICENSE create mode 100644 sdk/wren-langchain/README.md create mode 100644 sdk/wren-langchain/examples/langchain_demo.py create mode 100644 sdk/wren-langchain/examples/langgraph_demo.py create mode 100644 sdk/wren-langchain/pyproject.toml create mode 100644 sdk/wren-langchain/src/wren_langchain/__init__.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_envelope.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_format.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_memory_api.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_prompt.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_providers/__init__.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_providers/connection.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_providers/mdl_source.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_providers/memory.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_toolkit.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_tools.py create mode 100644 sdk/wren-langchain/src/wren_langchain/_tools_memory.py create mode 100644 sdk/wren-langchain/src/wren_langchain/exceptions.py create mode 100644 sdk/wren-langchain/tests/__init__.py create mode 100644 sdk/wren-langchain/tests/conformance/__init__.py create mode 100644 sdk/wren-langchain/tests/conformance/test_langchain_contract.py create mode 100644 sdk/wren-langchain/tests/conftest.py create mode 100644 sdk/wren-langchain/tests/integration/__init__.py create mode 100644 sdk/wren-langchain/tests/integration/conftest.py create mode 100644 sdk/wren-langchain/tests/integration/test_langgraph_toolnode.py create mode 100644 sdk/wren-langchain/tests/integration/test_memory_tools.py create mode 100644 sdk/wren-langchain/tests/integration/test_runtime_tools.py create mode 100644 sdk/wren-langchain/tests/unit/__init__.py create mode 100644 sdk/wren-langchain/tests/unit/test_envelope.py create mode 100644 sdk/wren-langchain/tests/unit/test_exceptions.py create mode 100644 sdk/wren-langchain/tests/unit/test_memory_api.py create mode 100644 sdk/wren-langchain/tests/unit/test_prompt.py create mode 100644 sdk/wren-langchain/tests/unit/test_providers_connection.py create mode 100644 sdk/wren-langchain/tests/unit/test_providers_mdl.py create mode 100644 sdk/wren-langchain/tests/unit/test_providers_memory.py create mode 100644 sdk/wren-langchain/tests/unit/test_toolkit_init.py create mode 100644 sdk/wren-langchain/tests/unit/test_toolkit_runtime.py create mode 100644 sdk/wren-langchain/tests/unit/test_tools_memory.py create mode 100644 sdk/wren-langchain/tests/unit/test_tools_runtime.py diff --git a/.github/workflows/publish-wren-langchain.yml b/.github/workflows/publish-wren-langchain.yml new file mode 100644 index 000000000..3784b4128 --- /dev/null +++ b/.github/workflows/publish-wren-langchain.yml @@ -0,0 +1,120 @@ +name: Publish wren-langchain to PyPI + +# Mirrors .github/workflows/publish-wren.yml. Inputs come from +# workflow_call (orchestrated by release-please or manual dispatch wrapper), +# not from untrusted issue/PR/comment events, so direct `${{ inputs.* }}` +# interpolation is safe in run/env contexts here. If you ever wire this to +# a public event source, switch to `env:` + shell variables first. + +on: + workflow_call: + inputs: + version: + description: "Version number (e.g. 0.1.0 or 0.1.0rc1)" + required: true + type: string + tag_name: + description: "Git tag to checkout (e.g. wren-langchain-v0.1.0)" + required: true + type: string + pypi_target: + description: "Publish target (pypi or testpypi)" + required: false + type: string + default: "pypi" + +permissions: + # Workflow-wide default = read-only. The publish job scopes its own + # id-token: write below; build/validate-inputs run user code and don't + # need to mint OIDC tokens. + contents: read + +jobs: + validate-inputs: + name: Validate workflow inputs + runs-on: ubuntu-latest + steps: + # Reject typos like "test-pypi" / "PYPI" before they accidentally route + # to production: only the exact strings "pypi" or "testpypi" are valid. + - name: Check pypi_target + env: + PYPI_TARGET: ${{ inputs.pypi_target }} + run: | + case "$PYPI_TARGET" in + pypi|testpypi) ;; + *) + echo "::error::Invalid pypi_target '$PYPI_TARGET'. Expected 'pypi' or 'testpypi'." + exit 1 + ;; + esac + + build: + name: Build distribution + needs: validate-inputs + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ inputs.tag_name }} + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Set version + env: + VERSION: ${{ inputs.version }} + shell: python + run: | + import re, os, sys + version = os.environ["VERSION"] + if not re.fullmatch(r"\d+\.\d+\.\d+(rc\d+)?", version): + print(f"::error::Unsupported version format: {version!r}. Expected X.Y.Z or X.Y.ZrcN.") + sys.exit(1) + for path, pattern in [ + ("sdk/wren-langchain/pyproject.toml", r'^(version\s*=\s*)".*?"'), + ("sdk/wren-langchain/src/wren_langchain/__init__.py", r'^(__version__\s*=\s*)".*?"'), + ]: + # Explicit utf-8 — pyproject.toml description and __init__.py docstring + # both contain non-ASCII; relying on platform default could corrupt them. + text = open(path, encoding="utf-8").read() + text, n = re.subn(pattern, rf'\1"{version}"', text, count=1, flags=re.MULTILINE) + if n != 1: + print(f"::error::Failed to update version in {path}") + sys.exit(1) + open(path, "w", encoding="utf-8").write(text) + - name: Install build tool + run: pip install build + - name: Build sdist and wheel + run: python -m build + working-directory: sdk/wren-langchain + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: dist + path: sdk/wren-langchain/dist/ + + publish: + name: Publish to ${{ inputs.pypi_target }} + needs: build + runs-on: ubuntu-latest + # Scoped to publish only — least-privilege per least-privilege principle. + # `id-token: write` is what authenticates to PyPI via OIDC Trusted Publishing. + permissions: + contents: read + id-token: write + environment: + name: ${{ inputs.pypi_target }} + url: ${{ inputs.pypi_target == 'pypi' && 'https://pypi.org/project/wren-langchain/' || 'https://test.pypi.org/project/wren-langchain/' }} + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + name: dist + path: dist/ + - name: List artifacts + run: ls -lhR dist/ + - name: Publish + uses: pypa/gh-action-pypi-publish@release/v1 + with: + repository-url: ${{ inputs.pypi_target == 'testpypi' && 'https://test.pypi.org/legacy/' || 'https://upload.pypi.org/legacy/' }} + packages-dir: dist/ + attestations: false diff --git a/.github/workflows/release-please.yml b/.github/workflows/release-please.yml index c3f0f284e..6e2c67d35 100644 --- a/.github/workflows/release-please.yml +++ b/.github/workflows/release-please.yml @@ -22,6 +22,9 @@ jobs: wren-core-wasm--release_created: ${{ steps.release.outputs['core/wren-core-wasm--release_created'] }} wren-core-wasm--tag_name: ${{ steps.release.outputs['core/wren-core-wasm--tag_name'] }} wren-core-wasm--version: ${{ steps.release.outputs['core/wren-core-wasm--version'] }} + wren-langchain--release_created: ${{ steps.release.outputs['sdk/wren-langchain--release_created'] }} + wren-langchain--tag_name: ${{ steps.release.outputs['sdk/wren-langchain--tag_name'] }} + wren-langchain--version: ${{ steps.release.outputs['sdk/wren-langchain--version'] }} steps: - uses: googleapis/release-please-action@v4 id: release @@ -61,3 +64,17 @@ jobs: permissions: contents: read id-token: write + + # release-please outputs come from conventional-commit aggregation, not from + # untrusted issue/PR event payloads, so direct ${{ ... }} interpolation is + # safe here. Same pattern as the publish jobs above. + publish-wren-langchain: + needs: release-please + if: needs.release-please.outputs['wren-langchain--release_created'] == 'true' + uses: ./.github/workflows/publish-wren-langchain.yml + with: + version: ${{ needs.release-please.outputs['wren-langchain--version'] }} + tag_name: ${{ needs.release-please.outputs['wren-langchain--tag_name'] }} + permissions: + contents: read + id-token: write diff --git a/.github/workflows/sdk-langchain-ci.yml b/.github/workflows/sdk-langchain-ci.yml new file mode 100644 index 000000000..a2e8e4832 --- /dev/null +++ b/.github/workflows/sdk-langchain-ci.yml @@ -0,0 +1,96 @@ +name: wren-langchain CI + +# Trigger sources are controlled events (pull_request, push to main, +# workflow_dispatch). No untrusted issue / comment / external payload is +# interpolated into run/env contexts in this file. ${{ matrix.* }} and +# ${{ github.* }} usages below are confined to YAML-level fields +# (name, with, group), not run scripts. + +permissions: + # Lint / test / build only — no PR comments, statuses, or labels written. + contents: read + +on: + pull_request: + paths: + - 'sdk/wren-langchain/**' + - '.github/workflows/sdk-langchain-ci.yml' + push: + branches: + - main + paths: + - 'sdk/wren-langchain/**' + - '.github/workflows/sdk-langchain-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: true + +defaults: + run: + working-directory: sdk/wren-langchain + +jobs: + lint: + name: lint + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + # Ubuntu's system Python is PEP 668 externally-managed; use a hosted + # toolcache Python so `uv pip install --system` is allowed to write. + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - uses: astral-sh/setup-uv@v4 + - name: Install package + dev deps + run: uv pip install --system -e ".[dev]" + - name: ruff check + run: ruff check . + - name: ruff format check + run: ruff format --check . + + test: + name: tests (py${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.11', '3.12'] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - uses: astral-sh/setup-uv@v4 + - name: Install package + dev deps + run: uv pip install --system -e ".[dev]" + - name: Run default test suite + # pytest config skips tests marked `slow` (LanceDB + sentence-transformer + # heavyweights). Slow tests are opt-in locally; we add a dedicated CI + # job for them later if worth the runner cost. + run: pytest -v + + build: + name: build sdist + wheel + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - uses: astral-sh/setup-uv@v4 + - name: Install build tool + run: uv pip install --system build + - name: Build distributions + run: python -m build + - name: Verify LICENSE bundled in wheel + run: | + set -euo pipefail + unzip -l dist/wren_langchain-*.whl | grep -q "dist-info/licenses/LICENSE" + echo "LICENSE present in wheel" + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: wren-langchain-dist + path: sdk/wren-langchain/dist/ diff --git a/.release-please-manifest.json b/.release-please-manifest.json index c13247215..ac0f1bf93 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,5 +1,6 @@ { "core/wren-core-py": "0.5.0", "core/wren": "0.5.0", - "core/wren-core-wasm": "0.3.0" + "core/wren-core-wasm": "0.3.0", + "sdk/wren-langchain": "0.1.0" } diff --git a/LICENSE b/LICENSE index d849ffe9a..675a2e6a5 100644 --- a/LICENSE +++ b/LICENSE @@ -12,8 +12,8 @@ Path → License map | Path | License | | ----------------------------- | --------------------------------------------- | | `core/**` | Apache License 2.0 (see LICENSE-APACHE-2.0) | + | `sdk/**` | Apache License 2.0 | | `skills/**` | Apache License 2.0 | - | `sdks/integrations/**` | Apache License 2.0 | | `examples/**` | Apache License 2.0 | | `docs/**` | CC BY 4.0 International (see LICENSE-CC-BY-4.0) | | Everything else (root files) | Apache License 2.0 | diff --git a/README.md b/README.md index 4e4e1a57f..c5a6915a6 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,7 @@ See the connector API docs in the project documentation for the latest connectio |   [`core/wren/`](./core/wren) | Python SDK + `wren` CLI (PyPI: `wren-engine`). | |   [`core/wren-mdl/`](./core/wren-mdl) | MDL JSON schema. | | [`skills/`](./skills) | CLI-based agent skills (`wren-generate-mdl`, `wren-usage`, `wren-dlt-connector`, `wren-onboarding`). | -| [`sdks/integrations/`](./sdks) | Framework integrations (LangChain, CrewAI, Pydantic-AI, Goose, LlamaIndex, Mastra) — _coming soon_. | +| [`sdk/`](./sdk) | Framework integrations. [`sdk/wren-langchain/`](./sdk/wren-langchain) (PyPI: `wren-langchain`) is shipped; CrewAI / Pydantic-AI / Goose / LlamaIndex / Mastra are _coming soon_. | | [`examples/`](./examples) | End-to-end example projects — _coming soon_. | | [`docs/core/`](./docs/core) | Module documentation. | @@ -120,7 +120,7 @@ See the connector API docs in the project documentation for the latest connectio WrenAI is multi-licensed: -- **`core/**`, `skills/**`, `sdks/integrations/**`, `examples/**`, root-level files** — [Apache License 2.0](LICENSE-APACHE-2.0) +- **`core/**`, `sdk/**`, `skills/**`, `examples/**`, root-level files** — [Apache License 2.0](LICENSE-APACHE-2.0) - **`docs/**`** — [Creative Commons Attribution 4.0 International (CC BY 4.0)](LICENSE-CC-BY-4.0) Future modules may be introduced under [GNU Affero General Public License v3.0](LICENSE-AGPL-3.0); the full text is committed here pre-emptively. See [LICENSE](LICENSE) for the authoritative path-to-license map. diff --git a/release-please-config.json b/release-please-config.json index cb002c0b9..0d8a7605b 100644 --- a/release-please-config.json +++ b/release-please-config.json @@ -26,6 +26,14 @@ "extra-files": [ "Cargo.toml" ] + }, + "sdk/wren-langchain": { + "component": "wren-langchain", + "release-type": "python", + "bump-minor-pre-major": true, + "extra-files": [ + "src/wren_langchain/__init__.py" + ] } } } diff --git a/sdk/wren-langchain/.gitignore b/sdk/wren-langchain/.gitignore new file mode 100644 index 000000000..e0a7c52a5 --- /dev/null +++ b/sdk/wren-langchain/.gitignore @@ -0,0 +1,14 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +.coverage +dist/ +build/ +.venv/ + +# Local scratch scripts users create at the package root for ad-hoc testing +# (e.g. agent.py, agent_with_memory_check.py from the local-testing-guide). +# Anything in tests/ and src/wren_langchain/ stays tracked. +/agent*.py diff --git a/sdk/wren-langchain/LICENSE b/sdk/wren-langchain/LICENSE new file mode 100644 index 000000000..a2d260125 --- /dev/null +++ b/sdk/wren-langchain/LICENSE @@ -0,0 +1,201 @@ + 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 + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2024 Canner, Inc. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/sdk/wren-langchain/README.md b/sdk/wren-langchain/README.md new file mode 100644 index 000000000..9f4b317a6 --- /dev/null +++ b/sdk/wren-langchain/README.md @@ -0,0 +1,247 @@ +# wren-langchain + +LangChain and LangGraph integration for [Wren AI Core](https://github.com/Canner/WrenAI). + +Attach a CLI-prepared Wren project to a LangChain agent in three lines: + +```python +from wren_langchain import WrenToolkit +from langchain.agents import create_agent + +toolkit = WrenToolkit.from_project("./analytics_db") +agent = create_agent( + model="openai:gpt-4o", + tools=toolkit.get_tools(), + system_prompt=toolkit.system_prompt(), +) +``` + +Complete runnable demos: + +- [`examples/langchain_demo.py`](./examples/langchain_demo.py) — uses + ``langchain.agents.create_agent``, the high-level factory. Smallest + amount of code; recommended starting point. +- [`examples/langgraph_demo.py`](./examples/langgraph_demo.py) — builds the + ReAct loop from LangGraph primitives (`StateGraph` + `ToolNode` + + conditional edges). Use this when you need custom routing, state, or + streaming. + +## Prerequisites + +This package assumes you have already used the Wren CLI to prepare a project: + +```bash +wren context init +wren context build +wren memory index # optional but recommended +wren profile add ... +``` + +If you haven't installed the CLI yet, install `wren-engine` first: + +```bash +pip install "wren-engine[memory,postgres]" +``` + +## Installation + +`wren-langchain` exposes datasource and memory extras that pass through to +the matching `wren-engine` extras, so you only have to install once: + +```bash +# Match the datasource your wren_project.yml uses (DuckDB needs no extra): +pip install "wren-langchain[mysql]" +pip install "wren-langchain[postgres,memory]" +pip install "wren-langchain[bigquery,memory]" + +# Available datasource extras: postgres, mysql, bigquery, snowflake, +# clickhouse, trino, mssql, databricks, redshift, spark, athena, oracle. + +# `memory` extra enables the three memory tools (wren_fetch_context, +# wren_recall_queries, wren_store_query). Without it the toolkit exposes +# only the three runtime tools. + +# Install everything for experimentation: +pip install "wren-langchain[all,memory]" +``` + +If you prefer to install `wren-engine` separately (e.g. you already use the +CLI), the bare package is enough and your existing `wren-engine` extras carry +over: + +```bash +pip install wren-langchain +``` + +## What you get + +`WrenToolkit.from_project(path)` exposes: + +- **6 LLM-facing tools** (3 runtime + 3 memory when `.wren/memory/` exists): + - `wren_query` — execute SQL through Wren's context layer, returns rows + - `wren_dry_plan` — plan SQL without execution to verify it targets models correctly + - `wren_list_models` — list project models with column counts and descriptions + - `wren_fetch_context` — retrieve relevant schema/business context for a question + - `wren_recall_queries` — surface similar past NL→SQL pairs as few-shot examples + - `wren_store_query` — persist a confirmed NL→SQL pair for future recall +- **Direct Python API**: + ```python + toolkit.query("SELECT ...") # → pyarrow.Table + toolkit.dry_plan("SELECT ...") # → str (target-dialect SQL) + toolkit.dry_run("SELECT ...") # → None (validation only) + toolkit.memory.fetch("revenue trends") + toolkit.memory.recall("top customers") + toolkit.memory.store(nl="...", sql="...", tags=["..."]) + ``` +- **`toolkit.system_prompt()`** — Wren-aware system prompt that adapts to enabled tools and includes your project's `instructions.md` when present. + +## Configuration + +```python +WrenToolkit.from_project( + path, # required — path to your prepared Wren project + profile="prod", # optional — picks a named profile (default: active) +) + +toolkit.get_tools( + include_memory_write=True, # set False to keep memory read-only + raise_on_error=False, # set True to surface exceptions to LangChain retry +) +``` + +Memory is **auto-detected** from the project: present `/.wren/memory/` +exposes the 3 memory tools alongside the 3 runtime tools; absent → only the +runtime tools. To enable, run `wren memory index` from the project root; to +disable, delete the directory. There is no override kwarg. + +`include_memory_write=False` removes `wren_store_query` from the returned +list while keeping `wren_fetch_context` and `wren_recall_queries`. Use this +when you want the agent to read curated past pairs but never persist new +ones (e.g., a shared / pinned memory). When memory is disabled, this flag +has no effect — no memory tools are returned regardless. + +### How `path`, `profile`, and `.env` interact + +Three pieces of state combine to produce a connection. Understanding which +one drives what avoids surprises: + +| Source | Holds | Resolved by | +|---|---|---| +| `path/wren_project.yml` + `target/mdl.json` | MDL models, schema, `data_source` | `from_project(path)` | +| `path/.env` | Secret values (`MYSQL_HOST`, `MYSQL_PASSWORD`, …) | `from_project(path)` auto-loads it | +| `~/.wren/profiles.yml` | Connection template (`host: ${MYSQL_HOST}`, …) | `profile=` kwarg or fallback chain | + +**`profile=` resolution chain** (highest priority first): + +1. Explicit `profile=""` kwarg passed to `from_project`. +2. The `profile:` field inside the project's `wren_project.yml`, e.g.: + ```yaml + schema_version: 3 + data_source: mysql + profile: test-project3 # locks this project to a specific profile + ``` +3. The globally active profile (`wren profile switch `). + +### When does `profile=` actually change which database you connect to? + +This is subtle, because **profile values are templates that resolve from the +project's `.env`**, not standalone connection records. Three scenarios: + +**Scenario A — `profile=` is a no-op (most common)** + +Your `~/.wren/profiles.yml` has multiple profiles that all use the same +placeholder names: + +```yaml +profiles: + test-project3: + datasource: mysql + host: ${MYSQL_HOST} + database: ${MYSQL_DATABASE} + test-project4: + datasource: mysql + host: ${MYSQL_HOST} # ← same placeholder + database: ${MYSQL_DATABASE} +``` + +Because `from_project("/path/to/test-project3")` loads `test-project3/.env` +into the environment, **both profiles resolve to the same connection** — +`${MYSQL_HOST}` reads from project3's `.env` regardless of which profile +name you picked. Profile selection is cosmetic in this layout. + +**Scenario B — `profile=` selects different placeholders** + +```yaml +profiles: + dev: + host: ${DEV_HOST} + prod: + host: ${PROD_HOST} +``` + +Now `profile="dev"` and `profile="prod"` read different env vars from the +same `.env`, so the choice matters. + +**Scenario C — `profile=` selects hardcoded values or different datasources** + +```yaml +profiles: + local: + datasource: duckdb + url: /tmp/local.duckdb + format: duckdb + remote: + datasource: postgres + host: prod-db.example.com # hardcoded + port: 5432 +``` + +`profile="local"` vs `profile="remote"` connect to genuinely different +databases. Note: if `wren_project.yml` specifies `data_source:` and you +pick a profile with a different `datasource:`, the connection will fail — +the project's MDL is built against one specific dialect. + +### Recommendation: one project, one profile + +If you follow the common pattern of **one Wren project per database** +(each project gets its own `.env` and points at its own DB), set the +profile inside `wren_project.yml` and stop passing `profile=`: + +```yaml +# wren_project.yml +schema_version: 3 +name: test-project3 +data_source: mysql +profile: test-project3 +``` + +```python +toolkit = WrenToolkit.from_project("/path/to/test-project3") +``` + +This pins the project to its intended profile, no more "is the active +profile what I think it is?" — and it survives `wren profile switch` +elsewhere on the same machine. + +## Compatibility matrix + +| `wren-langchain` | `wren-engine` | `langchain` | `langgraph` | +|---|---|---|---| +| 0.1.0 | >= 0.5.0 | >= 1.0 | >= 1.0 | + +## Known limitations (v0.1) + +- **Synchronous tools only.** LangChain auto-bridges to a thread pool when tools run in async LangGraph; multi-tenant servers serving > ~32 concurrent users may exhaust the default executor pool. +- **One toolkit per agent.** If you need to query multiple Wren projects, build separate agents. +- **Memory is auto-detected** from `.wren/memory/` and there is no kwarg to override. To enable, run `wren memory index`; to disable, delete the directory. +- **No hot reload mechanism.** `target/mdl.json` is re-read on every tool call, so `wren context build` updates from CLI are picked up automatically. Profile changes require constructing a new toolkit. +- **Don't run `wren memory index` while an agent is using the same project.** The index operation drops and recreates the LanceDB schema table; concurrent reads may transiently fail. + +## License + +Apache License 2.0. See [LICENSE](./LICENSE) for the full text, or the +[repository-level LICENSE](../../LICENSE) for the path-to-license map. + +The names "Wren", "WrenAI", and the project's logos are trademarks of +Canner, Inc. and are not licensed under Apache 2.0; their use is governed +separately. diff --git a/sdk/wren-langchain/examples/langchain_demo.py b/sdk/wren-langchain/examples/langchain_demo.py new file mode 100644 index 000000000..992a5a30f --- /dev/null +++ b/sdk/wren-langchain/examples/langchain_demo.py @@ -0,0 +1,138 @@ +"""End-to-end LangChain agent demo using wren-langchain. + +Shows the minimum viable flow: + 1. Build a toolkit from a CLI-prepared Wren project. + 2. Pass its tools and system prompt straight into a LangChain agent. + 3. Ask a data question and let the agent decide which Wren tools to call. + +The agent picks up the full Wren workflow from ``toolkit.system_prompt()``: +recall past pairs → fetch context → write SQL → dry_plan if complex → +execute → store the NL/SQL pair. Memory tools are auto-enabled when +``.wren/memory/`` exists; otherwise the agent runs with the 3 runtime +tools only. + +Prerequisites +============= + - A CLI-prepared Wren project. Either follow the README quickstart, or + see ``temp-docs/v0.1-langchain-langgraph-sdk-local-testing-guide.md`` §3 + for a one-shot DuckDB-backed demo project. + - ``langchain-openai`` installed in the active venv: + uv pip install langchain-openai + - ``OPENAI_API_KEY`` set in the environment. + +Usage +===== + export OPENAI_API_KEY=sk-... + export PROJECT_PATH=/path/to/your-wren-project + python examples/langchain_demo.py + + # Custom question: + QUESTION="What's the gender distribution of users?" \\ + python examples/langchain_demo.py + +A note on the agent factory +=========================== +This demo uses ``langchain.agents.create_agent`` (the langgraph 1.0+ +recommended entrypoint). The older ``langgraph.prebuilt.create_react_agent`` +still works but is scheduled for removal in langgraph 2.0. +""" + +from __future__ import annotations + +import os +import sys +from collections import Counter + +try: + from langchain_openai import ChatOpenAI +except ImportError: + sys.exit( + "langchain-openai is not installed.\n" + "Run: uv pip install langchain-openai\n" + "(or substitute any other LangChain-compatible chat model below)." + ) + +try: + from langchain.agents import create_agent +except ImportError: + sys.exit( + "langchain (>= 1.0) is not installed in this venv.\n" + "It is declared as a dependency of wren-langchain, but if you used an\n" + "editable install before that pin was added, you need to re-sync deps:\n" + ' uv pip install -e ".[dev]"' + ) + +from wren_langchain import WrenToolkit + + +def main() -> None: + project_path = os.environ.get("PROJECT_PATH") + if not project_path: + sys.exit( + "PROJECT_PATH is required. Example:\n" + " PROJECT_PATH=/Users/you/my-wren-project python examples/langchain_demo.py" + ) + if not os.environ.get("OPENAI_API_KEY"): + sys.exit("OPENAI_API_KEY is required.") + + question = os.environ.get( + "QUESTION", + "List the models available in this project and summarize what each one tracks.", + ) + + # 1) Build the toolkit. ``from_project`` validates prerequisites eagerly: + # wren_project.yml + target/mdl.json must exist, profile must resolve, + # project-local .env is loaded automatically. Memory is auto-detected + # from .wren/memory/. + toolkit = WrenToolkit.from_project(project_path) + tools = toolkit.get_tools() + prompt = toolkit.system_prompt() + + print(f"Project: {project_path}") + print(f"Memory enabled: {toolkit._memory.enabled}") + print(f"Tools exposed: {[t.name for t in tools]}") + print(f"Question: {question}") + print() + + # 2) Build the agent. Any LangChain-compatible chat model works here. + agent = create_agent( + model=ChatOpenAI(model="gpt-4o", temperature=0), + tools=tools, + system_prompt=prompt, + ) + + # 3) Run and print the conversation. The agent decides which Wren tools + # to call based on the system prompt's workflow rules. + response = agent.invoke({"messages": [{"role": "user", "content": question}]}) + + bar = "=" * 64 + print(bar) + print("Conversation") + print(bar) + + tool_count: Counter[str] = Counter() + for msg in response["messages"]: + kind = type(msg).__name__ + content = (getattr(msg, "content", None) or "").strip() + tool_calls = getattr(msg, "tool_calls", None) or [] + + print(f"--- {kind} ---") + if content: + print(content) + for tc in tool_calls: + tool_count[tc["name"]] += 1 + print(f" -> {tc['name']}({tc['args']})") + print() + + print(bar) + print("Tool call summary") + print(bar) + if tool_count: + for name in sorted(tool_count): + print(f" {name:25s} called {tool_count[name]}x") + else: + print(" (no tool calls — the agent answered from prior knowledge only)") + + +if __name__ == "__main__": + main() diff --git a/sdk/wren-langchain/examples/langgraph_demo.py b/sdk/wren-langchain/examples/langgraph_demo.py new file mode 100644 index 000000000..4e7b34e1b --- /dev/null +++ b/sdk/wren-langchain/examples/langgraph_demo.py @@ -0,0 +1,187 @@ +"""LangGraph demo: build a Wren-aware ReAct agent from primitives. + +The companion ``langchain_demo.py`` calls ``langchain.agents.create_agent``, +which is the high-level factory that hides the agent loop. This demo +hand-builds the same loop with LangGraph primitives so you can customize: + + - Routing — decide between tool execution and finishing per turn. + - State — add fields beyond messages (e.g. ``user_id``, ``session_meta``). + - Streaming — yield intermediate node outputs to a UI as the agent works. + - Per-turn middleware — logging, telemetry, retry, approval gates. + +The graph topology is the standard ReAct pattern:: + + ┌──────────┐ + │ START │ + └────┬─────┘ + ▼ + ┌──────────┐ ┌────────────┐ + │ agent │ ──────► │ has tool │ ── no ──► END + │ (model) │ │ calls? │ + └──────────┘ └─────┬──────┘ + ▲ │ yes + │ ▼ + │ ┌──────────────┐ + └──────────────│ tools │ + │ (ToolNode) │ + └──────────────┘ + +Prereqs match ``langchain_demo.py``: a CLI-prepared Wren project, OPENAI_API_KEY, +and ``langchain-openai`` installed. + +Usage +===== + export OPENAI_API_KEY=sk-... + export PROJECT_PATH=/path/to/your-wren-project + python examples/langgraph_demo.py + + # Custom question + streaming view: + QUESTION="..." STREAM=1 python examples/langgraph_demo.py +""" + +from __future__ import annotations + +import os +import sys +from collections import Counter +from typing import Annotated, TypedDict + +try: + from langchain_openai import ChatOpenAI +except ImportError: + sys.exit("langchain-openai is not installed.\nRun: uv pip install langchain-openai") + +from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import add_messages +from langgraph.prebuilt import ToolNode + +from wren_langchain import WrenToolkit + + +class AgentState(TypedDict): + """Conversation state. ``add_messages`` appends new messages instead of + replacing the list — this is what makes the ReAct loop accumulate context. + + Add your own keys here (e.g. ``user_id: str``, ``trace: list[dict]``, + ``approved: bool``) when you need state beyond the conversation history. + """ + + messages: Annotated[list[BaseMessage], add_messages] + + +def build_app(toolkit: WrenToolkit, model_name: str = "gpt-4o"): + """Compile a ReAct graph that uses Wren tools.""" + tools = toolkit.get_tools() + system_prompt = toolkit.system_prompt() + model_with_tools = ChatOpenAI(model=model_name, temperature=0).bind_tools(tools) + + def agent_node(state: AgentState) -> dict: + """Call the model. Inject the Wren system prompt only on the first turn.""" + messages = state["messages"] + if not messages or not isinstance(messages[0], SystemMessage): + messages = [SystemMessage(content=system_prompt), *messages] + response = model_with_tools.invoke(messages) + return {"messages": [response]} + + def should_continue(state: AgentState) -> str: + """If the last AIMessage requested tool calls, run them; otherwise finish.""" + last = state["messages"][-1] + if isinstance(last, AIMessage) and last.tool_calls: + return "tools" + return END + + graph = StateGraph(AgentState) + graph.add_node("agent", agent_node) + graph.add_node("tools", ToolNode(tools)) + graph.add_edge(START, "agent") + graph.add_conditional_edges("agent", should_continue, {"tools": "tools", END: END}) + graph.add_edge("tools", "agent") # loop back after tool execution + return graph.compile() + + +def _print_message(msg: BaseMessage) -> Counter[str]: + """Pretty-print a single message; return a counter of any tool calls it made.""" + counts: Counter[str] = Counter() + kind = type(msg).__name__ + content = (getattr(msg, "content", None) or "").strip() + print(f"--- {kind} ---") + if content: + print(content) + for tc in getattr(msg, "tool_calls", None) or []: + counts[tc["name"]] += 1 + print(f" -> {tc['name']}({tc['args']})") + if isinstance(msg, ToolMessage): + # ToolMessage has its own name / tool_call_id worth surfacing. + print(f" (tool: {msg.name}, id: {msg.tool_call_id})") + print() + return counts + + +def main() -> None: + project_path = os.environ.get("PROJECT_PATH") + if not project_path: + sys.exit( + "PROJECT_PATH is required. Example:\n" + " PROJECT_PATH=/Users/you/my-wren-project python examples/langgraph_demo.py" + ) + if not os.environ.get("OPENAI_API_KEY"): + sys.exit("OPENAI_API_KEY is required.") + + question = os.environ.get( + "QUESTION", + "List the models in this project and pick one to summarize.", + ) + stream_mode = os.environ.get("STREAM") not in (None, "", "0") + + toolkit = WrenToolkit.from_project(project_path) + print(f"Project: {project_path}") + print(f"Memory enabled: {toolkit._memory.enabled}") + print(f"Tools exposed: {[t.name for t in toolkit.get_tools()]}") + print(f"Question: {question}") + print(f"Stream mode: {stream_mode}") + print() + + app = build_app(toolkit) + initial_state: AgentState = {"messages": [HumanMessage(content=question)]} + + bar = "=" * 64 + print(bar) + print("Conversation") + print(bar) + + tool_count: Counter[str] = Counter() + + if stream_mode: + # ``stream(..., stream_mode="updates")`` yields one dict per node + # invocation, where the dict is ``{node_name: state_update}``. This + # is the natural place to push events into a UI / log / telemetry. + for event in app.stream(initial_state, stream_mode="updates"): + for node_name, update in event.items(): + print(f"[node: {node_name}]") + for msg in update.get("messages", []): + tool_count.update(_print_message(msg)) + else: + # ``invoke`` returns the final state after the loop terminates. + final_state = app.invoke(initial_state) + for msg in final_state["messages"]: + tool_count.update(_print_message(msg)) + + print(bar) + print("Tool call summary") + print(bar) + if tool_count: + for name in sorted(tool_count): + print(f" {name:25s} called {tool_count[name]}x") + else: + print(" (no tool calls — model answered directly)") + + +if __name__ == "__main__": + main() diff --git a/sdk/wren-langchain/pyproject.toml b/sdk/wren-langchain/pyproject.toml new file mode 100644 index 000000000..06380ec6f --- /dev/null +++ b/sdk/wren-langchain/pyproject.toml @@ -0,0 +1,125 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "wren-langchain" +version = "0.1.0" +description = "LangChain and LangGraph integration for Wren AI Core — attach a CLI-prepared Wren project to your agent in three lines." +readme = "README.md" +requires-python = ">=3.11" +license = { text = "Apache-2.0" } +authors = [{ name = "Wren AI", email = "contact@getwren.ai" }] +keywords = [ + "wrenai", "wren", "langchain", "langgraph", "agent", "ai-agent", + "sql", "context-layer", "semantic", "mdl", "data-modeling", "analytics", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Database", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", + "Typing :: Typed", +] +dependencies = [ + # Core: provides WrenEngine, MemoryStore, profile resolution, MDL handling. + "wren-engine>=0.5.0", + # langchain >=1.0 ships the recommended `langchain.agents.create_agent` + # entrypoint. Older `langgraph.prebuilt.create_react_agent` still works + # but emits deprecation warnings and is scheduled for removal in + # langgraph 2.0. + "langchain>=1.0", + "langchain-core>=0.3", + "langgraph>=1.0", + "pydantic>=2", +] + +[project.optional-dependencies] +# ── Datasource pass-through extras ──────────────────────────────────────── +# Pick the one matching your project's `data_source` (in wren_project.yml). +# These chain through to the matching `wren-engine[]` extra so +# you don't have to install the connector dependency twice. +# +# pip install "wren-langchain[mysql]" +# pip install "wren-langchain[mysql,memory]" +# +# DuckDB is built into wren-engine — no extra needed. +postgres = ["wren-engine[postgres]>=0.5.0"] +mysql = ["wren-engine[mysql]>=0.5.0"] +bigquery = ["wren-engine[bigquery]>=0.5.0"] +snowflake = ["wren-engine[snowflake]>=0.5.0"] +clickhouse = ["wren-engine[clickhouse]>=0.5.0"] +trino = ["wren-engine[trino]>=0.5.0"] +mssql = ["wren-engine[mssql]>=0.5.0"] +databricks = ["wren-engine[databricks]>=0.5.0"] +redshift = ["wren-engine[redshift]>=0.5.0"] +spark = ["wren-engine[spark]>=0.5.0"] +athena = ["wren-engine[athena]>=0.5.0"] +oracle = ["wren-engine[oracle]>=0.5.0"] + +# ── Memory extra ────────────────────────────────────────────────────────── +# Required to use the three memory tools (wren_fetch_context, +# wren_recall_queries, wren_store_query). Pulls in LanceDB and the +# sentence-transformer model deps via wren-engine. +# +# Without this, the toolkit auto-detects no memory and exposes only the +# three runtime tools. +memory = ["wren-engine[memory]>=0.5.0"] + +# ── Convenience aggregate ───────────────────────────────────────────────── +# Installs every datasource + memory at once. Useful for CI / examples, +# heavyweight for production where you usually want only one datasource. +all = ["wren-engine[all]>=0.5.0"] + +# ── Dev / test deps ─────────────────────────────────────────────────────── +dev = [ + "pytest>=8", + "pytest-mock>=3", + "ruff>=0.4", + # Tests run real DuckDB + LanceDB integrations, so memory deps are required. + "wren-engine[memory]>=0.5.0", +] + +[project.urls] +Homepage = "https://getwren.ai" +Documentation = "https://github.com/Canner/WrenAI/tree/main/sdk/wren-langchain#readme" +Repository = "https://github.com/Canner/WrenAI" +Issues = "https://github.com/Canner/WrenAI/issues" +Changelog = "https://github.com/Canner/WrenAI/blob/main/sdk/wren-langchain/CHANGELOG.md" + +[tool.hatch.build.targets.wheel] +packages = ["src/wren_langchain"] + +[tool.ruff] +line-length = 88 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "I", "PLC"] +ignore = ["E501"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +# Default run skips tests marked `slow` (currently the LanceDB integration +# tests which load a sentence-transformer model, ~30-40s). Run them explicitly +# with `pytest -m slow` or `pytest -m "slow or not slow"` for everything. +addopts = "-ra -m \"not slow\"" +markers = [ + "slow: marks tests as slow (deselect with '-m \"not slow\"')", +] +# Suppress deprecation warnings emitted by langgraph's own internal imports. +# These are not from our code — they fire when `from langgraph.prebuilt import +# ToolNode` triggers langgraph's package init. We cannot fix them upstream-side +# and do not use the deprecated APIs ourselves. Review quarterly per +# sdk/MAINTENANCE.md §3. +filterwarnings = [ + "ignore::langchain_core._api.deprecation.LangChainPendingDeprecationWarning", + "ignore:AgentStatePydantic has been moved", +] diff --git a/sdk/wren-langchain/src/wren_langchain/__init__.py b/sdk/wren-langchain/src/wren_langchain/__init__.py new file mode 100644 index 000000000..9d9ba7835 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/__init__.py @@ -0,0 +1,15 @@ +"""LangChain and LangGraph integration for Wren AI Core.""" + +from wren_langchain._toolkit import WrenToolkit +from wren_langchain.exceptions import ( + MemoryNotEnabledError, + WrenToolkitInitError, +) + +__version__ = "0.1.0" + +__all__ = [ + "WrenToolkit", + "WrenToolkitInitError", + "MemoryNotEnabledError", +] diff --git a/sdk/wren-langchain/src/wren_langchain/_envelope.py b/sdk/wren-langchain/src/wren_langchain/_envelope.py new file mode 100644 index 000000000..f464b4c0a --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_envelope.py @@ -0,0 +1,118 @@ +"""Envelope construction and error formatting for LLM-facing tools. + +The envelope shape is the contract between SDK tools and LangChain agents. +Success: {"ok": True, "content": str, "data": dict, "warnings": list[str]} +Error: {"ok": False, "content": str, "error": {"code", "phase", "message", "metadata"}} +""" + +import datetime as _dt +import json +from decimal import Decimal +from typing import Any + +from wren.model.error import WrenError + +_SECRET_PATTERNS = ("password", "secret", "token", "credential") +_DEFAULT_METADATA_CAP = 4 * 1024 +_DEFAULT_CONTENT_CAP = 16 * 1024 + + +def json_safe(value: Any) -> Any: + """Recursively convert non-JSON-serializable values to JSON-friendly forms. + + - datetime/date/time → ISO 8601 string + - Decimal → string (preserves precision) + - dict / list / tuple → recurse + - other types fall back to str(value) when not natively JSON-encodable + """ + if isinstance(value, dict): + return {k: json_safe(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [json_safe(v) for v in value] + if isinstance(value, (_dt.datetime, _dt.date, _dt.time)): + return value.isoformat() + if isinstance(value, Decimal): + return str(value) + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) + + +def cap_size(data: dict[str, Any], max_bytes: int) -> dict[str, Any]: + """Return *data* unchanged if JSON-serialized size <= max_bytes. + + When over-limit, return a sentinel marker dict instead. Callers should + treat the marker as opaque and surface to logs/middleware, not LLMs. + """ + encoded = json.dumps(data, default=str).encode("utf-8") + if len(encoded) <= max_bytes: + return data + return { + "_truncated": True, + "original_size_bytes": len(encoded), + } + + +def redact_secrets(data: dict[str, Any]) -> dict[str, Any]: + """Replace values whose keys contain secret patterns with '***'. + + Match is case-insensitive substring match against the key. Recursively + walks nested dicts and lists so a payload like + ``{"connection_info": {"password": "..."}}`` is also redacted. + Input is not mutated. + """ + + def _walk(value: Any, key_hint: str | None = None) -> Any: + if key_hint and any(pat in key_hint.lower() for pat in _SECRET_PATTERNS): + return "***" + if isinstance(value, dict): + return {k: _walk(v, k) for k, v in value.items()} + if isinstance(value, list): + return [_walk(v, key_hint) for v in value] + return value + + return {k: _walk(v, k) for k, v in data.items()} + + +def make_success( + content: str, + data: dict[str, Any], + warnings: list[str] | None = None, +) -> dict[str, Any]: + """Construct a success envelope.""" + return { + "ok": True, + "content": content, + "data": data, + "warnings": warnings or [], + } + + +def format_error(exc: Exception) -> dict[str, Any]: + """Convert an exception into the structured error dict used in envelopes.""" + if isinstance(exc, WrenError): + metadata = redact_secrets(exc.metadata or {}) + metadata = json_safe(metadata) + metadata = cap_size(metadata, max_bytes=_DEFAULT_METADATA_CAP) + return { + "code": exc.error_code.name, + "phase": exc.phase.name if exc.phase else None, + "message": exc.message, + "metadata": metadata, + } + return { + "code": "SDK_ERROR", + "phase": None, + "message": str(exc), + "metadata": {}, + } + + +def make_error(exc: Exception) -> dict[str, Any]: + """Construct a failure envelope from an exception.""" + error = format_error(exc) + return { + "ok": False, + "content": error["message"], + "error": error, + } diff --git a/sdk/wren-langchain/src/wren_langchain/_format.py b/sdk/wren-langchain/src/wren_langchain/_format.py new file mode 100644 index 000000000..82d308519 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_format.py @@ -0,0 +1,149 @@ +"""Content formatters per tool. Used by tool wrappers to produce the +LLM-facing ``content`` field of the envelope.""" + +from __future__ import annotations + +import json +from typing import Any + +import pyarrow as pa + +CONTENT_CAP_BYTES = 16 * 1024 + + +def format_query_content( + table: pa.Table, total_rows: int | None = None +) -> tuple[str, list[str]]: + """Render query rows as JSON, truncating to fit ``CONTENT_CAP_BYTES``. + + Returns ``(content, warnings)``. The content is a JSON array of row dicts + plus an optional ``(showing N of M rows)`` footer when truncated. + """ + rows = table.to_pylist() + full = json.dumps(rows, default=str) + encoded = full.encode("utf-8") + if len(encoded) <= CONTENT_CAP_BYTES: + return full, [] + + # Binary search-ish for largest prefix that fits. + keep = len(rows) + while keep > 0: + partial = json.dumps(rows[:keep], default=str) + if len(partial.encode("utf-8")) + 64 <= CONTENT_CAP_BYTES: + break + keep -= max(1, keep // 4) + + total = total_rows if total_rows is not None else len(rows) + + # Edge case: even the first row alone exceeds the content cap. Returning + # `[]` with "showing 0 of N rows" would be misleading — the LLM would + # think the query returned nothing. Surface a clear error message + # instead, with concrete guidance on how to recover. + if keep == 0: + msg = ( + f"(0 of {total} rows shown — a single row exceeds the " + f"{CONTENT_CAP_BYTES}-byte content cap. Use SELECT to pick fewer / " + "narrower columns, or aggregate in SQL.)" + ) + warning = ( + f"content truncated: row 1 alone exceeds {CONTENT_CAP_BYTES}-byte " + "cap; no rows shown." + ) + return msg, [warning] + + shown = rows[:keep] + body = json.dumps(shown, default=str) + footer = f"\n(showing {keep} of {total} rows)" + warning = f"content truncated: showed {keep} of {total} rows due to size cap" + return body + footer, [warning] + + +def format_dry_plan_content(sql: str) -> str: + """Wrap dialect SQL in a markdown code fence.""" + return f"```sql\n{sql}\n```" + + +def format_fetch_context_content(result: dict[str, Any]) -> str: + """Render get_context output as readable text. + + Strategy ``full`` returns the schema text directly. + Strategy ``search`` returns a numbered list of items. + """ + strategy = result.get("strategy") + if strategy == "full": + text = result.get("schema", "") + return _cap_to_bytes(text, suffix="\n\n...[truncated]") + + items = result.get("results", []) or [] + if not items: + return "_No relevant context items found._" + + lines = [] + for i, item in enumerate(items, start=1): + item_type = item.get("item_type", "item") + name = item.get("name", "") + summary = item.get("summary") or item.get("text") or "" + if len(summary) > 120: + summary = summary[:117] + "..." + lines.append(f"{i}. [{item_type}] {name} — {summary}") + return _cap_to_bytes("\n".join(lines), suffix="\n...[truncated]") + + +def _cap_to_bytes(text: str, *, suffix: str) -> str: + """Truncate *text* so its UTF-8 size never exceeds CONTENT_CAP_BYTES. + + Truncation is byte-aware (cuts on a UTF-8 boundary, not a char count) so + multibyte chars cannot push the encoded size past the cap. The trailing + *suffix* is reserved before slicing. + """ + encoded = text.encode("utf-8") + if len(encoded) <= CONTENT_CAP_BYTES: + return text + suffix_bytes = suffix.encode("utf-8") + budget = max(0, CONTENT_CAP_BYTES - len(suffix_bytes)) + return encoded[:budget].decode("utf-8", errors="ignore") + suffix + + +def format_recall_content(rows: list[dict[str, Any]]) -> str: + """Render recalled NL→SQL pairs as a numbered list with code fences.""" + if not rows: + return "_No similar past queries found._" + + chunks = [] + for i, row in enumerate(rows, start=1): + nl = row.get("nl_query") or row.get("nl") or "" + sql = row.get("sql_query") or row.get("sql") or "" + chunks.append(f'{i}. "{nl}"\n ```sql\n {sql}\n ```') + return "\n".join(chunks) + + +def format_store_content(nl: str, sql: str, tags: list[str] | None) -> str: + """One-liner ``Stored: "" → (N tags)``.""" + sql_preview = sql.strip().split("\n")[0] + if len(sql_preview) > 80: + sql_preview = sql_preview[:77] + "..." + tag_count = len(tags) if tags else 0 + return f'Stored: "{nl}" → {sql_preview} ({tag_count} tags)' + + +def format_list_models_content(manifest: dict[str, Any]) -> str: + """Render manifest models as a compact markdown table. + + Columns: model | cols | description. + """ + models = manifest.get("models", []) or [] + if not models: + return "_No models defined in this Wren project._" + + lines = ["| model | cols | description |", "|---|---|---|"] + for m in models: + name = m.get("name", "") + col_count = len(m.get("columns", []) or []) + desc = ( + (m.get("properties") or {}).get("description") or m.get("description") or "" + ) + # Trim long descriptions to keep table compact. + if len(desc) > 80: + desc = desc[:77] + "..." + lines.append(f"| {name} | {col_count} | {desc} |") + return "\n".join(lines) diff --git a/sdk/wren-langchain/src/wren_langchain/_memory_api.py b/sdk/wren-langchain/src/wren_langchain/_memory_api.py new file mode 100644 index 000000000..37545ed59 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_memory_api.py @@ -0,0 +1,98 @@ +"""Direct Python subscope for memory operations. + +Exposed as ``toolkit.memory``. Operations raise ``MemoryNotEnabledError`` +when memory is disabled — distinct error model from LLM tools, which +silently filter out memory tools when disabled. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from wren_langchain.exceptions import MemoryNotEnabledError + +if TYPE_CHECKING: + from wren.memory.store import MemoryStore + + from wren_langchain._toolkit import WrenToolkit + + +class _MemoryAPI: + """Thin wrapper around ``wren.memory.MemoryStore`` bound to a toolkit.""" + + def __init__(self, toolkit: WrenToolkit): + self._toolkit = toolkit + + def fetch( + self, + question: str, + *, + limit: int = 5, + item_type: str | None = None, + model: str | None = None, + threshold: int | None = None, + ) -> dict[str, Any]: + """Return schema/business context relevant to *question*.""" + store = self._store() + manifest = self._toolkit._mdl_source.load_manifest() + kwargs: dict[str, Any] = { + "query": question, + "manifest": manifest, + "limit": limit, + } + if item_type is not None: + kwargs["item_type"] = item_type + if model is not None: + kwargs["model_name"] = model + if threshold is not None: + kwargs["threshold"] = threshold + return store.get_context(**kwargs) + + def recall( + self, + question: str, + *, + limit: int = 3, + ) -> list[dict[str, Any]]: + """Return up to *limit* past NL→SQL pairs similar to *question*.""" + store = self._store() + return store.recall_queries(query=question, limit=limit) + + def store( + self, + nl: str, + sql: str, + *, + tags: list[str] | None = None, + ) -> None: + """Persist a confirmed NL→SQL pair for future recall. + + Tags are joined with commas before storage; Core's MemoryStore stores + them as an opaque string. Tags must therefore not contain commas + themselves — a tag like ``"revenue, Q1"`` would be silently split into + two tags on any future consumer that splits on the separator. We + reject such inputs early with ``ValueError`` rather than corrupt + the round-trip. + Empty/None tags map to no tag. + """ + if tags: + for tag in tags: + if "," in tag: + raise ValueError( + f"tag {tag!r} contains a comma; commas are reserved as the " + "separator for the underlying storage format. " + "Replace commas with dashes or spaces." + ) + store = self._store() + tag_str = ",".join(tags) if tags else None + store.store_query(nl_query=nl, sql_query=sql, tags=tag_str) + + def _store(self) -> MemoryStore: + if not self._toolkit._memory.enabled: + raise MemoryNotEnabledError( + "memory is not enabled for this toolkit. " + "Run `wren memory index` in your project to enable it." + ) + if self._toolkit._memory_store_cache is None: + self._toolkit._memory_store_cache = self._toolkit._memory.open() + return self._toolkit._memory_store_cache diff --git a/sdk/wren-langchain/src/wren_langchain/_prompt.py b/sdk/wren-langchain/src/wren_langchain/_prompt.py new file mode 100644 index 000000000..8d1706a68 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_prompt.py @@ -0,0 +1,224 @@ +"""Build a Wren-aware system prompt for LangChain/LangGraph agents. + +The workflow distilled here mirrors the Wren CLI's `wren-usage` skill — +recall → fetch context → write SQL → dry_plan if complex → execute → store — +adapted to the SDK's tool surface. + +Defaults are deliberately strong ("recall by default", "store by default") +because empirical testing showed soft phrasing ("for non-trivial questions", +"when useful") was almost always interpreted by GPT-4o as "skip". The CLI +skill takes the same stance and bakes the strong defaults into its workflow. + +The prompt is **derived from the actual tool list** so it stays in sync with +``toolkit.get_tools(include_memory_write=...)``. If a caller hides +``wren_store_query``, the workflow drops the persistence step rather than +instructing the agent to call a tool that no longer exists. + +Three markdown sections (any may be empty): + 1. Workflow rules — auto-adapted to the supplied tool list. + 2. Available tools — bullet list rendered from the same list. + 3. Project-specific instructions — content of ``/instructions.md`` + when present; silently omitted when absent. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Iterable + +if TYPE_CHECKING: + from wren_langchain._toolkit import WrenToolkit + + +_INTRO = ( + "You use Wren Engine as the semantic layer for data querying. SQL " + "targets MDL model names (defined in `target/mdl.json`); the engine " + "translates to the target database dialect." +) + +# Workflow step blobs. Composed conditionally based on which tools are present. +_STEP_RECALL = """1. Recall similar past NL→SQL pairs: + `wren_recall_queries(question="", limit=3)` + Use the results as few-shot examples. Empty results are fine — continue + to the next step. Do NOT skip this step on the grounds that the question + seems simple; past pairs may use better joins, filters, or column names + than you would write from scratch.""" + +_STEP_FETCH = """2. Fetch schema and business context: + `wren_fetch_context(question="")` + Optionally narrow scope with `model=""` or + `item_type="model" | "column" | "relationship" | "view"`.""" + +_STEP_LIST_MODELS_FALLBACK = """1. If you don't already know the available models, call `wren_list_models()` + to enumerate them.""" + +_STEP_COMPOSE = ( + "{n}. Compose SQL targeting Wren model names — NEVER raw database tables." +) + +_STEP_DRY_PLAN = """{n}. (Complex queries only) Verify with `wren_dry_plan(sql="...")` before + executing. "Complex" = subqueries, multi-step CTEs, or JOINs not + already defined as MDL relationships. Simple GROUP BY or + model-defined JOINs can skip this step.""" + +_STEP_QUERY = """{n}. Execute: `wren_query(sql="...", limit=100)`. Raise the limit only when + you genuinely need more rows.""" + +_STEP_STORE = """{n}. Persist the NL→SQL pair: `wren_store_query(nl="", sql="", tags=[...])`. + + Store BY DEFAULT after a successful query. Skip ONLY when: + - The query failed (`ok=false`). + - The user said the result is wrong. + - The SQL is exploratory (e.g. `SELECT * FROM x LIMIT 10` with no + analytical clauses). + - There is no natural-language question (e.g. the user pasted raw SQL). + - The user explicitly said don't save. + + The `nl` value should be the user's original question, not a paraphrase.""" + + +def _build_workflow_section(tool_names: set[str]) -> str: + """Compose the workflow header from tool blobs based on which tools are + actually available. Numbering is dynamic so steps stay sequential when + optional ones are dropped.""" + has_recall = "wren_recall_queries" in tool_names + has_fetch = "wren_fetch_context" in tool_names + has_store = "wren_store_query" in tool_names + has_dry_plan = "wren_dry_plan" in tool_names + has_list_models = "wren_list_models" in tool_names + + steps: list[str] = [] + if has_recall: + steps.append(_STEP_RECALL) + if has_fetch: + steps.append(_STEP_FETCH) + elif has_list_models and not has_recall: + steps.append(_STEP_LIST_MODELS_FALLBACK) + + next_n = len(steps) + 1 + steps.append(_STEP_COMPOSE.format(n=next_n)) + next_n += 1 + if has_dry_plan: + steps.append(_STEP_DRY_PLAN.format(n=next_n)) + next_n += 1 + steps.append(_STEP_QUERY.format(n=next_n)) + next_n += 1 + if has_store: + steps.append(_STEP_STORE.format(n=next_n)) + + intro = "Run these steps in order:" if len(steps) > 2 else "" + body = "\n\n".join(steps) + if intro: + body = f"{intro}\n\n{body}" + return f"# Workflow for every data question\n\n{body}" + + +def _error_recovery_section(*, has_fetch_context: bool, has_list_models: bool) -> str: + if has_fetch_context: + find_name_hint = ( + 'use `wren_fetch_context(question="", ' + 'item_type="model")` (or `item_type="column"`) to find the ' + "correct one." + ) + elif has_list_models: + find_name_hint = ( + "use `wren_list_models()` to enumerate models and inspect their columns." + ) + else: + find_name_hint = "consult the project's MDL files directly." + return f"""# Error recovery + +If a tool returns `ok=false`, inspect `error.phase` and `error.message`: + +- `SQL_PARSING` → SQL syntax error. Read the message, fix, and retry. +- `METADATA_FETCHING` / `MDL_EXTRACTION` → wrong model or column name; + {find_name_hint} +- `SQL_EXECUTION` → database-side error. `error.metadata.dialect_sql` shows + the translated SQL — diagnose against the message (type mismatch, missing + function, permission, timeout). Add explicit `CAST` or simplify the query + if needed. + +Don't silently abandon. Either fix and retry, or report the failure to the +user along with what you tried.""" + + +def _things_to_avoid_section(tool_names: set[str]) -> str: + bullets = [] + if "wren_fetch_context" in tool_names: + bullets.append( + "- Don't guess model or column names — call `wren_fetch_context` first." + ) + elif "wren_list_models" in tool_names: + bullets.append( + "- Don't guess model or column names — call `wren_list_models()` " + "first when in doubt." + ) + if "wren_recall_queries" in tool_names: + bullets.append( + '- Don\'t skip `wren_recall_queries` on questions that seem "simple" — ' + "past pairs are often the most accurate template." + ) + if "wren_store_query" in tool_names: + bullets.append( + "- Don't store failed queries, queries the user said are wrong, " + "or exploratory queries." + ) + bullets.append("- Don't store SQL that has no clear natural-language question.") + bullets.append( + "- Don't write SQL against raw database tables — always use MDL model names." + ) + return "# Things to avoid\n\n" + "\n".join(bullets) + + +def build_system_prompt(toolkit: WrenToolkit, *, tools: Iterable | None = None) -> str: + """Render the system prompt. + + ``tools`` is the actual tool list used by the agent. When ``None``, the + toolkit's default ``get_tools()`` output is used. Pass the same list you + pass to ``create_agent(..., tools=...)`` so the prompt stays in sync — + e.g., when ``include_memory_write=False`` is used to drop + ``wren_store_query``, the workflow's persistence step is dropped too. + """ + tool_list = list(tools) if tools is not None else list(toolkit.get_tools()) + tool_names = {t.name for t in tool_list} + + sections: list[str] = [_INTRO] + sections.append(_build_workflow_section(tool_names)) + sections.append( + _error_recovery_section( + has_fetch_context="wren_fetch_context" in tool_names, + has_list_models="wren_list_models" in tool_names, + ) + ) + sections.append(_things_to_avoid_section(tool_names)) + + tools_section = _build_tools_section(tool_list) + if tools_section: + sections.append(tools_section) + + instructions_section = _build_instructions_section(toolkit._project_path) + if instructions_section: + sections.append(instructions_section) + + return "\n\n".join(sections) + + +def _build_tools_section(tool_list: list) -> str: + if not tool_list: + return "" + lines = ["## Available tools"] + for tool in tool_list: + description = (tool.description or "").strip().split("\n")[0] + lines.append(f"- `{tool.name}`: {description}") + return "\n".join(lines) + + +def _build_instructions_section(project_path: Path) -> str: + instructions_file = project_path / "instructions.md" + if not instructions_file.exists(): + return "" + body = instructions_file.read_text().strip() + if not body: + return "" + return f"## Project-specific instructions\n\n{body}" diff --git a/sdk/wren-langchain/src/wren_langchain/_providers/__init__.py b/sdk/wren-langchain/src/wren_langchain/_providers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/sdk/wren-langchain/src/wren_langchain/_providers/connection.py b/sdk/wren-langchain/src/wren_langchain/_providers/connection.py new file mode 100644 index 000000000..e1a8f0e29 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_providers/connection.py @@ -0,0 +1,80 @@ +"""Connection providers resolve a Wren profile to (datasource, connection_info). + +Layer order (highest priority first): + 1. ``explicit_profile=`` kwarg passed to ``WrenToolkit.from_project`` + 2. ``profile:`` field in the project's ``wren_project.yml`` + 3. The user's globally active profile (``wren profile switch ...``) + +Secrets in profile values (``${ENV_VAR}``) are expanded via Core's +``expand_profile_secrets`` before the connection is exposed. +""" + +from pathlib import Path +from typing import Any + +from wren.context import load_project_config +from wren.profile import ( # noqa: F401 re-exported for monkeypatching in tests + expand_profile_secrets, + get_active_profile, + list_profiles, +) + +from wren_langchain.exceptions import WrenToolkitInitError + + +class ProfileConnectionProvider: + """Resolves a profile via the 3-layer fallback and exposes connection info.""" + + def __init__( + self, + *, + project_path: Path, + explicit_profile: str | None = None, + ): + self._project_path = project_path + profile_dict = self._resolve_profile(explicit_profile) + profile_dict = expand_profile_secrets(profile_dict) + # ``datasource`` is part of the profile dict but logically separate. + self._datasource = profile_dict.pop("datasource", None) + self._connection_info = profile_dict + + def _resolve_profile(self, explicit: str | None) -> dict[str, Any]: + # Layer 1: explicit kwarg. Use `is not None` so a misconfigured caller + # passing `profile=""` raises a clear "profile not found" error instead + # of silently falling through to the project-config / active layers. + if explicit is not None: + return self._lookup_named_profile(explicit) + + # Layer 2: profile name from wren_project.yml + project_profile_name = self._project_config_profile() + if project_profile_name: + return self._lookup_named_profile(project_profile_name) + + # Layer 3: globally active profile + _, active = get_active_profile() + if not active: + raise WrenToolkitInitError( + "no active Wren profile found. " + "Run `wren profile add` and `wren profile switch` first." + ) + return dict(active) + + def _lookup_named_profile(self, name: str) -> dict[str, Any]: + profiles = list_profiles() + if name not in profiles: + raise WrenToolkitInitError( + f"profile {name!r} not found in ~/.wren/profiles.yml. " + f"Available: {sorted(profiles)}" + ) + return dict(profiles[name]) + + def _project_config_profile(self) -> str | None: + config = load_project_config(self._project_path) + value = config.get("profile") + return str(value) if value else None + + def datasource(self) -> str | None: + return self._datasource + + def connection_info(self) -> dict[str, Any]: + return dict(self._connection_info) diff --git a/sdk/wren-langchain/src/wren_langchain/_providers/mdl_source.py b/sdk/wren-langchain/src/wren_langchain/_providers/mdl_source.py new file mode 100644 index 000000000..b639ab5c8 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_providers/mdl_source.py @@ -0,0 +1,42 @@ +"""MDL sources resolve where the manifest comes from. + +v0.1 ships only ``ProjectMDLSource`` which reads ``target/mdl.json`` from a +prepared Wren project directory. Each ``load_manifest()`` call re-reads the +file from disk so that ``wren context build`` updates by an external CLI run +are picked up on the next tool invocation without needing ``toolkit.reload()``. +""" + +import json +from pathlib import Path +from typing import Any + +from wren_langchain.exceptions import WrenToolkitInitError + + +class ProjectMDLSource: + """Read the manifest from ``/target/mdl.json``.""" + + def __init__(self, *, project_path: Path): + self._project_path = project_path + self._mdl_path = project_path / "target" / "mdl.json" + + def load_manifest(self) -> dict[str, Any]: + if not self._mdl_path.exists(): + raise WrenToolkitInitError( + f"target/mdl.json not found at {self._mdl_path}. " + "Run `wren context build` first." + ) + try: + return json.loads(self._mdl_path.read_text()) + except json.JSONDecodeError as exc: + # Normalize malformed manifest into the common init-error contract + # so callers don't need to special-case JSON errors. + raise WrenToolkitInitError( + f"target/mdl.json at {self._mdl_path} is not valid JSON: {exc.msg} " + f"(line {exc.lineno}, col {exc.colno}). " + "Re-run `wren context build` to regenerate it." + ) from exc + + @property + def mdl_path(self) -> Path: + return self._mdl_path diff --git a/sdk/wren-langchain/src/wren_langchain/_providers/memory.py b/sdk/wren-langchain/src/wren_langchain/_providers/memory.py new file mode 100644 index 000000000..96faf1162 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_providers/memory.py @@ -0,0 +1,41 @@ +"""Memory providers resolve where the long-lived context store lives. + +v0.1 ships: + - ``LocalLanceDBMemoryProvider``: opens a ``MemoryStore`` against a local + ``.wren/memory/`` directory. + - ``NoopMemoryProvider``: signals that memory is disabled. Direct API calls + raise ``MemoryNotEnabledError``; LLM-facing tools are filtered out. + +Auto-selection is performed by ``WrenToolkit.from_project`` based on whether +``/.wren/memory/`` exists. +""" + +from pathlib import Path + +from wren.memory.store import MemoryStore + +from wren_langchain.exceptions import MemoryNotEnabledError + + +class LocalLanceDBMemoryProvider: + """Lazily opens a local LanceDB-backed ``MemoryStore`` on first use.""" + + enabled = True + + def __init__(self, *, memory_path: Path): + self._memory_path = memory_path + + def open(self) -> MemoryStore: + return MemoryStore(path=self._memory_path) + + +class NoopMemoryProvider: + """Inert provider used when no ``.wren/memory/`` exists in the project.""" + + enabled = False + + def open(self) -> MemoryStore: + raise MemoryNotEnabledError( + "memory is not enabled for this toolkit. " + "Run `wren memory index` in your project to enable it." + ) diff --git a/sdk/wren-langchain/src/wren_langchain/_toolkit.py b/sdk/wren-langchain/src/wren_langchain/_toolkit.py new file mode 100644 index 000000000..ada34344f --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_toolkit.py @@ -0,0 +1,229 @@ +"""WrenToolkit: facade over an existing CLI-prepared Wren project.""" + +from __future__ import annotations + +import base64 +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from wren.engine import WrenEngine + +from wren_langchain._memory_api import _MemoryAPI +from wren_langchain._prompt import build_system_prompt +from wren_langchain._providers.connection import ProfileConnectionProvider +from wren_langchain._providers.mdl_source import ProjectMDLSource +from wren_langchain._providers.memory import ( + LocalLanceDBMemoryProvider, + NoopMemoryProvider, +) +from wren_langchain._tools import build_runtime_tools +from wren_langchain._tools_memory import build_memory_tools +from wren_langchain.exceptions import WrenToolkitInitError + +if TYPE_CHECKING: + import pyarrow as pa + + +class WrenToolkit: + """Adapter that exposes an existing Wren project as LangChain tools.""" + + def __init__( + self, + *, + project_path: Path, + mdl_source: ProjectMDLSource, + connection_provider: ProfileConnectionProvider, + memory_provider: LocalLanceDBMemoryProvider | NoopMemoryProvider, + ): + self._project_path = project_path + self._mdl_source = mdl_source + self._connection = connection_provider + self._memory = memory_provider + # Connector is cached at the toolkit level to avoid reconnecting on + # every query. The engine itself is rebuilt per call so manifest + # changes are picked up read-through. + self._connector_cache: Any = None + # MemoryStore is heavy (loads sentence-transformer model) — cache + # the instance and let LanceDB handle data versioning internally. + self._memory_store_cache: Any = None + + # ── Memory subscope (exposed as toolkit.memory) ──────────────────────── + + @property + def memory(self): + if not hasattr(self, "_memory_api"): + self._memory_api = _MemoryAPI(self) + return self._memory_api + + # ── Direct Python API ────────────────────────────────────────────────── + + def query(self, sql: str, limit: int | None = None) -> pa.Table: + """Execute SQL through the Wren context layer. Returns a pyarrow Table.""" + engine = self._build_engine() + try: + result = engine.query(sql, limit=limit) + finally: + self._connector_cache = engine._connector + return result + + def dry_plan(self, sql: str) -> str: + """Plan SQL through MDL and return the expanded SQL in target dialect.""" + return self._build_engine().dry_plan(sql) + + def dry_run(self, sql: str) -> None: + """Validate SQL by planning and asking the DB to plan it without executing.""" + engine = self._build_engine() + try: + engine.dry_run(sql) + finally: + self._connector_cache = engine._connector + + # ── LangChain adapter ────────────────────────────────────────────────── + + def get_tools( + self, + *, + include_memory_write: bool = True, + raise_on_error: bool = False, + ) -> list: + """Return LangChain-compatible tools bound to this toolkit. + + Memory tools are auto-filtered when memory is disabled (no + ``.wren/memory/`` directory in the project). ``include_memory_write=False`` + removes ``wren_store_query`` while keeping the read-only memory tools + (``wren_fetch_context``, ``wren_recall_queries``). + + When memory is disabled, ``include_memory_write`` has no effect — no + memory tools are added regardless of its value. + """ + tools = build_runtime_tools(self, raise_on_error=raise_on_error) + if self._memory.enabled: + tools.extend( + build_memory_tools( + self, + raise_on_error=raise_on_error, + include_write=include_memory_write, + ) + ) + return tools + + def system_prompt(self, *, tools=None) -> str: + """Return a Wren-aware system prompt suitable for LangChain agents. + + Composition: + 1. Workflow rules — derived from the supplied tool list so the + prompt stays in sync with what the agent actually has. + 2. Available tools — bullet list rendered from the same list. + 3. Project-specific instructions (from ``instructions.md`` if present). + + Pass the same ``tools`` you give to ``create_agent`` if you customized + the toolset — e.g. ``get_tools(include_memory_write=False)`` — so the + workflow drops the persistence step instead of instructing the agent + to call a tool that no longer exists. When ``tools`` is omitted, the + toolkit's default ``get_tools()`` output is used. + + To extend, concatenate with your own instructions:: + + tools = toolkit.get_tools(include_memory_write=False) + prompt = ( + f"You are a finance analyst.\\n\\n" + f"{toolkit.system_prompt(tools=tools)}" + ) + """ + return build_system_prompt(self, tools=tools) + + # ── Internal ─────────────────────────────────────────────────────────── + + def _build_engine(self) -> WrenEngine: + """Construct a fresh WrenEngine with a read-through manifest. + + The connector is reused across calls when available so DB authentication + only happens once per toolkit lifetime. + """ + manifest = self._mdl_source.load_manifest() + manifest_str = base64.b64encode(json.dumps(manifest).encode("utf-8")).decode() + engine = WrenEngine( + manifest_str=manifest_str, + data_source=self._connection.datasource(), + connection_info=self._connection.connection_info(), + ) + if self._connector_cache is not None: + engine._connector = self._connector_cache + return engine + + @classmethod + def from_project( + cls, + path: str | Path, + *, + profile: str | None = None, + ) -> WrenToolkit: + """Build a toolkit from a CLI-prepared Wren project directory. + + Memory is auto-detected from ``/.wren/memory/``: present → + memory tools are exposed, absent → only the 3 runtime tools. + To enable, run ``wren memory index`` in the project; to disable, + delete the directory. There is no kwarg to override. + """ + project_path = Path(path).expanduser().resolve() + + if not (project_path / "wren_project.yml").exists(): + raise WrenToolkitInitError( + f"wren_project.yml not found at {project_path}. " + "Is this a Wren project? Run `wren context init` to create one." + ) + + if not (project_path / "target" / "mdl.json").exists(): + raise WrenToolkitInitError( + f"target/mdl.json not found at {project_path}/target/mdl.json. " + "Run `wren context build` first." + ) + + cls._load_project_dotenv(project_path) + + mdl_source = ProjectMDLSource(project_path=project_path) + connection = ProfileConnectionProvider( + project_path=project_path, + explicit_profile=profile, + ) + memory_provider = cls._resolve_memory_provider(project_path) + + return cls( + project_path=project_path, + mdl_source=mdl_source, + connection_provider=connection, + memory_provider=memory_provider, + ) + + @staticmethod + def _load_project_dotenv(project_path: Path) -> None: + """Load ``/.env`` into ``os.environ`` if present. + + Required for SDK ergonomics: when a caller passes + ``from_project("/some/path")`` from anywhere on the filesystem, they + expect that project's secrets to resolve. Core's ``expand_profile_secrets`` + discovers ``.env`` relative to CWD, which doesn't help here. + + Uses ``override=False`` so values the user already exported in their + shell still win, matching Core's policy. + """ + env_path = project_path / ".env" + if not env_path.exists(): + return + try: + from dotenv import load_dotenv # noqa: PLC0415 + except ImportError: + return + load_dotenv(env_path, override=False) + + @staticmethod + def _resolve_memory_provider( + project_path: Path, + ) -> LocalLanceDBMemoryProvider | NoopMemoryProvider: + memory_dir = project_path / ".wren" / "memory" + # Require a directory (not a regular file or broken symlink) so we + # never construct LocalLanceDBMemoryProvider against an invalid root. + if memory_dir.is_dir(): + return LocalLanceDBMemoryProvider(memory_path=memory_dir) + return NoopMemoryProvider() diff --git a/sdk/wren-langchain/src/wren_langchain/_tools.py b/sdk/wren-langchain/src/wren_langchain/_tools.py new file mode 100644 index 000000000..600602680 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_tools.py @@ -0,0 +1,120 @@ +"""LangChain tool wrappers exposing toolkit operations to LLMs. + +Each tool returns a JSON-serializable envelope: success or error. +On error, the envelope's ``ok=False`` allows the agent to inspect ``error.phase`` +and ``error.code`` to recover. Set ``raise_on_error=True`` on get_tools() to +raise instead. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from langchain_core.tools import tool + +from wren_langchain._envelope import make_error, make_success +from wren_langchain._format import ( + format_dry_plan_content, + format_list_models_content, + format_query_content, +) + +if TYPE_CHECKING: + from wren_langchain._toolkit import WrenToolkit + + +# Hard cap for the LLM-facing `wren_query` tool. The 16 KB content cap +# already truncates the rendered preview, but `data.rows` materializes every +# row via `to_pylist()` — a runaway `limit` value (typo, hallucinated huge +# number) would still balloon memory before that cap fires. 1000 rows leaves +# headroom over the 100-default while keeping payloads bounded. Direct API +# (`toolkit.query`) keeps no cap on purpose; that's a Python-programmer surface. +MAX_QUERY_ROWS = 1000 + + +def build_runtime_tools(toolkit: WrenToolkit, *, raise_on_error: bool) -> list: + """Return wren_query, wren_dry_plan, wren_list_models bound to toolkit.""" + return [ + _build_wren_query(toolkit, raise_on_error=raise_on_error), + _build_wren_dry_plan(toolkit, raise_on_error=raise_on_error), + _build_wren_list_models(toolkit, raise_on_error=raise_on_error), + ] + + +def _build_wren_query(toolkit: WrenToolkit, *, raise_on_error: bool): + @tool("wren_query") + def wren_query(sql: str, limit: int = 100) -> dict[str, Any]: + """Execute SQL through the Wren semantic layer and return rows. + + Use this after wren_dry_plan looks correct. Default limit is 100 rows; + increase only when you need more. Hard cap is 1000 rows — beyond that, + aggregate in SQL instead. + """ + if limit < 1 or limit > MAX_QUERY_ROWS: + err = ValueError( + f"limit must be between 1 and {MAX_QUERY_ROWS} (got {limit}). " + "Aggregate in SQL if you need more rows." + ) + if raise_on_error: + raise err + return make_error(err) + + try: + table = toolkit.query(sql, limit=limit) + except Exception as exc: + if raise_on_error: + raise + return make_error(exc) + + content, warnings = format_query_content(table, total_rows=table.num_rows) + data = { + "columns": table.column_names, + "rows": table.to_pylist(), + "row_count": table.num_rows, + "content_truncated": bool(warnings), + } + return make_success(content=content, data=data, warnings=warnings) + + return wren_query + + +def _build_wren_dry_plan(toolkit: WrenToolkit, *, raise_on_error: bool): + @tool("wren_dry_plan") + def wren_dry_plan(sql: str) -> dict[str, Any]: + """Plan SQL through MDL and return the expanded target-dialect SQL. + + Use this to verify your SQL targets Wren models correctly before + running wren_query. Cheap (no DB round-trip). + """ + try: + dialect_sql = toolkit.dry_plan(sql) + except Exception as exc: + if raise_on_error: + raise + return make_error(exc) + + return make_success( + content=format_dry_plan_content(dialect_sql), + data={"dialect_sql": dialect_sql}, + ) + + return wren_dry_plan + + +def _build_wren_list_models(toolkit: WrenToolkit, *, raise_on_error: bool): + @tool("wren_list_models") + def wren_list_models() -> dict[str, Any]: + """List all models defined in this Wren project with column counts and descriptions.""" + try: + manifest = toolkit._mdl_source.load_manifest() + except Exception as exc: + if raise_on_error: + raise + return make_error(exc) + + return make_success( + content=format_list_models_content(manifest), + data={"models": manifest.get("models", []) or []}, + ) + + return wren_list_models diff --git a/sdk/wren-langchain/src/wren_langchain/_tools_memory.py b/sdk/wren-langchain/src/wren_langchain/_tools_memory.py new file mode 100644 index 000000000..b337b41e9 --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/_tools_memory.py @@ -0,0 +1,123 @@ +"""LangChain tool wrappers for memory operations. + +Three tools, all enabled when memory is detected (``.wren/memory/`` exists): + - ``wren_fetch_context``: schema/business context via embedding search + - ``wren_recall_queries``: similar past NL→SQL pairs + - ``wren_store_query``: persist a confirmed pair (filtered by include_write) +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +from langchain_core.tools import tool + +from wren_langchain._envelope import make_error, make_success +from wren_langchain._format import ( + format_fetch_context_content, + format_recall_content, + format_store_content, +) + +if TYPE_CHECKING: + from wren_langchain._toolkit import WrenToolkit + + +def build_memory_tools( + toolkit: WrenToolkit, + *, + raise_on_error: bool, + include_write: bool, +) -> list: + tools = [ + _build_fetch_context(toolkit, raise_on_error=raise_on_error), + _build_recall_queries(toolkit, raise_on_error=raise_on_error), + ] + if include_write: + tools.append(_build_store_query(toolkit, raise_on_error=raise_on_error)) + return tools + + +def _build_fetch_context(toolkit: WrenToolkit, *, raise_on_error: bool): + @tool("wren_fetch_context") + def wren_fetch_context( + question: str, + limit: int = 5, + item_type: Literal["model", "column", "relationship", "view"] | None = None, + model: str | None = None, + ) -> dict[str, Any]: + """Fetch relevant schema and business context for an analytical question. + + Call this BEFORE writing SQL so you query the correct Wren models and + columns. Use ``item_type`` to narrow scope (e.g. only columns) and + ``model`` to narrow to a single model when known. + """ + try: + result = toolkit.memory.fetch( + question, limit=limit, item_type=item_type, model=model + ) + except Exception as exc: + if raise_on_error: + raise + return make_error(exc) + + return make_success( + content=format_fetch_context_content(result), + data=result, + ) + + return wren_fetch_context + + +def _build_recall_queries(toolkit: WrenToolkit, *, raise_on_error: bool): + @tool("wren_recall_queries") + def wren_recall_queries(question: str, limit: int = 3) -> dict[str, Any]: + """Recall up to *limit* past NL→SQL pairs similar to *question*. + + Useful as few-shot examples before writing new SQL. Pairs are + previously confirmed by users (or seeded for the project). + """ + try: + rows = toolkit.memory.recall(question, limit=limit) + except Exception as exc: + if raise_on_error: + raise + return make_error(exc) + + return make_success( + content=format_recall_content(rows), + data={"results": rows}, + ) + + return wren_recall_queries + + +def _build_store_query(toolkit: WrenToolkit, *, raise_on_error: bool): + @tool("wren_store_query") + def wren_store_query( + nl: str, + sql: str, + tags: list[str] | None = None, + ) -> dict[str, Any]: + """Save a confirmed natural-language → SQL pair for future recall. + + Call this AFTER ``wren_query`` succeeds and the result was useful, + so future agent runs can recall the example via ``wren_recall_queries``. + """ + # Normalize once up-front so every downstream caller (memory.store, + # format_store_content, the data payload) sees a list — no hidden + # `None`-handling expectations to leak. + tags_list = tags or [] + try: + toolkit.memory.store(nl=nl, sql=sql, tags=tags_list) + except Exception as exc: + if raise_on_error: + raise + return make_error(exc) + + return make_success( + content=format_store_content(nl, sql, tags_list), + data={"nl": nl, "sql": sql, "tags": tags_list}, + ) + + return wren_store_query diff --git a/sdk/wren-langchain/src/wren_langchain/exceptions.py b/sdk/wren-langchain/src/wren_langchain/exceptions.py new file mode 100644 index 000000000..f1c1e78bc --- /dev/null +++ b/sdk/wren-langchain/src/wren_langchain/exceptions.py @@ -0,0 +1,18 @@ +"""SDK-specific exception types for wren-langchain.""" + + +class WrenToolkitInitError(Exception): + """Raised when ``WrenToolkit.from_project(...)`` cannot validate prerequisites. + + Examples include missing ``wren_project.yml``, missing ``target/mdl.json``, + or unresolvable profile. + """ + + +class MemoryNotEnabledError(Exception): + """Raised when memory operations are called but no memory provider is active. + + Triggered by direct API access to ``toolkit.memory.*`` when the toolkit + was initialized against a project without ``.wren/memory/``. LLM-facing + tools handle this case via tool filtering, not by raising. + """ diff --git a/sdk/wren-langchain/tests/__init__.py b/sdk/wren-langchain/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/sdk/wren-langchain/tests/conformance/__init__.py b/sdk/wren-langchain/tests/conformance/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/sdk/wren-langchain/tests/conformance/test_langchain_contract.py b/sdk/wren-langchain/tests/conformance/test_langchain_contract.py new file mode 100644 index 000000000..cfc9de625 --- /dev/null +++ b/sdk/wren-langchain/tests/conformance/test_langchain_contract.py @@ -0,0 +1,122 @@ +"""LangChain BaseTool conformance: each Wren tool must satisfy the contract. + +A LangChain agent (and LangGraph ToolNode) relies on: +- ``tool.name`` (str, fixed identifier) +- ``tool.description`` (str, non-empty) +- ``tool.args_schema`` (Pydantic model with the expected fields) +- ``tool.invoke({...})`` returns a JSON-serializable dict envelope +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pyarrow as pa +import pytest +from langchain_core.tools import BaseTool + +from wren_langchain import WrenToolkit + +_EXPECTED_TOOL_NAMES = { + "wren_query", + "wren_dry_plan", + "wren_list_models", + "wren_fetch_context", + "wren_recall_queries", + "wren_store_query", +} + + +def _all_tools(tmp_project): + """Build a toolkit with memory enabled and stubbed engine + memory store.""" + (tmp_project / ".wren" / "memory").mkdir(parents=True) + fake_store = MagicMock(name="MemoryStore") + fake_store.get_context.return_value = {"strategy": "search", "results": []} + fake_store.recall_queries.return_value = [] + fake_engine = MagicMock(name="WrenEngine") + fake_engine.query.return_value = pa.table({"x": [1]}) + fake_engine.dry_plan.return_value = "SELECT 1" + fake_engine._connector = MagicMock() + + with ( + patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store), + patch("wren_langchain._toolkit.WrenEngine", return_value=fake_engine), + ): + toolkit = WrenToolkit.from_project(tmp_project) + yield toolkit + + +@pytest.fixture +def all_tools(tmp_project, fake_active_profile): + yield from _all_tools(tmp_project) + + +def test_get_tools_yields_all_expected_tools(all_tools): + names = {t.name for t in all_tools.get_tools()} + assert names == _EXPECTED_TOOL_NAMES + + +@pytest.mark.parametrize("expected_name", sorted(_EXPECTED_TOOL_NAMES)) +def test_each_tool_is_a_basetool(all_tools, expected_name): + tools_by_name = {t.name: t for t in all_tools.get_tools()} + tool = tools_by_name[expected_name] + assert isinstance(tool, BaseTool) + + +@pytest.mark.parametrize("expected_name", sorted(_EXPECTED_TOOL_NAMES)) +def test_each_tool_has_non_empty_description(all_tools, expected_name): + tools_by_name = {t.name: t for t in all_tools.get_tools()} + tool = tools_by_name[expected_name] + assert tool.description + assert len(tool.description.strip()) > 10 + + +@pytest.mark.parametrize( + "tool_name,required_args", + [ + ("wren_query", {"sql"}), + ("wren_dry_plan", {"sql"}), + ("wren_list_models", set()), + ("wren_fetch_context", {"question"}), + ("wren_recall_queries", {"question"}), + ("wren_store_query", {"nl", "sql"}), + ], +) +def test_each_tool_args_schema_includes_expected_fields( + all_tools, tool_name, required_args +): + tools_by_name = {t.name: t for t in all_tools.get_tools()} + tool = tools_by_name[tool_name] + schema = tool.args_schema + assert schema is not None + fields = set(schema.model_fields.keys()) + for arg in required_args: + assert arg in fields, f"{tool_name} missing arg {arg!r}; fields={fields}" + + +@pytest.mark.parametrize( + "tool_name,invoke_args", + [ + ("wren_query", {"sql": "SELECT 1"}), + ("wren_dry_plan", {"sql": "SELECT 1"}), + ("wren_list_models", {}), + ("wren_fetch_context", {"question": "what models exist?"}), + ("wren_recall_queries", {"question": "top customers"}), + ("wren_store_query", {"nl": "x", "sql": "SELECT 1"}), + ], +) +def test_each_tool_invoke_returns_envelope_dict(all_tools, tool_name, invoke_args): + tools_by_name = {t.name: t for t in all_tools.get_tools()} + tool = tools_by_name[tool_name] + result = tool.invoke(invoke_args) + + assert isinstance(result, dict) + assert "ok" in result + if result["ok"]: + assert "content" in result + assert "data" in result + assert "warnings" in result + else: + assert "content" in result + assert "error" in result + assert "code" in result["error"] diff --git a/sdk/wren-langchain/tests/conftest.py b/sdk/wren-langchain/tests/conftest.py new file mode 100644 index 000000000..e7e84ce51 --- /dev/null +++ b/sdk/wren-langchain/tests/conftest.py @@ -0,0 +1,34 @@ +"""Shared pytest fixtures for wren-langchain tests.""" + +import json + +import pytest + + +@pytest.fixture +def tmp_project(tmp_path): + """A minimal valid Wren project directory. + + Layout: + / + wren_project.yml + target/mdl.json + """ + (tmp_path / "wren_project.yml").write_text("schema_version: 1\n") + target = tmp_path / "target" + target.mkdir() + (target / "mdl.json").write_text(json.dumps({"models": []})) + return tmp_path + + +@pytest.fixture +def fake_active_profile(monkeypatch): + """Patch profile resolution to return a duckdb in-memory active profile.""" + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: {"test": {"datasource": "duckdb", "path": ":memory:"}}, + ) + monkeypatch.setattr( + "wren_langchain._providers.connection.get_active_profile", + lambda: ("test", {"datasource": "duckdb", "path": ":memory:"}), + ) diff --git a/sdk/wren-langchain/tests/integration/__init__.py b/sdk/wren-langchain/tests/integration/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/sdk/wren-langchain/tests/integration/conftest.py b/sdk/wren-langchain/tests/integration/conftest.py new file mode 100644 index 000000000..ca7985dfb --- /dev/null +++ b/sdk/wren-langchain/tests/integration/conftest.py @@ -0,0 +1,78 @@ +"""Integration test fixtures: real DuckDB-backed Wren project.""" + +from __future__ import annotations + +import json + +import duckdb +import pytest + + +@pytest.fixture +def duckdb_project(tmp_path, monkeypatch): + """Build a real Wren project backed by an attached DuckDB file with sample data. + + Layout: + tmp_path/ + wren_project.yml + target/mdl.json + db/sample.duckdb (one table: customers) + """ + db_dir = tmp_path / "db" + db_dir.mkdir() + db_path = db_dir / "sample.duckdb" + + con = duckdb.connect(str(db_path)) + con.execute("CREATE TABLE customers(id INTEGER, name VARCHAR);") + con.execute("INSERT INTO customers VALUES (1, 'Acme'), (2, 'Globex');") + con.close() + + manifest = { + "catalog": "wren", + "schema": "public", + "models": [ + { + "name": "customers", + "tableReference": { + "catalog": "sample", + "schema": "main", + "table": "customers", + }, + "columns": [ + {"name": "id", "type": "integer"}, + {"name": "name", "type": "varchar"}, + ], + "primaryKey": "id", + "properties": {"description": "Customer master data"}, + } + ], + } + + (tmp_path / "wren_project.yml").write_text("schema_version: 1\n") + target = tmp_path / "target" + target.mkdir() + (target / "mdl.json").write_text(json.dumps(manifest)) + + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: { + "test": { + "datasource": "duckdb", + "url": str(db_dir), + "format": "duckdb", + } + }, + ) + monkeypatch.setattr( + "wren_langchain._providers.connection.get_active_profile", + lambda: ( + "test", + { + "datasource": "duckdb", + "url": str(db_dir), + "format": "duckdb", + }, + ), + ) + + return tmp_path diff --git a/sdk/wren-langchain/tests/integration/test_langgraph_toolnode.py b/sdk/wren-langchain/tests/integration/test_langgraph_toolnode.py new file mode 100644 index 000000000..224297258 --- /dev/null +++ b/sdk/wren-langchain/tests/integration/test_langgraph_toolnode.py @@ -0,0 +1,76 @@ +"""LangGraph integration: WrenToolkit tools work inside a compiled graph. + +This is a structural smoke test — no real LLM. It builds a minimal +``StateGraph`` containing a ``ToolNode`` over the toolkit's tools, then +hand-constructs an ``AIMessage`` with ``tool_calls`` (the shape an LLM +would emit) and verifies the graph routes the call to our tool and +produces a ``ToolMessage`` with the envelope content. + +Note on langgraph 1.0+: ``ToolNode`` can no longer be invoked standalone — +it must run inside a compiled ``StateGraph`` (or be reached via +``langchain.agents.create_agent`` which sets that up internally). +""" + +from __future__ import annotations + +from langchain_core.messages import AIMessage, ToolMessage +from langgraph.graph import START, MessagesState, StateGraph +from langgraph.prebuilt import ToolNode + +from wren_langchain import WrenToolkit + + +def _build_graph(tools): + graph = StateGraph(MessagesState) + graph.add_node("tools", ToolNode(tools)) + graph.add_edge(START, "tools") + return graph.compile() + + +def test_toolnode_invokes_wren_query_via_simulated_llm_tool_call(duckdb_project): + """A simulated LLM tool_call against a real DuckDB project flows through ToolNode.""" + toolkit = WrenToolkit.from_project(duckdb_project) + app = _build_graph(toolkit.get_tools()) + + ai_message = AIMessage( + content="", + tool_calls=[ + { + "name": "wren_query", + "args": {"sql": "SELECT id, name FROM customers ORDER BY id"}, + "id": "call_1", + "type": "tool_call", + } + ], + ) + + state = app.invoke({"messages": [ai_message]}) + + tool_messages = [m for m in state["messages"] if isinstance(m, ToolMessage)] + assert len(tool_messages) == 1 + assert tool_messages[0].tool_call_id == "call_1" + assert "Acme" in tool_messages[0].content + assert "Globex" in tool_messages[0].content + + +def test_toolnode_invokes_wren_list_models(duckdb_project): + """No-arg tool also works through ToolNode-in-graph.""" + toolkit = WrenToolkit.from_project(duckdb_project) + app = _build_graph(toolkit.get_tools()) + + ai_message = AIMessage( + content="", + tool_calls=[ + { + "name": "wren_list_models", + "args": {}, + "id": "call_2", + "type": "tool_call", + } + ], + ) + + state = app.invoke({"messages": [ai_message]}) + tool_messages = [m for m in state["messages"] if isinstance(m, ToolMessage)] + assert len(tool_messages) == 1 + assert "customers" in tool_messages[0].content diff --git a/sdk/wren-langchain/tests/integration/test_memory_tools.py b/sdk/wren-langchain/tests/integration/test_memory_tools.py new file mode 100644 index 000000000..0b0cf1e6e --- /dev/null +++ b/sdk/wren-langchain/tests/integration/test_memory_tools.py @@ -0,0 +1,68 @@ +"""End-to-end integration: memory tools running against real LanceDB. + +These tests load a sentence-transformer model (~30-40s startup) and are +marked ``slow`` so they don't run in the default ``pytest`` invocation. +Run them explicitly with ``pytest -m slow``. +""" + +from __future__ import annotations + +import json +import shutil + +import pytest +from wren.memory.store import MemoryStore + +from wren_langchain import WrenToolkit + +pytestmark = pytest.mark.slow + + +@pytest.fixture +def project_with_memory(duckdb_project): + """Augment the duckdb_project fixture with an indexed .wren/memory dir.""" + memory_dir = duckdb_project / ".wren" / "memory" + memory_dir.mkdir(parents=True) + + # Eagerly create the LanceDB tables by indexing the (small) manifest so + # subsequent fetch/recall calls don't crash on first read. + manifest = json.loads((duckdb_project / "target" / "mdl.json").read_text()) + store = MemoryStore(path=memory_dir) + store.index_schema(manifest, replace=True, seed_queries=False) + + yield duckdb_project + + # Best-effort cleanup; tmp_path is auto-cleaned but LanceDB may leave open + # handles on Windows. On macOS/Linux this is a no-op safety net. + shutil.rmtree(memory_dir, ignore_errors=True) + + +def test_fetch_context_runs_against_real_lancedb(project_with_memory): + toolkit = WrenToolkit.from_project(project_with_memory) + fetch = next(t for t in toolkit.get_tools() if t.name == "wren_fetch_context") + + envelope = fetch.invoke({"question": "customers"}) + + assert envelope["ok"] is True + assert envelope["data"]["strategy"] in {"full", "search"} + + +def test_store_then_recall_round_trip(project_with_memory): + toolkit = WrenToolkit.from_project(project_with_memory) + store = next(t for t in toolkit.get_tools() if t.name == "wren_store_query") + recall = next(t for t in toolkit.get_tools() if t.name == "wren_recall_queries") + + store_env = store.invoke( + { + "nl": "list all customers", + "sql": "SELECT id, name FROM customers", + "tags": ["demo"], + } + ) + assert store_env["ok"] is True + + recall_env = recall.invoke({"question": "customers list"}) + assert recall_env["ok"] is True + assert len(recall_env["data"]["results"]) >= 1 + nl_values = [r.get("nl_query") for r in recall_env["data"]["results"]] + assert "list all customers" in nl_values diff --git a/sdk/wren-langchain/tests/integration/test_runtime_tools.py b/sdk/wren-langchain/tests/integration/test_runtime_tools.py new file mode 100644 index 000000000..6789dd9aa --- /dev/null +++ b/sdk/wren-langchain/tests/integration/test_runtime_tools.py @@ -0,0 +1,53 @@ +"""End-to-end integration: WrenToolkit running real queries against DuckDB.""" + +from __future__ import annotations + +from wren_langchain import WrenToolkit + + +def test_query_executes_against_duckdb_returns_arrow_table(duckdb_project): + """toolkit.query against a real DuckDB-backed project returns rows.""" + toolkit = WrenToolkit.from_project(duckdb_project) + + table = toolkit.query("SELECT id, name FROM customers ORDER BY id") + + assert table.column_names == ["id", "name"] + assert table.to_pylist() == [ + {"id": 1, "name": "Acme"}, + {"id": 2, "name": "Globex"}, + ] + + +def test_wren_query_tool_returns_envelope_with_real_rows(duckdb_project): + """The LangChain wren_query tool returns a full envelope from a real query.""" + toolkit = WrenToolkit.from_project(duckdb_project) + wren_query = next(t for t in toolkit.get_tools() if t.name == "wren_query") + + envelope = wren_query.invoke({"sql": "SELECT id, name FROM customers ORDER BY id"}) + + assert envelope["ok"] is True + assert envelope["data"]["row_count"] == 2 + assert envelope["data"]["rows"][0]["name"] == "Acme" + + +def test_wren_list_models_tool_renders_manifest(duckdb_project): + toolkit = WrenToolkit.from_project(duckdb_project) + list_models = next(t for t in toolkit.get_tools() if t.name == "wren_list_models") + + envelope = list_models.invoke({}) + + assert envelope["ok"] is True + assert "customers" in envelope["content"] + assert envelope["data"]["models"][0]["name"] == "customers" + + +def test_connector_reused_across_queries(duckdb_project): + """After the first query, ``_connector_cache`` is populated and reused.""" + toolkit = WrenToolkit.from_project(duckdb_project) + + toolkit.query("SELECT 1") + first_connector = toolkit._connector_cache + assert first_connector is not None + + toolkit.query("SELECT 2") + assert toolkit._connector_cache is first_connector diff --git a/sdk/wren-langchain/tests/unit/__init__.py b/sdk/wren-langchain/tests/unit/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/sdk/wren-langchain/tests/unit/test_envelope.py b/sdk/wren-langchain/tests/unit/test_envelope.py new file mode 100644 index 000000000..8585b44bb --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_envelope.py @@ -0,0 +1,187 @@ +"""Tests for envelope construction and error formatting.""" + +import datetime as dt +from decimal import Decimal + +from wren.model.error import ErrorCode, ErrorPhase, WrenError + +from wren_langchain._envelope import ( + cap_size, + format_error, + json_safe, + make_error, + make_success, +) + + +def test_make_success_returns_ok_envelope(): + """make_success returns a JSON-serializable envelope with ok=True.""" + result = make_success(content="hello", data={"foo": "bar"}) + + assert result == { + "ok": True, + "content": "hello", + "data": {"foo": "bar"}, + "warnings": [], + } + + +def test_make_success_includes_warnings_when_provided(): + """Warnings list is preserved in the envelope.""" + result = make_success( + content="ok", + data={}, + warnings=["content truncated: showed 32 of 100 rows"], + ) + + assert result["warnings"] == ["content truncated: showed 32 of 100 rows"] + + +def test_format_error_extracts_wren_error_fields(): + """format_error pulls code, phase, message, metadata from WrenError.""" + exc = WrenError( + error_code=ErrorCode.INVALID_SQL, + message="syntax error near 'SELEC'", + phase=ErrorPhase.SQL_PARSING, + metadata={"position": 7}, + ) + + result = format_error(exc) + + assert result["code"] == "INVALID_SQL" + assert result["phase"] == "SQL_PARSING" + assert result["message"] == "syntax error near 'SELEC'" + assert result["metadata"] == {"position": 7} + + +def test_format_error_handles_non_wren_exception(): + """Generic exceptions become SDK_ERROR with stringified message.""" + result = format_error(ValueError("something went wrong")) + + assert result["code"] == "SDK_ERROR" + assert result["phase"] is None + assert result["message"] == "something went wrong" + assert result["metadata"] == {} + + +def test_format_error_redacts_nested_secret_keys_in_metadata(): + """Nested secrets inside dicts and lists must also be redacted.""" + exc = WrenError( + error_code=ErrorCode.GET_CONNECTION_ERROR, + message="connect failed", + metadata={ + "connection_info": { + "host": "db.example.com", + "password": "hunter2", + "credentials": {"token": "deeply-nested"}, + }, + "history": [ + {"event": "connect", "auth_token": "should-be-hidden"}, + ], + }, + ) + + result = format_error(exc) + md = result["metadata"] + + assert md["connection_info"]["host"] == "db.example.com" + assert md["connection_info"]["password"] == "***" + assert md["connection_info"]["credentials"] == "***" + assert md["history"][0]["event"] == "connect" + assert md["history"][0]["auth_token"] == "***" + + +def test_format_error_redacts_secret_keys_in_metadata(): + """Keys matching secret patterns get replaced with '***'.""" + exc = WrenError( + error_code=ErrorCode.GET_CONNECTION_ERROR, + message="connect failed", + metadata={ + "host": "db.example.com", + "password": "hunter2", + "API_TOKEN": "abc123", + "auth_secret": "shh", + "user_credential": "pwd", + "harmless": "value", + }, + ) + + result = format_error(exc) + + assert result["metadata"]["host"] == "db.example.com" + assert result["metadata"]["harmless"] == "value" + assert result["metadata"]["password"] == "***" + assert result["metadata"]["API_TOKEN"] == "***" + assert result["metadata"]["auth_secret"] == "***" + assert result["metadata"]["user_credential"] == "***" + + +def test_json_safe_converts_datetime_to_iso_string(): + """datetime objects are serialized to ISO strings.""" + result = json_safe({"when": dt.datetime(2026, 5, 6, 12, 0, 0)}) + assert result["when"] == "2026-05-06T12:00:00" + + +def test_json_safe_converts_decimal_to_string(): + """Decimal objects are serialized to strings to preserve precision.""" + result = json_safe({"amount": Decimal("123.45")}) + assert result["amount"] == "123.45" + + +def test_json_safe_recurses_into_nested_dict_and_list(): + """Conversion recurses through nested structures.""" + payload = { + "rows": [ + {"date": dt.date(2026, 5, 6), "amount": Decimal("1.5")}, + {"date": dt.date(2026, 5, 7), "amount": Decimal("2.5")}, + ], + } + result = json_safe(payload) + assert result["rows"][0]["date"] == "2026-05-06" + assert result["rows"][0]["amount"] == "1.5" + assert result["rows"][1]["amount"] == "2.5" + + +def test_cap_size_returns_input_unchanged_when_under_limit(): + """cap_size leaves data alone when JSON-serialized size is below limit.""" + payload = {"a": 1, "b": "hello"} + result = cap_size(payload, max_bytes=4096) + assert result == payload + + +def test_cap_size_truncates_and_marks_when_over_limit(): + """When over-limit, cap_size returns a marker dict with the original size note.""" + payload = {"big": "x" * 10000} + result = cap_size(payload, max_bytes=1024) + assert result["_truncated"] is True + assert "original_size_bytes" in result + + +def test_format_error_caps_metadata_at_4kb(): + """Metadata exceeding 4KB is replaced by a truncation marker.""" + big_metadata = {"sql": "SELECT * FROM x WHERE y = '" + ("a" * 5000) + "'"} + exc = WrenError( + error_code=ErrorCode.INVALID_SQL, + message="bad sql", + metadata=big_metadata, + ) + + result = format_error(exc) + + assert result["metadata"].get("_truncated") is True + + +def test_make_error_returns_full_envelope(): + """make_error returns a full {ok: False, content, error} envelope.""" + exc = WrenError( + error_code=ErrorCode.INVALID_SQL, + message="syntax error", + phase=ErrorPhase.SQL_PARSING, + ) + + result = make_error(exc) + + assert result["ok"] is False + assert result["content"] == "syntax error" + assert result["error"]["code"] == "INVALID_SQL" + assert result["error"]["phase"] == "SQL_PARSING" diff --git a/sdk/wren-langchain/tests/unit/test_exceptions.py b/sdk/wren-langchain/tests/unit/test_exceptions.py new file mode 100644 index 000000000..a8ecfdf92 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_exceptions.py @@ -0,0 +1,21 @@ +"""Tests for SDK-specific exception types.""" + +import pytest + +from wren_langchain.exceptions import ( + MemoryNotEnabledError, + WrenToolkitInitError, +) + + +def test_wren_toolkit_init_error_carries_message(): + with pytest.raises(WrenToolkitInitError, match="missing target/mdl.json"): + raise WrenToolkitInitError("missing target/mdl.json") + + +def test_memory_not_enabled_error_is_distinct_type(): + """MemoryNotEnabledError is its own class, not a generic ValueError.""" + err = MemoryNotEnabledError("memory provider not configured") + assert isinstance(err, MemoryNotEnabledError) + assert isinstance(err, Exception) + assert "memory provider" in str(err) diff --git a/sdk/wren-langchain/tests/unit/test_memory_api.py b/sdk/wren-langchain/tests/unit/test_memory_api.py new file mode 100644 index 000000000..24f88a307 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_memory_api.py @@ -0,0 +1,112 @@ +"""Tests for the _MemoryAPI subscope (toolkit.memory.*).""" + +from unittest.mock import MagicMock, patch + +import pytest + +from wren_langchain import WrenToolkit +from wren_langchain.exceptions import MemoryNotEnabledError + + +def _enable_memory(tmp_project): + """Helper: create .wren/memory directory so memory auto-enables.""" + (tmp_project / ".wren" / "memory").mkdir(parents=True) + return tmp_project + + +def test_memory_fetch_calls_get_context_with_manifest(tmp_project, fake_active_profile): + """toolkit.memory.fetch passes the loaded manifest to MemoryStore.get_context.""" + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + fake_store.get_context.return_value = {"strategy": "search", "results": []} + + toolkit = WrenToolkit.from_project(project) + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + result = toolkit.memory.fetch("revenue trends", limit=3) + + assert result == {"strategy": "search", "results": []} + fake_store.get_context.assert_called_once() + kwargs = fake_store.get_context.call_args.kwargs + assert kwargs["query"] == "revenue trends" + assert kwargs["limit"] == 3 + # manifest is loaded read-through and passed through. + assert "manifest" in kwargs + + +def test_memory_recall_calls_recall_queries(tmp_project, fake_active_profile): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + fake_store.recall_queries.return_value = [{"nl": "x", "sql": "SELECT 1"}] + + toolkit = WrenToolkit.from_project(project) + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + result = toolkit.memory.recall("top customers", limit=5) + + assert result == [{"nl": "x", "sql": "SELECT 1"}] + fake_store.recall_queries.assert_called_once_with(query="top customers", limit=5) + + +def test_memory_store_calls_store_query(tmp_project, fake_active_profile): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + + toolkit = WrenToolkit.from_project(project) + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit.memory.store( + nl="top customers", + sql="SELECT * FROM customers ORDER BY revenue DESC LIMIT 10", + tags=["revenue", "ranking"], + ) + + fake_store.store_query.assert_called_once() + kwargs = fake_store.store_query.call_args.kwargs + assert kwargs["nl_query"] == "top customers" + assert kwargs["sql_query"].startswith("SELECT") + # SDK joins list[str] tags into a Core-compatible comma-separated string. + assert kwargs["tags"] == "revenue,ranking" + + +def test_memory_fetch_raises_when_memory_disabled(tmp_project, fake_active_profile): + """Direct API access when memory is disabled raises MemoryNotEnabledError.""" + toolkit = WrenToolkit.from_project(tmp_project) + + with pytest.raises(MemoryNotEnabledError): + toolkit.memory.fetch("anything") + + +def test_memory_store_rejects_tags_containing_commas(tmp_project, fake_active_profile): + """Commas separate tags in the underlying storage format. A tag like + "revenue, Q1" would silently corrupt the round-trip if passed through — + we reject early with ValueError instead.""" + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + + toolkit = WrenToolkit.from_project(project) + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + with pytest.raises(ValueError, match="comma"): + toolkit.memory.store(nl="x", sql="SELECT 1", tags=["revenue, Q1"]) + + fake_store.store_query.assert_not_called() + + +def test_memory_store_caches_across_calls(tmp_project, fake_active_profile): + """The MemoryStore instance is constructed once and reused across operations.""" + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + fake_store.get_context.return_value = {} + fake_store.recall_queries.return_value = [] + + toolkit = WrenToolkit.from_project(project) + + with patch( + "wren_langchain._providers.memory.MemoryStore", return_value=fake_store + ) as ctor: + toolkit.memory.fetch("x") + toolkit.memory.recall("y") + toolkit.memory.fetch("z") + + assert ctor.call_count == 1 # constructed exactly once diff --git a/sdk/wren-langchain/tests/unit/test_prompt.py b/sdk/wren-langchain/tests/unit/test_prompt.py new file mode 100644 index 000000000..79609e5ff --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_prompt.py @@ -0,0 +1,154 @@ +"""Tests for the system_prompt builder.""" + +from unittest.mock import MagicMock, patch + +from wren_langchain import WrenToolkit + + +def _enable_memory(tmp_project): + (tmp_project / ".wren" / "memory").mkdir(parents=True) + return tmp_project + + +def test_system_prompt_returns_str(tmp_project, fake_active_profile): + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + assert isinstance(prompt, str) + assert len(prompt) > 0 + + +def test_system_prompt_includes_workflow_section(tmp_project, fake_active_profile): + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + assert "Wren" in prompt + # Workflow rule about Wren model names being preferred over raw tables. + assert "model" in prompt.lower() + + +def test_system_prompt_lists_enabled_tools_in_summary(tmp_project, fake_active_profile): + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + assert "wren_query" in prompt + assert "wren_dry_plan" in prompt + assert "wren_list_models" in prompt + + +def test_system_prompt_omits_memory_tools_when_disabled( + tmp_project, fake_active_profile +): + """When memory is off, prompt should not reference memory tools.""" + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + assert "wren_fetch_context" not in prompt + assert "wren_recall_queries" not in prompt + assert "wren_store_query" not in prompt + + +def test_system_prompt_includes_memory_tools_when_enabled( + tmp_project, fake_active_profile +): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + prompt = toolkit.system_prompt() + + assert "wren_fetch_context" in prompt + assert "wren_recall_queries" in prompt + assert "wren_store_query" in prompt + + +def test_system_prompt_appends_project_instructions_when_present( + tmp_project, fake_active_profile +): + (tmp_project / "instructions.md").write_text( + "# Domain\n\nThis project tracks B2B SaaS revenue.\n" + ) + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + + assert "B2B SaaS revenue" in prompt + assert "Project-specific instructions" in prompt + + +def test_system_prompt_silently_skips_instructions_when_absent( + tmp_project, fake_active_profile +): + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + # No "Project-specific instructions" section header should appear. + assert "Project-specific instructions" not in prompt + + +def test_memory_workflow_uses_strong_default_language(tmp_project, fake_active_profile): + """Memory-enabled prompt must use 'by default'/'only when' phrasing, + not hedge words like 'non-trivial' or 'useful', because empirical testing + showed soft phrasing causes GPT-4o to skip recall/store reliably.""" + project = tmp_project + (project / ".wren" / "memory").mkdir(parents=True) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + prompt = toolkit.system_prompt() + + # Strong-default phrasing must appear. + assert "by default" in prompt.lower() + assert "only when" in prompt.lower() + + # Hedges that we explicitly removed must NOT appear. + assert "non-trivial" not in prompt.lower() + assert "if helpful" not in prompt.lower() + assert "useful" not in prompt.lower() + + +def test_error_phase_guidance_present_in_prompt(tmp_project, fake_active_profile): + """The prompt must instruct the agent how to react to ok=false envelopes + by phase, so it can fix-and-retry instead of silently abandoning.""" + toolkit = WrenToolkit.from_project(tmp_project) + prompt = toolkit.system_prompt() + + assert "SQL_PARSING" in prompt + assert "SQL_EXECUTION" in prompt + + +def test_system_prompt_respects_include_memory_write_false( + tmp_project, fake_active_profile +): + """When the caller passes a tool list with `wren_store_query` filtered out, + the workflow must drop the persistence step and the tools section must not + list it. Otherwise the prompt would tell the LLM to call a tool the agent + doesn't actually have.""" + project = tmp_project + (project / ".wren" / "memory").mkdir(parents=True) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + tools_no_write = toolkit.get_tools(include_memory_write=False) + prompt = toolkit.system_prompt(tools=tools_no_write) + + # Read tools (fetch + recall) still mentioned. + assert "wren_fetch_context" in prompt + assert "wren_recall_queries" in prompt + # Write tool dropped both from workflow steps and tools listing. + assert "wren_store_query" not in prompt + assert "Persist the NL→SQL pair" not in prompt + + +def test_system_prompt_default_uses_full_tool_set(tmp_project, fake_active_profile): + """Without an explicit tools= override, the prompt mirrors get_tools() + defaults — full memory workflow when memory is enabled.""" + project = tmp_project + (project / ".wren" / "memory").mkdir(parents=True) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + prompt = toolkit.system_prompt() + + # Default get_tools() includes all 6, so the workflow has all 3 memory steps. + assert "wren_fetch_context" in prompt + assert "wren_recall_queries" in prompt + assert "wren_store_query" in prompt diff --git a/sdk/wren-langchain/tests/unit/test_providers_connection.py b/sdk/wren-langchain/tests/unit/test_providers_connection.py new file mode 100644 index 000000000..ca373d696 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_providers_connection.py @@ -0,0 +1,84 @@ +"""Tests for ConnectionProvider implementations.""" + +import pytest + +from wren_langchain._providers.connection import ProfileConnectionProvider +from wren_langchain.exceptions import WrenToolkitInitError + + +def test_explicit_profile_kwarg_resolves_first(monkeypatch, tmp_path): + """Layer 1: explicit profile= kwarg wins over project config and active.""" + fake_profiles = { + "prod": {"datasource": "postgres", "host": "prod.db", "port": 5432}, + "dev": {"datasource": "duckdb", "path": ":memory:"}, + } + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: fake_profiles, + ) + monkeypatch.setattr( + "wren_langchain._providers.connection.get_active_profile", + lambda: ("dev", fake_profiles["dev"]), + ) + + provider = ProfileConnectionProvider( + project_path=tmp_path, + explicit_profile="prod", + ) + + assert provider.datasource() == "postgres" + assert provider.connection_info() == {"host": "prod.db", "port": 5432} + + +def test_project_config_profile_field_resolves_second(monkeypatch, tmp_path): + """Layer 2: wren_project.yml's `profile:` field used when no explicit kwarg.""" + fake_profiles = { + "from_project": {"datasource": "mysql", "host": "from-project.db"}, + "active": {"datasource": "duckdb", "path": ":memory:"}, + } + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: fake_profiles, + ) + monkeypatch.setattr( + "wren_langchain._providers.connection.get_active_profile", + lambda: ("active", fake_profiles["active"]), + ) + (tmp_path / "wren_project.yml").write_text("profile: from_project\n") + + provider = ProfileConnectionProvider(project_path=tmp_path) + + assert provider.datasource() == "mysql" + assert provider.connection_info() == {"host": "from-project.db"} + + +def test_active_profile_resolves_third_when_no_explicit_or_project( + monkeypatch, tmp_path +): + """Layer 3: globally active profile is used when no explicit and no project config.""" + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: {"only": {"datasource": "snowflake"}}, + ) + monkeypatch.setattr( + "wren_langchain._providers.connection.get_active_profile", + lambda: ("only", {"datasource": "snowflake", "account": "abc"}), + ) + + provider = ProfileConnectionProvider(project_path=tmp_path) + + assert provider.datasource() == "snowflake" + assert provider.connection_info() == {"account": "abc"} + + +def test_unknown_profile_name_raises(monkeypatch, tmp_path): + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: {"prod": {"datasource": "postgres", "host": "x"}}, + ) + + with pytest.raises(WrenToolkitInitError, match="profile.*not found"): + ProfileConnectionProvider( + project_path=tmp_path, + explicit_profile="nonexistent", + ) diff --git a/sdk/wren-langchain/tests/unit/test_providers_mdl.py b/sdk/wren-langchain/tests/unit/test_providers_mdl.py new file mode 100644 index 000000000..e4a258593 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_providers_mdl.py @@ -0,0 +1,58 @@ +"""Tests for MDLSource implementations.""" + +import json + +import pytest + +from wren_langchain._providers.mdl_source import ProjectMDLSource +from wren_langchain.exceptions import WrenToolkitInitError + + +def test_project_mdl_source_reads_target_mdl_json(tmp_path): + """ProjectMDLSource reads the project's target/mdl.json on every load.""" + target = tmp_path / "target" + target.mkdir() + manifest = {"models": [{"name": "orders"}]} + (target / "mdl.json").write_text(json.dumps(manifest)) + + source = ProjectMDLSource(project_path=tmp_path) + + assert source.load_manifest() == manifest + + +def test_project_mdl_source_picks_up_file_changes_between_calls(tmp_path): + """Subsequent load_manifest() calls reflect on-disk changes (read-through).""" + target = tmp_path / "target" + target.mkdir() + mdl_file = target / "mdl.json" + mdl_file.write_text(json.dumps({"models": [{"name": "v1"}]})) + + source = ProjectMDLSource(project_path=tmp_path) + first = source.load_manifest() + + mdl_file.write_text(json.dumps({"models": [{"name": "v2"}]})) + second = source.load_manifest() + + assert first["models"][0]["name"] == "v1" + assert second["models"][0]["name"] == "v2" + + +def test_project_mdl_source_raises_on_missing_target(tmp_path): + """A missing target/mdl.json raises WrenToolkitInitError when loading.""" + source = ProjectMDLSource(project_path=tmp_path) + + with pytest.raises(WrenToolkitInitError, match="target/mdl.json"): + source.load_manifest() + + +def test_project_mdl_source_normalizes_malformed_json_to_init_error(tmp_path): + """Malformed mdl.json must surface as WrenToolkitInitError, not raw JSONDecodeError, + so callers don't need to special-case JSON internals.""" + target = tmp_path / "target" + target.mkdir() + (target / "mdl.json").write_text("{not valid json") + + source = ProjectMDLSource(project_path=tmp_path) + + with pytest.raises(WrenToolkitInitError, match="not valid JSON"): + source.load_manifest() diff --git a/sdk/wren-langchain/tests/unit/test_providers_memory.py b/sdk/wren-langchain/tests/unit/test_providers_memory.py new file mode 100644 index 000000000..e838b667d --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_providers_memory.py @@ -0,0 +1,43 @@ +"""Tests for MemoryProvider implementations.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from wren_langchain._providers.memory import ( + LocalLanceDBMemoryProvider, + NoopMemoryProvider, +) +from wren_langchain.exceptions import MemoryNotEnabledError + + +def test_noop_memory_provider_is_disabled(): + """NoopMemoryProvider reports as disabled and raises on open().""" + provider = NoopMemoryProvider() + + assert provider.enabled is False + + with pytest.raises(MemoryNotEnabledError): + provider.open() + + +def test_local_lancedb_provider_is_enabled(tmp_path): + """A local provider is enabled (regardless of whether the dir exists yet).""" + provider = LocalLanceDBMemoryProvider(memory_path=tmp_path / ".wren" / "memory") + assert provider.enabled is True + + +def test_local_lancedb_provider_open_constructs_memory_store(tmp_path): + """open() lazily constructs a wren.memory.MemoryStore at the given path.""" + memory_path = tmp_path / ".wren" / "memory" + provider = LocalLanceDBMemoryProvider(memory_path=memory_path) + + fake_store = MagicMock(name="MemoryStore") + with patch( + "wren_langchain._providers.memory.MemoryStore", + return_value=fake_store, + ) as ctor: + store = provider.open() + + assert store is fake_store + ctor.assert_called_once_with(path=memory_path) diff --git a/sdk/wren-langchain/tests/unit/test_toolkit_init.py b/sdk/wren-langchain/tests/unit/test_toolkit_init.py new file mode 100644 index 000000000..41b2beb95 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_toolkit_init.py @@ -0,0 +1,104 @@ +"""Tests for WrenToolkit construction (from_project).""" + +import pytest +from wren.profile import _reset_env_loaded_for_tests + +from wren_langchain import WrenToolkit +from wren_langchain._providers.memory import ( + LocalLanceDBMemoryProvider, + NoopMemoryProvider, +) +from wren_langchain.exceptions import WrenToolkitInitError + + +def test_from_project_raises_when_project_yml_missing(tmp_path, fake_active_profile): + """A directory without wren_project.yml is not a Wren project.""" + with pytest.raises(WrenToolkitInitError, match="wren_project.yml"): + WrenToolkit.from_project(tmp_path) + + +def test_from_project_raises_when_target_mdl_missing(tmp_path, fake_active_profile): + """A project without target/mdl.json hasn't been built.""" + (tmp_path / "wren_project.yml").write_text("schema_version: 1\n") + with pytest.raises(WrenToolkitInitError, match="target/mdl.json"): + WrenToolkit.from_project(tmp_path) + + +def test_from_project_returns_toolkit_when_prereqs_met( + tmp_project, fake_active_profile +): + """from_project returns a WrenToolkit when all prerequisites exist.""" + toolkit = WrenToolkit.from_project(tmp_project) + assert isinstance(toolkit, WrenToolkit) + + +def test_from_project_relative_path_resolves( + tmp_project, fake_active_profile, monkeypatch +): + """from_project accepts relative paths and resolves them.""" + monkeypatch.chdir(tmp_project.parent) + toolkit = WrenToolkit.from_project(tmp_project.name) + assert isinstance(toolkit, WrenToolkit) + + +def test_memory_auto_detect_disabled_when_dir_missing(tmp_project, fake_active_profile): + """Without .wren/memory/, memory auto-detects as Noop.""" + toolkit = WrenToolkit.from_project(tmp_project) + assert isinstance(toolkit._memory, NoopMemoryProvider) + + +def test_memory_auto_detect_enabled_when_dir_exists(tmp_project, fake_active_profile): + """With .wren/memory/, memory auto-detects as LocalLanceDB.""" + (tmp_project / ".wren" / "memory").mkdir(parents=True) + toolkit = WrenToolkit.from_project(tmp_project) + assert isinstance(toolkit._memory, LocalLanceDBMemoryProvider) + + +def test_from_project_loads_dotenv_from_project_path(tmp_project, monkeypatch): + """from_project loads /.env so ${VAR} secrets resolve regardless of CWD. + + Regression: previously the SDK relied on Core's CWD-relative .env discovery, + which fails when the user runs Python from anywhere other than the project + directory. + """ + # Stage 1: a profile that references an env var the caller's shell does NOT have. + sentinel_var = "WREN_LANGCHAIN_TEST_HOST_DOES_NOT_EXIST_IN_SHELL" + monkeypatch.delenv(sentinel_var, raising=False) + monkeypatch.setattr( + "wren_langchain._providers.connection.list_profiles", + lambda: { + "test": { + "datasource": "duckdb", + "host": f"${{{sentinel_var}}}", + "format": "duckdb", + } + }, + ) + monkeypatch.setattr( + "wren_langchain._providers.connection.get_active_profile", + lambda: ( + "test", + { + "datasource": "duckdb", + "host": f"${{{sentinel_var}}}", + "format": "duckdb", + }, + ), + ) + + # Stage 2: place the var only inside the project's .env. + (tmp_project / ".env").write_text(f"{sentinel_var}=resolved-from-project-env\n") + + # Stage 3: run from a different CWD so Core's CWD-walk would NOT find the file. + monkeypatch.chdir(tmp_project.parent) + _reset_env_loaded_for_tests() + + try: + toolkit = WrenToolkit.from_project(tmp_project) + assert ( + toolkit._connection.connection_info()["host"] == "resolved-from-project-env" + ) + finally: + # Reset the global loader flag again so this test cannot leak its + # half-loaded state into whatever runs next in the session. + _reset_env_loaded_for_tests() diff --git a/sdk/wren-langchain/tests/unit/test_toolkit_runtime.py b/sdk/wren-langchain/tests/unit/test_toolkit_runtime.py new file mode 100644 index 000000000..edfc24426 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_toolkit_runtime.py @@ -0,0 +1,121 @@ +"""Tests for WrenToolkit runtime API: query, dry_plan, dry_run.""" + +import base64 +import json +from unittest.mock import MagicMock, patch + +import pyarrow as pa + +from wren_langchain import WrenToolkit + + +def test_query_invokes_wren_engine_with_resolved_manifest( + tmp_project, fake_active_profile +): + """toolkit.query reads the manifest fresh and delegates to WrenEngine.query.""" + fake_table = pa.table({"x": [1, 2, 3]}) + fake_engine = MagicMock(name="WrenEngine") + fake_engine.query.return_value = fake_table + fake_engine._connector = MagicMock(name="connector") + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch( + "wren_langchain._toolkit.WrenEngine", return_value=fake_engine + ) as engine_ctor: + result = toolkit.query("SELECT 1", limit=10) + + assert result is fake_table + fake_engine.query.assert_called_once_with("SELECT 1", limit=10) + # Engine constructed with manifest bytes + datasource + connection_info + engine_ctor.assert_called_once() + kwargs = engine_ctor.call_args.kwargs + assert kwargs["data_source"] == "duckdb" + assert kwargs["connection_info"] == {"path": ":memory:"} + + +def test_connector_is_reused_across_query_calls(tmp_project, fake_active_profile): + """Second query reuses the cached connector instead of reconnecting. + + Distinguishes "what engine 2 starts with" from "what gets injected" by + seeding each fresh engine with a different connector. If reuse is broken, + engine 2's `_connector` would remain its own initial mock; if reuse works, + it gets replaced with engine 1's connector before `query()` runs. + """ + first_connector = MagicMock(name="first_connector") + second_initial_connector = MagicMock(name="second_initial_connector") + engines = [] + + def make_engine(*args, **kwargs): + engine = MagicMock(name=f"engine{len(engines)}") + engine._connector = first_connector if not engines else second_initial_connector + engines.append(engine) + return engine + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch("wren_langchain._toolkit.WrenEngine", side_effect=make_engine): + toolkit.query("SELECT 1") + toolkit.query("SELECT 2") + + # Engine 1 keeps its connector; engine 2's initial connector got + # overwritten with the cached one from engine 1 before query() ran. + assert engines[0]._connector is first_connector + assert engines[1]._connector is first_connector + assert engines[1]._connector is not second_initial_connector + + +def test_manifest_is_read_through_on_every_call( + tmp_project, fake_active_profile, monkeypatch +): + """Each query re-reads target/mdl.json so external CLI rebuilds are picked up.""" + fake_engine = MagicMock(name="engine") + fake_engine._connector = MagicMock() + + toolkit = WrenToolkit.from_project(tmp_project) + + # Replace the manifest content between calls. + mdl_path = tmp_project / "target" / "mdl.json" + mdl_path.write_text('{"models": [{"name": "v1"}]}') + + with patch( + "wren_langchain._toolkit.WrenEngine", return_value=fake_engine + ) as engine_ctor: + toolkit.query("SELECT 1") + + # Simulate `wren context build` updating the file. + mdl_path.write_text('{"models": [{"name": "v2"}]}') + toolkit.query("SELECT 2") + + first_manifest_b64 = engine_ctor.call_args_list[0].kwargs["manifest_str"] + second_manifest_b64 = engine_ctor.call_args_list[1].kwargs["manifest_str"] + first = json.loads(base64.b64decode(first_manifest_b64)) + second = json.loads(base64.b64decode(second_manifest_b64)) + assert first["models"][0]["name"] == "v1" + assert second["models"][0]["name"] == "v2" + + +def test_dry_plan_delegates_to_engine(tmp_project, fake_active_profile): + fake_engine = MagicMock(name="engine") + fake_engine.dry_plan.return_value = "SELECT * FROM cte_orders" + fake_engine._connector = MagicMock() + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch("wren_langchain._toolkit.WrenEngine", return_value=fake_engine): + result = toolkit.dry_plan("SELECT * FROM orders") + + assert result == "SELECT * FROM cte_orders" + fake_engine.dry_plan.assert_called_once_with("SELECT * FROM orders") + + +def test_dry_run_delegates_to_engine(tmp_project, fake_active_profile): + fake_engine = MagicMock(name="engine") + fake_engine._connector = MagicMock() + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch("wren_langchain._toolkit.WrenEngine", return_value=fake_engine): + toolkit.dry_run("SELECT 1") + + fake_engine.dry_run.assert_called_once_with("SELECT 1") diff --git a/sdk/wren-langchain/tests/unit/test_tools_memory.py b/sdk/wren-langchain/tests/unit/test_tools_memory.py new file mode 100644 index 000000000..1a434ca9a --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_tools_memory.py @@ -0,0 +1,159 @@ +"""Tests for memory LLM-facing tools.""" + +from unittest.mock import MagicMock, patch + +from wren_langchain import WrenToolkit + + +def _enable_memory(tmp_project): + (tmp_project / ".wren" / "memory").mkdir(parents=True) + return tmp_project + + +def test_get_tools_returns_six_tools_when_memory_enabled( + tmp_project, fake_active_profile +): + """Memory enabled → 3 runtime + 3 memory tools.""" + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + names = sorted(t.name for t in toolkit.get_tools()) + + assert names == [ + "wren_dry_plan", + "wren_fetch_context", + "wren_list_models", + "wren_query", + "wren_recall_queries", + "wren_store_query", + ] + + +def test_include_memory_write_false_removes_store_query( + tmp_project, fake_active_profile +): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + names = sorted(t.name for t in toolkit.get_tools(include_memory_write=False)) + + assert "wren_store_query" not in names + assert "wren_fetch_context" in names + assert "wren_recall_queries" in names + + +def test_include_memory_write_true_is_no_op_when_memory_disabled( + tmp_project, fake_active_profile +): + """When the project has no .wren/memory/, include_memory_write=True must + silently produce no memory tools — not raise, not warn, not partially + add tools that would fail on first call. + """ + # tmp_project fixture does NOT create .wren/memory/, so memory auto-detects + # as disabled. + toolkit = WrenToolkit.from_project(tmp_project) + + tools_default = toolkit.get_tools() + tools_explicit_true = toolkit.get_tools(include_memory_write=True) + + # Memory is disabled either way; no memory tools regardless of include flag. + for tools in (tools_default, tools_explicit_true): + names = {t.name for t in tools} + assert "wren_store_query" not in names + assert "wren_fetch_context" not in names + assert "wren_recall_queries" not in names + # Runtime tools are still present. + assert "wren_query" in names + + +def test_wren_fetch_context_full_strategy(tmp_project, fake_active_profile): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + fake_store.get_context.return_value = { + "strategy": "full", + "schema": "Schema text describing models...", + } + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + tool = next(t for t in toolkit.get_tools() if t.name == "wren_fetch_context") + envelope = tool.invoke({"question": "what models exist?"}) + + assert envelope["ok"] is True + assert envelope["data"]["strategy"] == "full" + assert "Schema text" in envelope["content"] + + +def test_wren_fetch_context_search_strategy(tmp_project, fake_active_profile): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + fake_store.get_context.return_value = { + "strategy": "search", + "results": [ + {"item_type": "model", "name": "orders", "summary": "orders model"}, + {"item_type": "column", "name": "orders.id", "summary": "primary key"}, + ], + } + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + tool = next(t for t in toolkit.get_tools() if t.name == "wren_fetch_context") + envelope = tool.invoke({"question": "orders"}) + + assert envelope["ok"] is True + assert envelope["data"]["strategy"] == "search" + assert "[model] orders" in envelope["content"] + assert "[column] orders.id" in envelope["content"] + + +def test_wren_recall_queries_formats_as_numbered_list(tmp_project, fake_active_profile): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + fake_store.recall_queries.return_value = [ + {"nl_query": "top customers", "sql_query": "SELECT * FROM customers"}, + {"nl_query": "revenue by region", "sql_query": "SELECT region, SUM(revenue)"}, + ] + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + tool = next(t for t in toolkit.get_tools() if t.name == "wren_recall_queries") + envelope = tool.invoke({"question": "customer rankings"}) + + assert envelope["ok"] is True + assert "1." in envelope["content"] + assert "top customers" in envelope["content"] + assert "```sql" in envelope["content"] + assert len(envelope["data"]["results"]) == 2 + + +def test_wren_store_query_returns_short_success_message( + tmp_project, fake_active_profile +): + project = _enable_memory(tmp_project) + fake_store = MagicMock(name="MemoryStore") + + with patch("wren_langchain._providers.memory.MemoryStore", return_value=fake_store): + toolkit = WrenToolkit.from_project(project) + tool = next(t for t in toolkit.get_tools() if t.name == "wren_store_query") + envelope = tool.invoke( + { + "nl": "top customers", + "sql": "SELECT * FROM customers", + "tags": ["ranking", "demo"], + } + ) + + assert envelope["ok"] is True + assert "Stored" in envelope["content"] + # Use `assert_called_once_with(...)` so this test fails with a clear + # diff if the SDK ever switches from kwargs to positional args, instead + # of a confusing KeyError on `call_args.kwargs["tags"]`. + fake_store.store_query.assert_called_once_with( + nl_query="top customers", + sql_query="SELECT * FROM customers", + tags="ranking,demo", + ) diff --git a/sdk/wren-langchain/tests/unit/test_tools_runtime.py b/sdk/wren-langchain/tests/unit/test_tools_runtime.py new file mode 100644 index 000000000..7b9bd0823 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_tools_runtime.py @@ -0,0 +1,141 @@ +"""Tests for runtime LLM-facing tools: wren_query, wren_dry_plan, wren_list_models.""" + +import json +from unittest.mock import MagicMock, patch + +import pyarrow as pa +from wren.model.error import ErrorCode, ErrorPhase, WrenError + +from wren_langchain import WrenToolkit + + +def _get_tool(toolkit, name): + return next(t for t in toolkit.get_tools() if t.name == name) + + +def test_get_tools_returns_three_runtime_tools_when_memory_disabled( + tmp_project, fake_active_profile +): + """When memory is off, get_tools returns the 3 runtime-only tools.""" + toolkit = WrenToolkit.from_project(tmp_project) + + tools = toolkit.get_tools() + names = sorted(t.name for t in tools) + + assert names == ["wren_dry_plan", "wren_list_models", "wren_query"] + + +def test_wren_query_success_envelope(tmp_project, fake_active_profile): + fake_table = pa.table({"id": [1, 2], "name": ["a", "b"]}) + fake_engine = MagicMock(name="engine") + fake_engine.query.return_value = fake_table + fake_engine._connector = MagicMock() + + toolkit = WrenToolkit.from_project(tmp_project) + tool = _get_tool(toolkit, "wren_query") + + with patch("wren_langchain._toolkit.WrenEngine", return_value=fake_engine): + envelope = tool.invoke({"sql": "SELECT * FROM x", "limit": 50}) + + assert envelope["ok"] is True + assert envelope["data"]["columns"] == ["id", "name"] + assert envelope["data"]["rows"] == [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}] + assert envelope["data"]["row_count"] == 2 + assert envelope["data"]["content_truncated"] is False + assert json.loads(envelope["content"]) == [ + {"id": 1, "name": "a"}, + {"id": 2, "name": "b"}, + ] + + +def test_wren_query_rejects_limit_above_hard_cap(tmp_project, fake_active_profile): + """The LLM tool must guard against runaway `limit` values before + materializing rows (typo / hallucinated huge number → memory blow-up).""" + toolkit = WrenToolkit.from_project(tmp_project) + tool = _get_tool(toolkit, "wren_query") + + # Don't even need to patch WrenEngine — validation should fire before + # toolkit.query is called. + envelope = tool.invoke({"sql": "SELECT 1", "limit": 100_000}) + + assert envelope["ok"] is False + assert "1 and 1000" in envelope["content"] + + +def test_wren_query_rejects_zero_or_negative_limit(tmp_project, fake_active_profile): + """Zero / negative limits are nonsensical and would either return zero + rows or fail in the DB layer with a less helpful message.""" + toolkit = WrenToolkit.from_project(tmp_project) + tool = _get_tool(toolkit, "wren_query") + + envelope = tool.invoke({"sql": "SELECT 1", "limit": 0}) + + assert envelope["ok"] is False + assert "1 and 1000" in envelope["content"] + + +def test_wren_query_error_envelope_on_wren_error(tmp_project, fake_active_profile): + fake_engine = MagicMock(name="engine") + fake_engine.query.side_effect = WrenError( + error_code=ErrorCode.INVALID_SQL, + message="syntax error", + phase=ErrorPhase.SQL_PARSING, + ) + fake_engine._connector = MagicMock() + + toolkit = WrenToolkit.from_project(tmp_project) + tool = _get_tool(toolkit, "wren_query") + + with patch("wren_langchain._toolkit.WrenEngine", return_value=fake_engine): + envelope = tool.invoke({"sql": "SELEC * FROM x"}) + + assert envelope["ok"] is False + assert envelope["error"]["code"] == "INVALID_SQL" + assert envelope["error"]["phase"] == "SQL_PARSING" + assert "syntax error" in envelope["content"] + + +def test_wren_dry_plan_returns_sql_code_block(tmp_project, fake_active_profile): + fake_engine = MagicMock(name="engine") + fake_engine.dry_plan.return_value = ( + "WITH cte_orders AS (...) SELECT * FROM cte_orders" + ) + + toolkit = WrenToolkit.from_project(tmp_project) + tool = _get_tool(toolkit, "wren_dry_plan") + + with patch("wren_langchain._toolkit.WrenEngine", return_value=fake_engine): + envelope = tool.invoke({"sql": "SELECT * FROM orders"}) + + assert envelope["ok"] is True + assert envelope["content"].startswith("```sql") + assert envelope["content"].endswith("```") + assert "cte_orders" in envelope["data"]["dialect_sql"] + + +def test_wren_list_models_returns_markdown_table(tmp_project, fake_active_profile): + manifest = { + "models": [ + { + "name": "orders", + "columns": [{"name": "id"}, {"name": "customer"}], + "properties": {"description": "Customer orders"}, + }, + { + "name": "customers", + "columns": [{"name": "id"}], + }, + ] + } + (tmp_project / "target" / "mdl.json").write_text(json.dumps(manifest)) + + toolkit = WrenToolkit.from_project(tmp_project) + tool = _get_tool(toolkit, "wren_list_models") + + envelope = tool.invoke({}) + + assert envelope["ok"] is True + assert "| model | cols | description |" in envelope["content"] + assert "| orders | 2 | Customer orders |" in envelope["content"] + assert "| customers | 1 | |" in envelope["content"] + assert len(envelope["data"]["models"]) == 2