diff --git a/cognee/api/v1/search/routers/get_search_router.py b/cognee/api/v1/search/routers/get_search_router.py index cb3ef38a81..8538ac01b4 100644 --- a/cognee/api/v1/search/routers/get_search_router.py +++ b/cognee/api/v1/search/routers/get_search_router.py @@ -1,5 +1,6 @@ from uuid import UUID from datetime import datetime +from typing import Union from fastapi import Depends, APIRouter from fastapi.responses import JSONResponse from cognee.modules.search.types import SearchType @@ -12,6 +13,7 @@ class SearchPayloadDTO(InDTO): search_type: SearchType query: str + datasets: Union[list[str], str, None] = None def get_search_router() -> APIRouter: @@ -39,7 +41,10 @@ async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticat try: results = await cognee_search( - query_text=payload.query, query_type=payload.search_type, user=user + query_text=payload.query, + query_type=payload.search_type, + user=user, + datasets=payload.datasets ) return results