Skip to content

Commit c8f441c

Browse files
PawanThakurIBMpriyanshu-krishnan1GeetikaChughIBMlorenzejayDhruv Chaturvedi
authored
feat(crewai-tools): add IBM Db2 search tool (#5885)
* feat(crewai-tools): add db2 search tool * refactor(crewai-tools): improve db2 search tool implementation * feat(tools): improve DB2VectorSearchTool validation, security, and configurability * docs: add DB2SearchTool documentation * feat: add DB2 search tool * docs: update DB2SearchTool documentation * fix: address CodeRabbit review feedback * fix: validate non-empty filter_by in DB2ToolSchema * chore: trigger CodeRabbit re-review * feat: fortify DB2 tool; fixed JSON response shape, added input guards and config validation * refactor(db2): replace DB2Config with connection_string field * refactor(db2): remove dead _setup_db2 validator and importlib import * refactor(db2): remove dead guard in _connect as _disconnect() is called at the end of every _run, so self.connection is always None when _connect is called next. The 'if not self.connection' guard was dead code. * fix(db2): tighten _validate_identifier regex. Old regex allowed leading digits, multiple periods and dot-only strings (e.g. '.....' passed). * fix(db2): replace __import__ with importlib.import_module in _generate_embedding as keeping openai as a lazy optional import since it is not always required. * perf(db2): cache OpenAI client in _openai_client to avoid re-instantiation as OpenAI(api_key=...) was recreated on every _generate_embedding call. Extract into _get_openai_client() which lazily initialises and caches self._openai_client on first use, reusing it for all subsequent queries. * docs(db2): clarify tool description to mention embedding fallback * docs(db2): update README supported features to clarify embedding behaviour. 'OpenAI embedding fallback' implied it was optional. Replaced with 'Uses a custom embedding function if supplied, otherwise OpenAI embeddings.' * updated both code examples to use the correct import path and public run() method. * feat(crewai-tools): add db2 search tool * refactor(crewai-tools): improve db2 search tool implementation * feat(tools): improve DB2VectorSearchTool validation, security, and configurability * docs: add DB2SearchTool documentation * feat: add DB2 search tool * docs: update DB2SearchTool documentation * fix: address CodeRabbit review feedback * fix: validate non-empty filter_by in DB2ToolSchema * chore: trigger CodeRabbit re-review * feat: fortify DB2 tool; fixed JSON response shape, added input guards and config validation * fix(db2): address ruff and mypy linter errors * style(db2): apply ruff format to db2_search_tool.py * fix(db2-search-tool): address PR review comments - Restore DirectoryReadTool export accidentally removed; add DB2VectorSearchTool and DB2ToolSchema to crewai_tools.tools __init__ and __all__ - Align _ALLOWED_METRICS whitelist with Db2 VECTOR_DISTANCE API: replace DOT_PRODUCT/L2_DISTANCE with EUCLIDEAN_SQUARED/DOT/HAMMING/MANHATTAN - Replace ImportString fields for db2_package/db2_dbi_package with plain Any + lazy importlib.import_module in new _resolve_db2_packages() to avoid Pydantic default-validation gap where strings were never resolved at construction time - Move docs from frozen docs/v1.13.0/ snapshot to docs/edge/en/tools/database-data/ and register in docs/docs.json; update examples to match actual API (connection_string constructor, not DB2Config), correct return format, and align documented distance metrics with the whitelist * fix(db2-search-tool): resolve default and string db2 package imports dynamically * fix(db2-search-tool): export DB2VectorSearchTool and DB2ToolSchema from package-level crewai_tools * docs(db2-search-tool): fix installation command and import path in README --------- Co-authored-by: priyanshu-krishnan1 <priyanshu.krishnan1@ibm.com> Co-authored-by: GeetikaChugh24 <geetika@ibm.com> Co-authored-by: Lorenze Jay <63378463+lorenzejay@users.noreply.github.com> Co-authored-by: Dhruv Chaturvedi <dhruv_insights@Dhruvs-MacBook-Pro.local>
1 parent 3932d3f commit c8f441c

8 files changed

Lines changed: 1398 additions & 1 deletion

File tree

docs/docs.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,8 @@
271271
"edge/en/tools/database-data/qdrantvectorsearchtool",
272272
"edge/en/tools/database-data/weaviatevectorsearchtool",
273273
"edge/en/tools/database-data/mongodbvectorsearchtool",
274-
"edge/en/tools/database-data/singlestoresearchtool"
274+
"edge/en/tools/database-data/singlestoresearchtool",
275+
"edge/en/tools/database-data/db2searchtool"
275276
]
276277
},
277278
{
Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
---
2+
title: Db2 Vector Search Tool
3+
description: Semantic vector search for CrewAI agents using IBM Db2 native VECTOR_DISTANCE capabilities.
4+
icon: database
5+
mode: "wide"
6+
---
7+
8+
# `DB2VectorSearchTool`
9+
10+
## Description
11+
12+
Perform semantic vector similarity searches against IBM Db2 tables using the native `VECTOR_DISTANCE` function.
13+
Supports configurable distance metrics, OpenAI or custom embeddings, metadata filtering, and result shaping.
14+
15+
## Installation
16+
17+
```bash
18+
pip install ibm_db openai
19+
```
20+
21+
Or with uv:
22+
23+
```bash
24+
uv add ibm_db openai
25+
```
26+
27+
## Environment Variables
28+
29+
```bash
30+
OPENAI_API_KEY=your_openai_key # Required when using default OpenAI embeddings
31+
DB2_CONNECTION_STRING=DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;
32+
```
33+
34+
## Basic Usage
35+
36+
```python
37+
from crewai import Agent
38+
from crewai_tools import DB2VectorSearchTool
39+
40+
tool = DB2VectorSearchTool(
41+
connection_string="DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;",
42+
table_name="documents",
43+
vector_column="embedding",
44+
)
45+
46+
agent = Agent(
47+
role="Research Assistant",
48+
goal="Find relevant information in documents",
49+
tools=[tool],
50+
)
51+
```
52+
53+
## Full Semantic Search Workflow
54+
55+
```python
56+
import os
57+
from dotenv import load_dotenv
58+
from crewai import Agent, Task, Crew, Process
59+
from crewai_tools import DB2VectorSearchTool
60+
61+
load_dotenv()
62+
63+
db2_tool = DB2VectorSearchTool(
64+
connection_string=os.getenv("DB2_CONNECTION_STRING"),
65+
table_name="documents",
66+
vector_column="embedding",
67+
return_columns=["content", "category"],
68+
limit=3,
69+
distance_metric="COSINE",
70+
max_distance=0.35,
71+
)
72+
73+
search_agent = Agent(
74+
role="Senior Semantic Search Agent",
75+
goal="Find and analyse documents based on semantic search",
76+
backstory="You are an expert research assistant who can find relevant information using semantic search in a Db2 database.",
77+
tools=[db2_tool],
78+
verbose=True,
79+
)
80+
81+
answer_agent = Agent(
82+
role="Senior Answer Assistant",
83+
goal="Generate answers based on retrieved context",
84+
backstory="You are an expert assistant who generates answers from provided context.",
85+
tools=[db2_tool],
86+
verbose=True,
87+
)
88+
89+
search_task = Task(
90+
description="""Search for relevant documents about {query}.
91+
Include the relevant information found, vector distances, and returned fields.""",
92+
agent=search_agent,
93+
)
94+
95+
answer_task = Task(
96+
description="Given the retrieved Db2 context, generate a final answer.",
97+
agent=answer_agent,
98+
)
99+
100+
crew = Crew(
101+
agents=[search_agent, answer_agent],
102+
tasks=[search_task, answer_task],
103+
process=Process.sequential,
104+
verbose=True,
105+
)
106+
107+
result = crew.kickoff(inputs={"query": "What is the role of X in the document?"})
108+
print(result)
109+
```
110+
111+
## Tool Parameters
112+
113+
| Parameter | Type | Default | Description |
114+
|---|---|---|---|
115+
| `connection_string` | `str` | required | Db2 connection string. Format: `DATABASE=x;HOSTNAME=x;PORT=50000;PROTOCOL=TCPIP;UID=x;PWD=x;` |
116+
| `table_name` | `str` | `"documents"` | Table to search. Supports `schema.table` notation. |
117+
| `vector_column` | `str` | `"embedding"` | Column storing the vector embeddings. |
118+
| `embedding_model` | `str` | `"text-embedding-3-large"` | OpenAI model used when no custom embedding function is provided. |
119+
| `return_columns` | `list[str]` | `["content"]` | Columns to include in each result. Must contain at least one entry. |
120+
| `limit` | `int` | `3` | Maximum number of results (1–100). |
121+
| `distance_metric` | `str` | `"COSINE"` | Db2 distance metric. See supported values below. |
122+
| `max_distance` | `float \| None` | `None` | Drop results whose distance exceeds this value. |
123+
| `custom_embedding_fn` | `Callable[[str], list[float]] \| None` | `None` | Custom embedding function. Overrides OpenAI when provided. |
124+
125+
## Supported Distance Metrics
126+
127+
The following values map directly to the Db2 `VECTOR_DISTANCE` function:
128+
129+
- `COSINE`
130+
- `EUCLIDEAN`
131+
- `EUCLIDEAN_SQUARED`
132+
- `DOT`
133+
- `HAMMING`
134+
- `MANHATTAN`
135+
136+
Reference: [IBM Db2 VECTOR_DISTANCE documentation](https://www.ibm.com/docs/en/db2/12.1.x?topic=functions-vector-distance)
137+
138+
## Schema Parameters (per query)
139+
140+
| Parameter | Type | Required | Description |
141+
|---|---|---|---|
142+
| `query` | `str` || The search query. |
143+
| `filter_by` | `str \| None` || Column name for metadata filtering. Must be paired with `filter_value`. |
144+
| `filter_value` | `Any \| None` || Value to filter on. Must be paired with `filter_by`. |
145+
146+
## Return Format
147+
148+
```json
149+
{
150+
"success": true,
151+
"results": [
152+
{
153+
"distance": 0.1401,
154+
"data": {
155+
"content": "Document content here",
156+
"category": "research"
157+
}
158+
}
159+
]
160+
}
161+
```
162+
163+
On error:
164+
165+
```json
166+
{
167+
"success": false,
168+
"error": "Description of what went wrong",
169+
"error_type": "ExceptionClassName"
170+
}
171+
```
172+
173+
## Metadata Filtering
174+
175+
```python
176+
result = db2_tool.run(
177+
query="machine learning",
178+
filter_by="category",
179+
filter_value="research",
180+
)
181+
```
182+
183+
`filter_by` and `filter_value` must always be provided together. Providing only one raises a validation error.
184+
185+
## Custom Embeddings
186+
187+
Use any embedding model by supplying a `custom_embedding_fn`:
188+
189+
```python
190+
from sentence_transformers import SentenceTransformer
191+
from crewai_tools import DB2VectorSearchTool
192+
193+
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
194+
195+
def custom_embeddings(text: str) -> list[float]:
196+
return model.encode(text).tolist()
197+
198+
tool = DB2VectorSearchTool(
199+
connection_string="DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;",
200+
table_name="documents",
201+
custom_embedding_fn=custom_embeddings,
202+
)
203+
```
204+
205+
When `custom_embedding_fn` is provided, `OPENAI_API_KEY` is not required.
206+
207+
## Security Features
208+
209+
- SQL identifier validation (table, column names must match `^[A-Za-z][A-Za-z0-9_]*(\.[A-Za-z][A-Za-z0-9_]*)?$`)
210+
- Parameterised SQL queries — values never interpolated into SQL strings
211+
- Distance metric whitelist — only valid Db2 metric names accepted

lib/crewai-tools/src/crewai_tools/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,10 @@
5959
from crewai_tools.tools.databricks_query_tool.databricks_query_tool import (
6060
DatabricksQueryTool,
6161
)
62+
from crewai_tools.tools.db2_search_tool import (
63+
DB2ToolSchema,
64+
DB2VectorSearchTool,
65+
)
6266
from crewai_tools.tools.daytona_sandbox_tool import (
6367
DaytonaExecTool,
6468
DaytonaFileTool,
@@ -248,6 +252,8 @@
248252
"CrewaiPlatformTools",
249253
"DOCXSearchTool",
250254
"DallETool",
255+
"DB2ToolSchema",
256+
"DB2VectorSearchTool",
251257
"DatabricksQueryTool",
252258
"DaytonaExecTool",
253259
"DaytonaFileTool",

lib/crewai-tools/src/crewai_tools/tools/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@
5353
DaytonaFileTool,
5454
DaytonaPythonTool,
5555
)
56+
from crewai_tools.tools.db2_search_tool import (
57+
DB2ToolSchema,
58+
DB2VectorSearchTool,
59+
)
5660
from crewai_tools.tools.directory_read_tool.directory_read_tool import (
5761
DirectoryReadTool,
5862
)
@@ -234,6 +238,8 @@
234238
"DOCXSearchTool",
235239
"DallETool",
236240
"DatabricksQueryTool",
241+
"DB2ToolSchema",
242+
"DB2VectorSearchTool",
237243
"DaytonaExecTool",
238244
"DaytonaFileTool",
239245
"DaytonaPythonTool",
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
# DB2 Vector Search Tool
2+
3+
IBM DB2 Vector Search Tool for CrewAI.
4+
5+
Supports:
6+
7+
- IBM DB2 native VECTOR search
8+
- OpenAI embeddings
9+
- Custom embedding functions
10+
- Metadata filtering
11+
- Runtime dynamic imports
12+
- Standardized CrewAI tool architecture
13+
14+
---
15+
16+
# Installation
17+
18+
```bash
19+
uv add ibm_db openai
20+
```
21+
22+
---
23+
24+
# Environment Variables
25+
26+
```env
27+
OPENAI_API_KEY=your_openai_key
28+
29+
DB2_CONNECTION_STRING=DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;
30+
```
31+
32+
---
33+
34+
# Example Usage
35+
36+
```python
37+
from crewai_tools import DB2VectorSearchTool
38+
39+
tool = DB2VectorSearchTool(
40+
connection_string="DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;",
41+
table_name="documents",
42+
)
43+
44+
result = tool.run(
45+
query="What is machine learning?",
46+
)
47+
48+
print(result)
49+
```
50+
51+
---
52+
53+
# Example With Metadata Filtering
54+
55+
```python
56+
result = tool.run(
57+
query="AI papers",
58+
filter_by="category",
59+
filter_value="AI",
60+
)
61+
```
62+
63+
---
64+
65+
# Supported Features
66+
67+
- DB2 VECTOR datatype
68+
- VECTOR_DISTANCE search
69+
- COSINE similarity
70+
- Metadata filtering
71+
- Uses a custom embedding function if supplied, otherwise OpenAI embeddings
72+
73+
---
74+
75+
# Architecture
76+
77+
This tool follows the same architecture as:
78+
79+
- QdrantVectorSearchTool
80+
- WeaviateVectorSearchTool
81+
82+
Responsibilities:
83+
84+
- Generate query embeddings
85+
- Perform vector similarity search
86+
- Apply optional metadata filters
87+
- Return normalized JSON results
88+
89+
This tool is retrieval-only.
90+
91+
Document ingestion should be handled separately.
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
from crewai_tools.tools.db2_search_tool.db2_search_tool import (
2+
DB2ToolSchema,
3+
DB2VectorSearchTool,
4+
)
5+
6+
7+
__all__ = [
8+
"DB2ToolSchema",
9+
"DB2VectorSearchTool",
10+
]

0 commit comments

Comments
 (0)