From bf23ad87298b02f9aabd596cb582e0d99eeb6cae Mon Sep 17 00:00:00 2001 From: phact Date: Tue, 12 Aug 2025 12:37:55 -0400 Subject: [PATCH] search improvements (facets, limit, score threshold) --- frontend/components/ui/checkbox.tsx | 30 ++ frontend/components/ui/collapsible.tsx | 11 + frontend/package-lock.json | 62 +++ frontend/package.json | 2 + frontend/src/app/page.tsx | 708 +++++++++++++++++++++---- src/api/search.py | 6 +- src/auth_context.py | 31 +- src/services/search_service.py | 99 +++- 8 files changed, 839 insertions(+), 110 deletions(-) create mode 100644 frontend/components/ui/checkbox.tsx create mode 100644 frontend/components/ui/collapsible.tsx diff --git a/frontend/components/ui/checkbox.tsx b/frontend/components/ui/checkbox.tsx new file mode 100644 index 00000000..df61a138 --- /dev/null +++ b/frontend/components/ui/checkbox.tsx @@ -0,0 +1,30 @@ +"use client" + +import * as React from "react" +import * as CheckboxPrimitive from "@radix-ui/react-checkbox" +import { Check } from "lucide-react" + +import { cn } from "@/lib/utils" + +const Checkbox = React.forwardRef< + React.ElementRef, + React.ComponentPropsWithoutRef +>(({ className, ...props }, ref) => ( + + + + + +)) +Checkbox.displayName = CheckboxPrimitive.Root.displayName + +export { Checkbox } diff --git a/frontend/components/ui/collapsible.tsx b/frontend/components/ui/collapsible.tsx new file mode 100644 index 00000000..9fa48946 --- /dev/null +++ b/frontend/components/ui/collapsible.tsx @@ -0,0 +1,11 @@ +"use client" + +import * as CollapsiblePrimitive from "@radix-ui/react-collapsible" + +const Collapsible = CollapsiblePrimitive.Root + +const CollapsibleTrigger = CollapsiblePrimitive.CollapsibleTrigger + +const CollapsibleContent = CollapsiblePrimitive.CollapsibleContent + +export { Collapsible, CollapsibleTrigger, CollapsibleContent } diff --git a/frontend/package-lock.json b/frontend/package-lock.json index f016b6ab..10046beb 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -9,6 +9,8 @@ "version": "0.1.0", "dependencies": { "@radix-ui/react-avatar": "^1.1.10", + "@radix-ui/react-checkbox": "^1.3.2", + "@radix-ui/react-collapsible": "^1.1.11", "@radix-ui/react-dropdown-menu": "^2.1.15", "@radix-ui/react-label": "^2.1.7", "@radix-ui/react-navigation-menu": "^1.2.13", @@ -1086,6 +1088,66 @@ } } }, + "node_modules/@radix-ui/react-checkbox": { + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/@radix-ui/react-checkbox/-/react-checkbox-1.3.2.tgz", + "integrity": "sha512-yd+dI56KZqawxKZrJ31eENUwqc1QSqg4OZ15rybGjF2ZNwMO+wCyHzAVLRp9qoYJf7kYy0YpZ2b0JCzJ42HZpA==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.2", + "@radix-ui/react-compose-refs": "1.1.2", + "@radix-ui/react-context": "1.1.2", + "@radix-ui/react-presence": "1.1.4", + "@radix-ui/react-primitive": "2.1.3", + "@radix-ui/react-use-controllable-state": "1.2.2", + "@radix-ui/react-use-previous": "1.1.1", + "@radix-ui/react-use-size": "1.1.1" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-collapsible": { + "version": "1.1.11", + "resolved": "https://registry.npmjs.org/@radix-ui/react-collapsible/-/react-collapsible-1.1.11.tgz", + "integrity": "sha512-2qrRsVGSCYasSz1RFOorXwl0H7g7J1frQtgpQgYrt+MOidtPAINHn9CPovQXb83r8ahapdx3Tu0fa/pdFFSdPg==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.2", + "@radix-ui/react-compose-refs": "1.1.2", + "@radix-ui/react-context": "1.1.2", + "@radix-ui/react-id": "1.1.1", + "@radix-ui/react-presence": "1.1.4", + "@radix-ui/react-primitive": "2.1.3", + "@radix-ui/react-use-controllable-state": "1.2.2", + "@radix-ui/react-use-layout-effect": "1.1.1" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, "node_modules/@radix-ui/react-collection": { "version": "1.1.7", "resolved": "https://registry.npmjs.org/@radix-ui/react-collection/-/react-collection-1.1.7.tgz", diff --git a/frontend/package.json b/frontend/package.json index 0239b079..0e1dc341 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -10,6 +10,8 @@ }, "dependencies": { "@radix-ui/react-avatar": "^1.1.10", + "@radix-ui/react-checkbox": "^1.3.2", + "@radix-ui/react-collapsible": "^1.1.11", "@radix-ui/react-dropdown-menu": "^2.1.15", "@radix-ui/react-label": "^2.1.7", "@radix-ui/react-navigation-menu": "^1.2.13", diff --git a/frontend/src/app/page.tsx b/frontend/src/app/page.tsx index 9f1c41d5..dc6cff91 100644 --- a/frontend/src/app/page.tsx +++ b/frontend/src/app/page.tsx @@ -1,11 +1,13 @@ "use client" -import { useState } from "react" +import { useState, useEffect } from "react" import { Button } from "@/components/ui/button" import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card" import { Input } from "@/components/ui/input" import { Label } from "@/components/ui/label" -import { Search, Loader2, FileText, Zap } from "lucide-react" +import { Checkbox } from "@/components/ui/checkbox" +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible" +import { Search, Loader2, FileText, Zap, ChevronDown, ChevronUp, Filter, X } from "lucide-react" import { ProtectedRoute } from "@/components/protected-route" interface SearchResult { @@ -14,16 +16,56 @@ interface SearchResult { page: number text: string score: number + source_url?: string + owner?: string +} + +interface FacetBucket { + key: string + count: number +} + +interface Facets { + data_sources?: FacetBucket[] + document_types?: FacetBucket[] + owners?: FacetBucket[] +} + +interface SearchResponse { + results: SearchResult[] + aggregations: any + error?: string +} + +interface SelectedFilters { + data_sources: string[] + document_types: string[] + owners: string[] } function SearchPage() { const [query, setQuery] = useState("") const [loading, setLoading] = useState(false) const [results, setResults] = useState([]) + const [facets, setFacets] = useState({}) const [searchPerformed, setSearchPerformed] = useState(false) + const [selectedFilters, setSelectedFilters] = useState({ + data_sources: [], + document_types: [], + owners: [] + }) + const [openSections, setOpenSections] = useState({ + data_sources: true, + document_types: true, + owners: true + }) + const [sidebarOpen, setSidebarOpen] = useState(true) + const [resultLimit, setResultLimit] = useState(10) + const [scoreThreshold, setScoreThreshold] = useState(0) - const handleSearch = async (e: React.FormEvent) => { - e.preventDefault() + + const handleSearch = async (e?: React.FormEvent) => { + if (e) e.preventDefault() if (!query.trim()) return setLoading(true) @@ -35,28 +77,185 @@ function SearchPage() { headers: { "Content-Type": "application/json", }, - body: JSON.stringify({ query }), + body: JSON.stringify({ + query, + limit: resultLimit, + scoreThreshold, + ...(searchPerformed && { filters: selectedFilters }) + }), }) - const result = await response.json() + const result: SearchResponse = await response.json() if (response.ok) { setResults(result.results || []) + + // Process aggregations into facets + const aggs = result.aggregations + const processedFacets: Facets = {} + const newSelectedFilters: SelectedFilters = { + data_sources: [], + document_types: [], + owners: [] + } + + if (aggs && Object.keys(aggs).length > 0) { + processedFacets.data_sources = aggs.data_sources?.buckets?.map((b: any) => ({ key: b.key, count: b.doc_count })).filter((b: any) => b.count > 0) || [] + processedFacets.document_types = aggs.document_types?.buckets?.map((b: any) => ({ key: b.key, count: b.doc_count })).filter((b: any) => b.count > 0) || [] + processedFacets.owners = aggs.owners?.buckets?.map((b: any) => ({ key: b.key, count: b.doc_count })).filter((b: any) => b.count > 0) || [] + + // Set all filters as checked by default + newSelectedFilters.data_sources = processedFacets.data_sources?.map(f => f.key) || [] + newSelectedFilters.document_types = processedFacets.document_types?.map(f => f.key) || [] + newSelectedFilters.owners = processedFacets.owners?.map(f => f.key) || [] + } + + setFacets(processedFacets) + setSelectedFilters(newSelectedFilters) setSearchPerformed(true) } else { console.error("Search failed:", result.error) setResults([]) + setFacets({}) setSearchPerformed(true) } } catch (error) { console.error("Search error:", error) setResults([]) + setFacets({}) setSearchPerformed(true) } finally { setLoading(false) } } + const handleFilterChange = async (facetType: keyof SelectedFilters, value: string, checked: boolean) => { + const newFilters = { + ...selectedFilters, + [facetType]: checked + ? [...selectedFilters[facetType], value] + : selectedFilters[facetType].filter(item => item !== value) + } + + setSelectedFilters(newFilters) + + // Re-search immediately if search has been performed + if (searchPerformed && query.trim()) { + setLoading(true) + try { + const response = await fetch("/api/search", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + query, + limit: resultLimit, + scoreThreshold, + filters: newFilters + }), + }) + + const result: SearchResponse = await response.json() + + if (response.ok) { + setResults(result.results || []) + } else { + console.error("Search failed:", result.error) + setResults([]) + } + } catch (error) { + console.error("Search error:", error) + setResults([]) + } finally { + setLoading(false) + } + } + } + + const clearAllFilters = () => { + setSelectedFilters({ + data_sources: [], + document_types: [], + owners: [] + }) + } + + const selectAllFilters = () => { + setSelectedFilters({ + data_sources: facets.data_sources?.map(f => f.key) || [], + document_types: facets.document_types?.map(f => f.key) || [], + owners: facets.owners?.map(f => f.key) || [] + }) + } + + const toggleSection = (section: keyof typeof openSections) => { + setOpenSections(prev => ({ + ...prev, + [section]: !prev[section] + })) + } + + const getSelectedFilterCount = () => { + return selectedFilters.data_sources.length + + selectedFilters.document_types.length + + selectedFilters.owners.length + } + + const FacetSection = ({ + title, + buckets, + facetType, + isOpen, + onToggle + }: { + title: string + buckets: FacetBucket[] + facetType: keyof SelectedFilters + isOpen: boolean + onToggle: () => void + }) => { + if (!buckets || buckets.length === 0) return null + + return ( + + + + + + {buckets.map((bucket, index) => { + const isSelected = selectedFilters[facetType].includes(bucket.key) + return ( +
+ + handleFilterChange(facetType, bucket.key, checked as boolean) + } + /> + +
+ ) + })} + + + ) + } + return (
{/* Hero Section */} @@ -67,10 +266,10 @@ function SearchPage() {

- Find documents using semantic search + Find documents using hybrid search

- Enter your search query to find relevant documents using AI-powered semantic search across your document collection. + Enter your search query to find relevant documents using AI-powered semantic search combined with keyword matching across your document collection.

@@ -82,7 +281,7 @@ function SearchPage() { Search Documents - Enter your search query to find relevant documents using semantic search + Enter your search query to find relevant documents using hybrid search (semantic + keyword) @@ -91,115 +290,426 @@ function SearchPage() { - setQuery(e.target.value)} - className="h-12 bg-background/50 border-border/50 focus:border-blue-400/50 focus:ring-blue-400/20" - /> +
+ setQuery(e.target.value)} + className="h-12 bg-background/50 border-border/50 focus:border-blue-400/50 focus:ring-blue-400/20 flex-1" + /> + +
- - {/* Results Section */} -
- {searchPerformed ? ( -
-
+ {/* Search Results with Filters */} + {searchPerformed && ( +
+ {/* Filter Toggle - Always visible when filters are available */} + {(facets.data_sources?.length || facets.document_types?.length || facets.owners?.length) && ( +

Search Results

-
-
- - {results.length} result{results.length !== 1 ? 's' : ''} found - +
+
+
+ + {results.length} result{results.length !== 1 ? 's' : ''} returned + +
+
- {results.length === 0 ? ( - - -
-
- + )} + +
+ {/* Main Content */} +
+ {/* Active Filters Display */} + {getSelectedFilterCount() > 0 && getSelectedFilterCount() < (facets.data_sources?.length || 0) + (facets.document_types?.length || 0) + (facets.owners?.length || 0) && ( +
+
+

Active Filters

+
+ +
-

- No documents found -

-

- Try adjusting your search terms or check if documents have been indexed. -

- - - ) : ( -
- {results.map((result, index) => ( - - -
- -
- -
- {result.filename} - -
-
- - {result.score.toFixed(2)} - -
+
+ {Object.entries(selectedFilters).map(([facetType, values]) => + values.map((value: string) => ( +
+ {value} +
-
- - - {result.mimetype} - - - Page {result.page} - - - - -
-

- {result.text} + )) + )} +

+
+ )} + + {/* Results Section */} +
+ {results.length === 0 ? ( + + +
+
+ +
+

+ No documents found +

+

+ Try adjusting your search terms or check if documents have been indexed.

- ))} + ) : ( +
+ {results.map((result, index) => ( + + +
+ +
+ +
+ {result.filename} + +
+
+ + {result.score.toFixed(2)} + +
+
+
+ + + {result.mimetype} + + + Page {result.page} + + + + +
+

+ {result.text} +

+
+ + + ))} +
+ )} +
+
+ + {/* Right Sidebar - Filters */} + {(facets.data_sources?.length || facets.document_types?.length || facets.owners?.length) && sidebarOpen && ( +
+
+

+ + Filters +

+
+ + +
+
+ +
+ toggleSection('data_sources')} + /> + toggleSection('document_types')} + /> + toggleSection('owners')} + /> + + {/* Result Limit Control */} +
+
+
+ + { + const newLimit = Math.max(1, Math.min(1000, parseInt(e.target.value) || 1)) + setResultLimit(newLimit) + + // Re-search if search has been performed + if (searchPerformed && query.trim()) { + setLoading(true) + try { + const response = await fetch("/api/search", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + query, + limit: newLimit, + scoreThreshold, + filters: selectedFilters + }), + }) + + const result: SearchResponse = await response.json() + + if (response.ok) { + setResults(result.results || []) + } else { + console.error("Search failed:", result.error) + setResults([]) + } + } catch (error) { + console.error("Search error:", error) + setResults([]) + } finally { + setLoading(false) + } + } + }} + className="w-16 h-6 text-xs text-center" + /> +
+ { + const value = parseInt(e.target.value) + setResultLimit(value) + + // Re-search if search has been performed + if (searchPerformed && query.trim()) { + setLoading(true) + try { + const response = await fetch("/api/search", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + query, + limit: value, + scoreThreshold, + filters: selectedFilters + }), + }) + + const result: SearchResponse = await response.json() + + if (response.ok) { + setResults(result.results || []) + } else { + console.error("Search failed:", result.error) + setResults([]) + } + } catch (error) { + console.error("Search error:", error) + setResults([]) + } finally { + setLoading(false) + } + } + }} + className="w-full h-2 bg-gray-200 rounded-lg appearance-none cursor-pointer" + /> +
+ + {/* Score Threshold Control */} +
+
+ + { + const newThreshold = Math.max(0, Math.min(10, parseFloat(e.target.value) || 0)) + setScoreThreshold(newThreshold) + + // Re-search if search has been performed + if (searchPerformed && query.trim()) { + setLoading(true) + try { + const response = await fetch("/api/search", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + query, + limit: resultLimit, + scoreThreshold: newThreshold, + filters: selectedFilters + }), + }) + + const result: SearchResponse = await response.json() + + if (response.ok) { + setResults(result.results || []) + } else { + console.error("Search failed:", result.error) + setResults([]) + } + } catch (error) { + console.error("Search error:", error) + setResults([]) + } finally { + setLoading(false) + } + } + }} + className="w-16 h-6 text-xs text-center" + /> +
+ { + const value = parseFloat(e.target.value) + setScoreThreshold(value) + + // Re-search if search has been performed + if (searchPerformed && query.trim()) { + setLoading(true) + try { + const response = await fetch("/api/search", { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + query, + limit: resultLimit, + scoreThreshold: value, + filters: selectedFilters + }), + }) + + const result: SearchResponse = await response.json() + + if (response.ok) { + setResults(result.results || []) + } else { + console.error("Search failed:", result.error) + setResults([]) + } + } catch (error) { + console.error("Search error:", error) + setResults([]) + } finally { + setLoading(false) + } + } + }} + className="w-full h-2 bg-gray-200 rounded-lg appearance-none cursor-pointer" + /> +
+
+
)}
- ) : ( -
-

- Enter a search query above to get started -

-
- )} -
+
+ )} + + {/* Empty State */} + {!searchPerformed && ( +
+

+ Enter a search query above to get started +

+
+ )}
diff --git a/src/api/search.py b/src/api/search.py index 79f01274..386bb840 100644 --- a/src/api/search.py +++ b/src/api/search.py @@ -8,9 +8,13 @@ async def search(request: Request, search_service, session_manager): if not query: return JSONResponse({"error": "Query is required"}, status_code=400) + filters = payload.get("filters", {}) # Optional filters, defaults to empty dict + limit = payload.get("limit", 10) # Optional limit, defaults to 10 + score_threshold = payload.get("scoreThreshold", 0) # Optional score threshold, defaults to 0 + user = request.state.user # Extract JWT token from cookie for OpenSearch OIDC auth jwt_token = request.cookies.get("auth_token") - result = await search_service.search(query, user_id=user.user_id, jwt_token=jwt_token) + result = await search_service.search(query, user_id=user.user_id, jwt_token=jwt_token, filters=filters, limit=limit, score_threshold=score_threshold) return JSONResponse(result) \ No newline at end of file diff --git a/src/auth_context.py b/src/auth_context.py index 25647eab..3d9ec1f0 100644 --- a/src/auth_context.py +++ b/src/auth_context.py @@ -3,11 +3,14 @@ Authentication context for tool functions. Uses contextvars to safely pass user auth info through async calls. """ from contextvars import ContextVar -from typing import Optional +from typing import Optional, Dict, Any # Context variables for current request authentication _current_user_id: ContextVar[Optional[str]] = ContextVar('current_user_id', default=None) _current_jwt_token: ContextVar[Optional[str]] = ContextVar('current_jwt_token', default=None) +_current_search_filters: ContextVar[Optional[Dict[str, Any]]] = ContextVar('current_search_filters', default=None) +_current_search_limit: ContextVar[Optional[int]] = ContextVar('current_search_limit', default=10) +_current_score_threshold: ContextVar[Optional[float]] = ContextVar('current_score_threshold', default=0) def set_auth_context(user_id: str, jwt_token: str): """Set authentication context for the current async context""" @@ -24,4 +27,28 @@ def get_current_jwt_token() -> Optional[str]: def get_auth_context() -> tuple[Optional[str], Optional[str]]: """Get current authentication context (user_id, jwt_token)""" - return _current_user_id.get(), _current_jwt_token.get() \ No newline at end of file + return _current_user_id.get(), _current_jwt_token.get() + +def set_search_filters(filters: Dict[str, Any]): + """Set search filters for the current async context""" + _current_search_filters.set(filters) + +def get_search_filters() -> Optional[Dict[str, Any]]: + """Get current search filters from context""" + return _current_search_filters.get() + +def set_search_limit(limit: int): + """Set search limit for the current async context""" + _current_search_limit.set(limit) + +def get_search_limit() -> int: + """Get current search limit from context""" + return _current_search_limit.get() + +def set_score_threshold(threshold: float): + """Set score threshold for the current async context""" + _current_score_threshold.set(threshold) + +def get_score_threshold() -> float: + """Get current score threshold from context""" + return _current_score_threshold.get() \ No newline at end of file diff --git a/src/services/search_service.py b/src/services/search_service.py index 06e80bdb..b36df7b1 100644 --- a/src/services/search_service.py +++ b/src/services/search_service.py @@ -20,30 +20,99 @@ class SearchService: """ # Get authentication context from the current async context user_id, jwt_token = get_auth_context() + # Get search filters, limit, and score threshold from context + from auth_context import get_search_filters, get_search_limit, get_score_threshold + filters = get_search_filters() or {} + limit = get_search_limit() + score_threshold = get_score_threshold() # Embed the query resp = await clients.patched_async_client.embeddings.create(model=EMBED_MODEL, input=[query]) query_embedding = resp.data[0].embedding - # Base query structure + # Build filter clauses + filter_clauses = [] + if filters: + # Map frontend filter names to backend field names + field_mapping = { + "data_sources": "filename", + "document_types": "mimetype", + "owners": "owner" + } + + for filter_key, values in filters.items(): + if values is not None and isinstance(values, list): + # Map frontend key to backend field name + field_name = field_mapping.get(filter_key, filter_key) + + if len(values) == 0: + # Empty array means "match nothing" - use impossible filter + filter_clauses.append({"term": {field_name: "__IMPOSSIBLE_VALUE__"}}) + elif len(values) == 1: + # Single value filter + filter_clauses.append({"term": {field_name: values[0]}}) + else: + # Multiple values filter + filter_clauses.append({"terms": {field_name: values}}) + + # Hybrid search query structure (semantic + keyword) search_body = { "query": { "bool": { - "must": [ + "should": [ { "knn": { "chunk_embedding": { "vector": query_embedding, - "k": 10 + "k": 10, + "boost": 0.7 } } + }, + { + "multi_match": { + "query": query, + "fields": ["text^2", "filename^1.5"], + "type": "best_fields", + "fuzziness": "AUTO", + "boost": 0.3 + } } - ] + ], + "minimum_should_match": 1 + } + }, + "aggs": { + "data_sources": { + "terms": { + "field": "filename", + "size": 20 + } + }, + "document_types": { + "terms": { + "field": "mimetype", + "size": 10 + } + }, + "owners": { + "terms": { + "field": "owner", + "size": 10 + } } }, "_source": ["filename", "mimetype", "page", "text", "source_url", "owner", "allowed_users", "allowed_groups"], - "size": 10 + "size": limit } + # Add score threshold if specified + if score_threshold > 0: + search_body["min_score"] = score_threshold + + # Add filter clauses if any exist + if filter_clauses: + search_body["query"]["bool"]["filter"] = filter_clauses + # Authentication required - DLS will handle document filtering automatically if not user_id: return {"results": [], "error": "Authentication required"} @@ -52,7 +121,7 @@ class SearchService: opensearch_client = clients.create_user_opensearch_client(jwt_token) results = await opensearch_client.search(index=INDEX_NAME, body=search_body) - # Transform results + # Transform results (keep for backward compatibility) chunks = [] for hit in results["hits"]["hits"]: chunks.append({ @@ -64,13 +133,27 @@ class SearchService: "source_url": hit["_source"].get("source_url"), "owner": hit["_source"].get("owner") }) - return {"results": chunks} + + # Return both transformed results and aggregations + return { + "results": chunks, + "aggregations": results.get("aggregations", {}) + } - async def search(self, query: str, user_id: str = None, jwt_token: str = None) -> Dict[str, Any]: + async def search(self, query: str, user_id: str = None, jwt_token: str = None, filters: Dict[str, Any] = None, limit: int = 10, score_threshold: float = 0) -> Dict[str, Any]: """Public search method for API endpoints""" # Set auth context if provided (for direct API calls) if user_id and jwt_token: from auth_context import set_auth_context set_auth_context(user_id, jwt_token) + # Set filters and limit in context if provided + if filters: + from auth_context import set_search_filters + set_search_filters(filters) + + from auth_context import set_search_limit, set_score_threshold + set_search_limit(limit) + set_score_threshold(score_threshold) + return await self.search_tool(query) \ No newline at end of file