feat: add turbopuffer vector store integration - #6627
Conversation
Add turbopuffer as a RAG vector store provider alongside ChromaDB and Qdrant: - client/config/factory/types/utils implementing the BaseClient protocol (sync + async), with lazy namespace creation and $dist score normalization - register the provider in the RAG factory, discriminated config union, and optional-import placeholders - add the 'turbopuffer' optional dependency (extra); the default embedder uses fastembed all-MiniLM-L6-v2, loaded lazily on first embed so config construction never triggers a model download, and passing your own embedding_function skips fastembed entirely - expose turbopuffer through the crewai-tools RagTool provider list - document turbopuffer in the knowledge vector-store section - tests for the client, utils, RagTool routing, and default-embedder laziness Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Address review: point to the regions list instead of leaving gcp-us-central1 unexplained. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…tore feat(rag): add turbopuffer vector store provider
📝 WalkthroughWalkthroughAdds Turbopuffer as a RAG vector-store provider, including typed configuration, synchronous and asynchronous client operations, factory dispatch, CrewAI Tools integration, optional dependencies, documentation, and comprehensive tests. ChangesTurbopuffer RAG provider
Sequence Diagram(s)sequenceDiagram
participant RAGConfig
participant TurbopufferFactory
participant TurbopufferClient
participant TurbopufferNamespace
RAGConfig->>TurbopufferFactory: create_client(config)
TurbopufferFactory->>TurbopufferClient: construct configured wrapper
TurbopufferClient->>TurbopufferClient: generate embedding
TurbopufferClient->>TurbopufferNamespace: write or query vectors
TurbopufferNamespace-->>TurbopufferClient: return rows and distances
TurbopufferClient-->>RAGConfig: return SearchResult values
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment Warning |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py`:
- Line 247: Update RagTool._create_provider_config() to import TurbopufferConfig
through the existing optional-import fallback, using MissingTurbopufferConfig
when Turbopuffer is unavailable instead of raising ModuleNotFoundError. Match
the established ChromaDB/Qdrant optional embedding_function handling where
applicable.
- Around line 249-251: Update the embedding setup around build_embedder and
TurbopufferConfig to normalize custom embedding results into a flat list[float]
per document or query, unwrapping the single-vector [[...]] shape and validating
the result. Ensure Turbopuffer writes and rank_by queries receive flat vectors,
and add vector-shape assertions to test_rag_tool_with_custom_embeddings and the
corresponding second test.
In `@lib/crewai/src/crewai/rag/turbopuffer/utils.py`:
- Around line 87-99: Update the row construction in the metadata normalization
flow so metadata is merged before assigning the reserved fields, ensuring
metadata cannot overwrite “id”, “vector”, or CONTENT_KEY. Preserve the existing
list, dict, and empty-metadata handling, and add regression coverage for
collisions with each reserved field.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: b8100642-0719-423c-8b68-8dc50cc56251
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (22)
docs/edge/en/concepts/knowledge.mdxlib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.pylib/crewai-tools/src/crewai_tools/tools/rag/types.pylib/crewai-tools/tests/tools/rag/rag_tool_test.pylib/crewai/pyproject.tomllib/crewai/src/crewai/rag/config/optional_imports/base.pylib/crewai/src/crewai/rag/config/optional_imports/protocols.pylib/crewai/src/crewai/rag/config/optional_imports/providers.pylib/crewai/src/crewai/rag/config/optional_imports/types.pylib/crewai/src/crewai/rag/config/types.pylib/crewai/src/crewai/rag/factory.pylib/crewai/src/crewai/rag/turbopuffer/__init__.pylib/crewai/src/crewai/rag/turbopuffer/client.pylib/crewai/src/crewai/rag/turbopuffer/config.pylib/crewai/src/crewai/rag/turbopuffer/constants.pylib/crewai/src/crewai/rag/turbopuffer/factory.pylib/crewai/src/crewai/rag/turbopuffer/types.pylib/crewai/src/crewai/rag/turbopuffer/utils.pylib/crewai/tests/rag/turbopuffer/__init__.pylib/crewai/tests/rag/turbopuffer/test_client.pylib/crewai/tests/rag/turbopuffer/test_config.pylib/crewai/tests/rag/turbopuffer/test_utils.py
| return QdrantConfig(**kwargs) | ||
|
|
||
| if provider == "turbopuffer": | ||
| from crewai.rag.turbopuffer.config import TurbopufferConfig |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Locate files:"
fd -a 'rag_tool.py|types.py' . | sed 's#^\./##' | head -80
echo
echo "Relevant class/function outline around rag_tool.py:"
if [ -f lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py ]; then
wc -l lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py
ast-grep outline lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py --view expanded || true
fi
echo
echo "Search Turbopuffer related code:"
rg -n "TurbopufferConfig|MissingTurbopufferConfig|from crewai\.rag\.turbopuffer|import crewai\\.rag\\.turbopuffer|MissingTurbopuffer" -S .Repository: crewAIInc/crewAI
Length of output: 7843
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "rag_tool.py relevant section:"
sed -n '220,255p' lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py | cat -n
echo
echo "lib/crewai/src/crewai/rag/config/types.py relevant section:"
sed -n '1,90p' lib/crewai/src/crewai/rag/config/types.py | cat -n
echo
echo "lib/crewai/src/crewai/rag/config/optional_imports/providers.py relevant section:"
sed -n '1,70p' lib/crewai/src/crewai/rag/config/optional_imports/providers.py | cat -n
echo
echo "lib/crewai-tools/src/crewai_tools/tools/rag/types.py relevant section:"
sed -n '1,220p' lib/crewai-tools/src/crewai_tools/tools/rag/types.py | cat -n
echo
echo "Static call-graph check for _create_provider_config:"
rg -n "_create_provider_config|_validate_adapter|_resolve_adapter|RagToolConfig|TurbopufferConfig|from crewai\.rag\.turbopuffer\.config import" lib/crewai-tools/src lib/crewai/src/crewai -SRepository: crewAIInc/crewAI
Length of output: 10318
Route Turbopuffer through the optional-import fallback.
RagTool._create_provider_config() imports TurbopufferConfig directly, so Turbopuffer support outside crewai.rag.config.types.py still raises ModuleNotFoundError instead of using the existing MissingTurbopufferConfig placeholder path. Use the same optional fallback here, matching ChromaDB/Qdrant’s optional embedding_function handling if needed.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py` at line 247, Update
RagTool._create_provider_config() to import TurbopufferConfig through the
existing optional-import fallback, using MissingTurbopufferConfig when
Turbopuffer is unavailable instead of raising ModuleNotFoundError. Match the
established ChromaDB/Qdrant optional embedding_function handling where
applicable.
| kwargs = dict(provider_config) | ||
| if embedding_function is not None: | ||
| kwargs["embedding_function"] = embedding_function |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate file =="
git ls-files | rg '(^|/)rag_tool\.py$|turbopuffer|test.*rag|rag.*test' || true
echo
echo "== target file outline/size =="
wc -l lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py 2>/dev/null || true
ast-grep outline lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py 2>/dev/null | sed -n '1,220p' || true
echo
echo "== relevant source lines =="
cat -n lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py | sed -n '220,285p'Repository: crewAIInc/crewAI
Length of output: 6303
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== search build_embedder and embeddings usage in rag =="
rg -n "build_embedder|embedding_function|embeddings|Turbo|turbopuffer|rank_by|embedding_model" lib/crewai-tools/src/crewai_tools/tools/rag -S || true
echo
echo "== test files in rag area =="
git ls-files | rg 'test.*rag|rag.*test|rag_tool' | sed -n '1,120p'
echo
echo "== diff/stat for rag_tool.py if present =="
git diff --stat lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py || true
git diff -- lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py | sed -n '1,220p' || trueRepository: crewAIInc/crewAI
Length of output: 5771
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== rag_tool_test embeddings references =="
cat -n lib/crewai-tools/tests/tools/rag/rag_tool_test.py | sed -n '1,260p'
echo
echo "== turbopuffer config/client outlines =="
wc -l lib/crewai/src/crewai/rag/turbopuffer/*.py
for f in lib/crewai/src/crewai/rag/turbopuffer/{config,client,types,utils}.py; do
echo "--- $f"
ast-grep outline "$f" 2>/dev/null | sed -n '1,220p' || true
done
echo
echo "== turbopuffer source relevant snippets =="
for f in lib/crewai/src/crewai/rag/turbopuffer/{config,client,types,utils}.py; do
echo "--- $f"
rg -n "embeddings?|embedding|rank_by|vector|docs|insert" "$f" -C 3 || true
done
echo
echo "== turbopuffer tests relevant snippets =="
for f in lib/crewai/tests/rag/turbopuffer/{test_client,test_config,test_utils}.py; do
echo "--- $f"
rg -n "embeddings?|embedding|rank_by|vector|docs|insert" "$f" -C 3 || true
doneRepository: crewAIInc/crewAI
Length of output: 35185
🌐 Web query:
Turbopuffer Python SDK Namespace rank_by vector embedding flat list float
💡 Result:
In the Turbopuffer Python SDK, you perform a vector search by using the rank_by parameter within the namespace.query method [1][2]. The rank_by parameter accepts a tuple (or list) formatted as ("vector", "ANN", embedding), where the embedding is a flat list of floats representing your vector [1][3]. Example usage: ns = tpuf.namespace("your-namespace-name") result = ns.query( rank_by=("vector", "ANN", [0.1, 0.2, 0.3]), top_k=10, include_attributes=["name"]) Key details: - Embedding: The embedding should be a flat list of floats (e.g., [0.1, 0.2,...]) [1][2]. - ANN: Stands for "approximate nearest neighbor," which is the standard search mode for vector embedding queries [2]. - Namespace: The query is executed on a specific namespace instance, initialized via tpuf.namespace("name") [1][4]. - Data format: When writing data, vectors are represented as arrays of numbers or base64-encoded strings [5]. During a query, providing a flat list of floats for the rank_by parameter is the standard approach [1][3]. - Response: You can specify vector_encoding as "float" or "base64" to control how vectors are returned, with "float" being the default [2]. If "float" is selected, the vectors in the response are returned as arrays of numbers [2].
Citations:
- 1: https://github.com/turbopuffer/turbopuffer-python/blob/main/README.md
- 2: https://turbopuffer.com/docs/query
- 3: https://turbopuffer.com/docs/vector
- 4: https://pypi.org/project/turbopuffer/1.12.1/
- 5: https://turbopuffer.com/docs/write
Normalize embedding_model outputs for Turbopuffer.
build_embedder is passed into TurbopufferConfig unchanged, but the injected function in test_rag_tool_with_custom_embeddings returns [[...]]. Turbopuffer expects a flat list[float] for each document/query vector, so this can send nested vectors to writes and rank_by query calls. Unwrap/validate the single-vector shape and add a vector-shape assertion to both tests.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@lib/crewai-tools/src/crewai_tools/tools/rag/rag_tool.py` around lines 249 -
251, Update the embedding setup around build_embedder and TurbopufferConfig to
normalize custom embedding results into a flat list[float] per document or
query, unwrapping the single-vector [[...]] shape and validating the result.
Ensure Turbopuffer writes and rank_by queries receive flat vectors, and add
vector-shape assertions to test_rag_tool_with_custom_embeddings and the
corresponding second test.
| row: dict[str, Any] = { | ||
| "id": doc_id, | ||
| "vector": embedding, | ||
| CONTENT_KEY: doc["content"], | ||
| } | ||
|
|
||
| metadata = doc.get("metadata", {}) | ||
| if isinstance(metadata, list): | ||
| metadata = metadata[0] if metadata else {} | ||
| elif not isinstance(metadata, dict): | ||
| metadata = dict(metadata) if metadata else {} | ||
|
|
||
| row.update(metadata) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Prevent metadata from overwriting reserved row fields.
Line 99 lets metadata replace id, vector, or content, so a valid record can be persisted with the wrong identity, embedding, or body. Merge metadata first, then write the reserved fields; add regression coverage for each collision.
Proposed fix
- row: dict[str, Any] = {
- "id": doc_id,
- "vector": embedding,
- CONTENT_KEY: doc["content"],
- }
-
metadata = doc.get("metadata", {})
if isinstance(metadata, list):
metadata = metadata[0] if metadata else {}
elif not isinstance(metadata, dict):
metadata = dict(metadata) if metadata else {}
- row.update(metadata)
+ row = dict(metadata)
+ row.update(
+ {
+ "id": doc_id,
+ "vector": embedding,
+ CONTENT_KEY: doc["content"],
+ }
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| row: dict[str, Any] = { | |
| "id": doc_id, | |
| "vector": embedding, | |
| CONTENT_KEY: doc["content"], | |
| } | |
| metadata = doc.get("metadata", {}) | |
| if isinstance(metadata, list): | |
| metadata = metadata[0] if metadata else {} | |
| elif not isinstance(metadata, dict): | |
| metadata = dict(metadata) if metadata else {} | |
| row.update(metadata) | |
| metadata = doc.get("metadata", {}) | |
| if isinstance(metadata, list): | |
| metadata = metadata[0] if metadata else {} | |
| elif not isinstance(metadata, dict): | |
| metadata = dict(metadata) if metadata else {} | |
| row = dict(metadata) | |
| row.update( | |
| { | |
| "id": doc_id, | |
| "vector": embedding, | |
| CONTENT_KEY: doc["content"], | |
| } | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@lib/crewai/src/crewai/rag/turbopuffer/utils.py` around lines 87 - 99, Update
the row construction in the metadata normalization flow so metadata is merged
before assigning the reserved fields, ensuring metadata cannot overwrite “id”,
“vector”, or CONTENT_KEY. Preserve the existing list, dict, and empty-metadata
handling, and add regression coverage for collisions with each reserved field.
1 participant
Summary
Adds support for turbopuffer as a vector store integration
crewai.rag.turbopuffer— client (sync + async), config, factory, types, utilsRagToolrouting)