pip install ibm_db openai

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)

{
  "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",
)

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,
)