mirror of
https://github.com/Canner/WrenAI.git
synced 2026-08-30 18:00:36 +08:00
feat(sdk): add wren-langchain package for langchain/langgraph integration (#2247)
This commit is contained in:
Pin Hsu Chen
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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/
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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 `<path>/.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="<name>"` 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 <name>`).
|
||||
|
||||
### 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.
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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[<datasource>]` 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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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: "<nl>" → <sql preview> (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)
|
||||
@@ -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
|
||||
@@ -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 ``<project>/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="<user's 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="<user's question>")`
|
||||
Optionally narrow scope with `model="<name>"` 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="<user's original \
|
||||
question>", sql="<the SQL you ran>", 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="<bad name>", '
|
||||
'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}"
|
||||
@@ -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)
|
||||
@@ -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 ``<project>/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
|
||||
@@ -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
|
||||
``<project>/.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."
|
||||
)
|
||||
@@ -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 ``<path>/.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 ``<project>/.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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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.
|
||||
"""
|
||||
@@ -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"]
|
||||
@@ -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:
|
||||
<tmp_path>/
|
||||
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:"}),
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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 <path>/.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()
|
||||
@@ -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")
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user