Db2 向量搜索工具

使用 IBM Db2 原生 VECTOR_DISTANCE 能力，为 CrewAI 智能体提供语义向量搜索。

DB2VectorSearchTool

描述

使用原生 VECTOR_DISTANCE 函数对 IBM Db2 表执行语义向量相似度搜索。 支持可配置的距离度量、OpenAI 或自定义嵌入、元数据过滤以及结果处理。

安装

pip install ibm_db openai

或使用 uv：

uv add ibm_db openai

环境变量

OPENAI_API_KEY=your_openai_key          # Required when using default OpenAI embeddings
DB2_CONNECTION_STRING=DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;

基本用法

from crewai import Agent
from crewai_tools import DB2VectorSearchTool

tool = DB2VectorSearchTool(
    connection_string="DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;",
    table_name="documents",
    vector_column="embedding",
)

agent = Agent(
    role="Research Assistant",
    goal="Find relevant information in documents",
    tools=[tool],
)

完整语义搜索工作流

import os
from dotenv import load_dotenv
from crewai import Agent, Task, Crew, Process
from crewai_tools import DB2VectorSearchTool

load_dotenv()

db2_tool = DB2VectorSearchTool(
    connection_string=os.getenv("DB2_CONNECTION_STRING"),
    table_name="documents",
    vector_column="embedding",
    return_columns=["content", "category"],
    limit=3,
    distance_metric="COSINE",
    max_distance=0.35,
)

search_agent = Agent(
    role="Senior Semantic Search Agent",
    goal="Find and analyse documents based on semantic search",
    backstory="You are an expert research assistant who can find relevant information using semantic search in a Db2 database.",
    tools=[db2_tool],
    verbose=True,
)

answer_agent = Agent(
    role="Senior Answer Assistant",
    goal="Generate answers based on retrieved context",
    backstory="You are an expert assistant who generates answers from provided context.",
    tools=[db2_tool],
    verbose=True,
)

search_task = Task(
    description="""Search for relevant documents about {query}.
    Include the relevant information found, vector distances, and returned fields.""",
    agent=search_agent,
)

answer_task = Task(
    description="Given the retrieved Db2 context, generate a final answer.",
    agent=answer_agent,
)

crew = Crew(
    agents=[search_agent, answer_agent],
    tasks=[search_task, answer_task],
    process=Process.sequential,
    verbose=True,
)

result = crew.kickoff(inputs={"query": "What is the role of X in the document?"})
print(result)

工具参数

参数	类型	默认值	描述
connection_string	str	必需	Db2 连接字符串。格式：DATABASE=x;HOSTNAME=x;PORT=50000;PROTOCOL=TCPIP;UID=x;PWD=x;
table_name	str	"documents"	要搜索的表。支持 schema.table 表示法。
vector_column	str	"embedding"	存储向量嵌入的列。
embedding_model	str	"text-embedding-3-large"	未提供自定义嵌入函数时使用的 OpenAI 模型。
return_columns	list[str]	["content"]	每个结果中要包含的列。必须至少包含一个条目。
limit	int	3	最大结果数（1–100）。
distance_metric	str	"COSINE"	Db2 距离度量。参见下方支持的值。
max_distance	float | None	None	丢弃距离超过此值的结果。
custom_embedding_fn	Callable[[str], list[float]] | None	None	自定义嵌入函数。提供后将覆盖 OpenAI。

支持的距离度量

以下值直接映射到 Db2 VECTOR_DISTANCE 函数：

• COSINE
• EUCLIDEAN
• EUCLIDEAN_SQUARED
• DOT
• HAMMING
• MANHATTAN

参考：IBM Db2 VECTOR_DISTANCE 文档

Schema 参数（每次查询）

参数	类型	必需	描述
query	str	✅	搜索查询。
filter_by	str | None	❌	用于元数据过滤的列名。必须与 filter_value 配对使用。
filter_value	Any | None	❌	要过滤的值。必须与 filter_by 配对使用。

返回格式

{
  "success": true,
  "results": [
    {
      "distance": 0.1401,
      "data": {
        "content": "Document content here",
        "category": "research"
      }
    }
  ]
}

出错时：

{
  "success": false,
  "error": "Description of what went wrong",
  "error_type": "ExceptionClassName"
}

元数据过滤

result = db2_tool.run(
    query="machine learning",
    filter_by="category",
    filter_value="research",
)

filter_by 和 filter_value 必须始终一起提供。仅提供其中一个会引发验证错误。

自定义嵌入

通过提供 custom_embedding_fn 来使用任意嵌入模型：

from sentence_transformers import SentenceTransformer
from crewai_tools import DB2VectorSearchTool

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

def custom_embeddings(text: str) -> list[float]:
    return model.encode(text).tolist()

tool = DB2VectorSearchTool(
    connection_string="DATABASE=TESTDB;HOSTNAME=localhost;PORT=50000;PROTOCOL=TCPIP;UID=db2user;PWD=password;",
    table_name="documents",
    custom_embedding_fn=custom_embeddings,
)

提供 custom_embedding_fn 时，不需要 OPENAI_API_KEY。

安全功能

• SQL 标识符验证（表、列名必须匹配 ^[A-Za-z][A-Za-z0-9_]*(\.[A-Za-z][A-Za-z0-9_]*)?$）
• 参数化 SQL 查询 —— 值绝不会被插入到 SQL 字符串中
• 距离度量白名单 —— 仅接受有效的 Db2 度量名称