diff --git a/sdks/python/README.md b/sdks/python/README.md index 664a81d6b..6abac99f7 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -137,7 +137,9 @@ results = await client.search.query( "API documentation", filters=SearchFilters( data_sources=["api-docs.pdf"], - document_types=["application/pdf"] + document_types=["application/pdf"], + owners=["user@example.com"], + connector_types=["google_drive"], ), limit=5, score_threshold=0.5 diff --git a/sdks/python/openrag_sdk/chat.py b/sdks/python/openrag_sdk/chat.py index 749ca54ae..59ad4e9b4 100644 --- a/sdks/python/openrag_sdk/chat.py +++ b/sdks/python/openrag_sdk/chat.py @@ -1,7 +1,8 @@ """OpenRAG SDK chat client with streaming support.""" import json -from typing import TYPE_CHECKING, Any, AsyncIterator, Literal, overload +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Any, Literal, overload import httpx @@ -253,7 +254,8 @@ async def create( message: The message to send. stream: Whether to stream the response (default False). chat_id: ID of existing conversation to continue. - filters: Optional search filters (data_sources, document_types). + filters: Optional search filters (data_sources, document_types, owners, + connector_types). limit: Maximum number of search results (default 10). score_threshold: Minimum search score threshold (default 0). filter_id: Optional knowledge filter ID to apply. @@ -414,7 +416,8 @@ def stream( Args: message: The message to send. chat_id: ID of existing conversation to continue. - filters: Optional search filters (data_sources, document_types). + filters: Optional search filters (data_sources, document_types, owners, + connector_types). limit: Maximum number of search results (default 10). score_threshold: Minimum search score threshold (default 0). filter_id: Optional knowledge filter ID to apply. @@ -500,7 +503,3 @@ async def delete(self, chat_id: str) -> bool: return data.get("success", False) except NotFoundError: return False - - -# Import Literal for type hints -from typing import Literal diff --git a/sdks/python/openrag_sdk/models.py b/sdks/python/openrag_sdk/models.py index b0eee0864..a589ee9be 100644 --- a/sdks/python/openrag_sdk/models.py +++ b/sdks/python/openrag_sdk/models.py @@ -168,6 +168,8 @@ class SearchFilters(BaseModel): data_sources: list[str] | None = None document_types: list[str] | None = None + owners: list[str] | None = None + connector_types: list[str] | None = None # Settings update models diff --git a/sdks/python/openrag_sdk/search.py b/sdks/python/openrag_sdk/search.py index d084840ad..384b1adf5 100644 --- a/sdks/python/openrag_sdk/search.py +++ b/sdks/python/openrag_sdk/search.py @@ -2,8 +2,6 @@ from typing import TYPE_CHECKING, Any -import httpx - from .models import SearchFilters, SearchResponse, SearchResult if TYPE_CHECKING: @@ -30,7 +28,8 @@ async def query( Args: query: The search query text. - filters: Optional filters (data_sources, document_types). + filters: Optional filters (data_sources, document_types, owners, + connector_types). limit: Maximum number of results (default 10). score_threshold: Minimum score threshold (default 0). filter_id: Optional knowledge filter ID to apply. diff --git a/sdks/python/tests/test_models.py b/sdks/python/tests/test_models.py new file mode 100644 index 000000000..8ba5d2b68 --- /dev/null +++ b/sdks/python/tests/test_models.py @@ -0,0 +1,17 @@ +from openrag_sdk import SearchFilters + + +def test_search_filters_include_owner_and_connector_filters(): + filters = SearchFilters( + data_sources=["api-docs.pdf"], + document_types=["application/pdf"], + owners=["user@example.com"], + connector_types=["google_drive"], + ) + + assert filters.model_dump(exclude_none=True) == { + "data_sources": ["api-docs.pdf"], + "document_types": ["application/pdf"], + "owners": ["user@example.com"], + "connector_types": ["google_drive"], + } diff --git a/sdks/typescript/README.md b/sdks/typescript/README.md index 9526ca923..07c75d392 100644 --- a/sdks/typescript/README.md +++ b/sdks/typescript/README.md @@ -160,6 +160,8 @@ const results = await client.search.query("API documentation", { filters: { data_sources: ["api-docs.pdf"], document_types: ["application/pdf"], + owners: ["user@example.com"], + connector_types: ["google_drive"], }, limit: 5, scoreThreshold: 0.5, diff --git a/sdks/typescript/src/types.ts b/sdks/typescript/src/types.ts index 2d2336ee8..a1a5450f5 100644 --- a/sdks/typescript/src/types.ts +++ b/sdks/typescript/src/types.ts @@ -52,6 +52,8 @@ export interface SearchResponse { export interface SearchFilters { data_sources?: string[]; document_types?: string[]; + owners?: string[]; + connector_types?: string[]; } // Document types diff --git a/sdks/typescript/tests/types.test.ts b/sdks/typescript/tests/types.test.ts new file mode 100644 index 000000000..310c8c447 --- /dev/null +++ b/sdks/typescript/tests/types.test.ts @@ -0,0 +1,20 @@ +import { describe, expect, it } from "vitest"; +import type { SearchFilters } from "../src"; + +describe("SearchFilters", () => { + it("accepts owner and connector type filters", () => { + const filters: SearchFilters = { + data_sources: ["api-docs.pdf"], + document_types: ["application/pdf"], + owners: ["user@example.com"], + connector_types: ["google_drive"], + }; + + expect(filters).toEqual({ + data_sources: ["api-docs.pdf"], + document_types: ["application/pdf"], + owners: ["user@example.com"], + connector_types: ["google_drive"], + }); + }); +});