Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Weaviate multitenancy support #2199

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Weaviate multitenancy support #2199

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions docs/docs/deep-dive/retrieval_models_clients/WeaviateRM.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,20 @@ WeaviateRM(
)
```

## Using Multitenancy
If your Weaviate instance is tenant-aware, you can provide a tenant_id in the WeaviateRM constructor or as a keyword argument:

```python
retriever_model = WeaviateRM(
weaviate_collection_name="<WEAVIATE_COLLECTION>",
weaviate_client=weaviate_client,
tenant_id="tenant123"
)

results = retriever_model("Your query here", tenant_id="tenantXYZ")
```
When tenant_id is specified, this will scope all retrieval requests to the tenant ID provided.

## Under the Hood

`forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = None, **kwargs) -> dspy.Prediction`
Expand Down
25 changes: 12 additions & 13 deletions dspy/retrieve/weaviate_rm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,11 +51,13 @@ def __init__(
weaviate_client: Union[weaviate.WeaviateClient, weaviate.Client],
weaviate_collection_text_key: Optional[str] = "content",
k: int = 3,
tenant_id: Optional[str] = None,
):
self._weaviate_collection_name = weaviate_collection_name
self._weaviate_client = weaviate_client
self._weaviate_collection = self._weaviate_client.collections.get(self._weaviate_collection_name)
self._weaviate_collection_text_key = weaviate_collection_text_key
self._tenant_id = tenant_id

# Check the type of weaviate_client (this is added to support v3 and v4)
if hasattr(weaviate_client, "collections"):
Expand All @@ -82,26 +84,23 @@ def forward(self, query_or_queries: Union[str, List[str]], k: Optional[int] = No
queries = [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries
queries = [q for q in queries if q]
passages, parsed_results = [], []
tenant = kwargs.pop("tenant_id", self._tenant_id)
for query in queries:
if self._client_type == "WeaviateClient":
results = self._weaviate_collection.query.hybrid(
query=query,
limit=k,
**kwargs,
)
if tenant:
results = self._weaviate_collection.query.tenant(tenant).hybrid(query=query, limit=k, **kwargs)
else:
results = self._weaviate_collection.query.hybrid(query=query, limit=k, **kwargs)

parsed_results = [result.properties[self._weaviate_collection_text_key] for result in results.objects]

elif self._client_type == "Client":
results = (
self._weaviate_client.query.get(
self._weaviate_collection_name,
[self._weaviate_collection_text_key],
q = self._weaviate_client.query.get(
self._weaviate_collection_name, [self._weaviate_collection_text_key]
)
.with_hybrid(query=query)
.with_limit(k)
.do()
)
if tenant:
q = q.with_tenant(tenant)
results = q.with_hybrid(query=query).with_limit(k).do()

results = results["data"]["Get"][self._weaviate_collection_name]
parsed_results = [result[self._weaviate_collection_text_key] for result in results]
Expand Down