Skip to content

Commit ba268aa

Browse files
committed
fix: fix collection type
1 parent 27f68d0 commit ba268aa

File tree

2 files changed

+5
-1
lines changed

2 files changed

+5
-1
lines changed

taskingai/client/models/entities/collection.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from typing import Dict
1717

1818

19-
__all__ = ["Collection"]
19+
__all__ = ["Collection", "CollectionType"]
2020

2121

2222
class CollectionType(str, Enum):

taskingai/retrieval/collection.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ def create_collection(
107107
embedding_model_id: str,
108108
capacity: int = 1000,
109109
name: Optional[str] = None,
110+
type: Optional[CollectionType] = None,
110111
description: Optional[str] = None,
111112
metadata: Optional[Dict[str, str]] = None,
112113
) -> Collection:
@@ -125,6 +126,7 @@ def create_collection(
125126
embedding_model_id=embedding_model_id,
126127
capacity=capacity,
127128
name=name or "",
129+
type=type or CollectionType.TEXT,
128130
description=description or "",
129131
metadata=metadata or {},
130132
)
@@ -137,6 +139,7 @@ async def a_create_collection(
137139
embedding_model_id: str,
138140
capacity: int = 1000,
139141
name: Optional[str] = None,
142+
type: Optional[CollectionType] = None,
140143
description: Optional[str] = None,
141144
metadata: Optional[Dict[str, str]] = None,
142145
) -> Collection:
@@ -156,6 +159,7 @@ async def a_create_collection(
156159
embedding_model_id=embedding_model_id,
157160
capacity=capacity,
158161
name=name or "",
162+
type=type or CollectionType.TEXT,
159163
description=description or "",
160164
metadata=metadata or {},
161165
)

0 commit comments

Comments
 (0)