restrict to dandisets in semantic search in chat tools

Author: magland

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changes
 
+## June 10, 2025
+- Added optional dandisets parameter to DANDI semantic search for filtering results to specific datasets
+
 ## June 5, 2025
 - Modernized Python package structure with pyproject.toml configuration. Removed legacy setup.py, setup.cfg, and setup.cfg.j2 files
 - Added option for using local neurosift server with the CLI

nextjs/neurosift-chat-agent-tools/src/app/api/dandi_semantic_search/route.ts

@@ -4,13 +4,13 @@ import { dandiSemanticSearch } from '@/lib/tools/dandi_semantic_search';
 export async function POST(request: NextRequest) {
   try {
     const body = await request.json();
-    const { query, limit } = body;
+    const { query, limit, dandisets } = body;
 
     if (!query) {
       return NextResponse.json({ error: 'Query is required' }, { status: 400 });
     }
 
-    const results = await dandiSemanticSearch(query, limit);
+    const results = await dandiSemanticSearch(query, limit, dandisets);
     return NextResponse.json({ results });
   } catch (error) {
     console.error('Error in DANDI semantic search:', error);

nextjs/neurosift-chat-agent-tools/src/app/api/tools/route.ts

@@ -56,6 +56,13 @@ The output is the same format as dandi_search.
           type: "number",
           description: "Maximum number of results",
           default: 10
+        },
+        dandisets: {
+          type: "array",
+          items: {
+            type: "string"
+          },
+          description: "Optional list of dandiset IDs to restrict search results to"
         }
       }
     }

nextjs/neurosift-chat-agent-tools/src/lib/tools/dandi_semantic_search.ts

@@ -83,10 +83,12 @@ const doEmbedding = async (text: string, model: string) => {
 
 const findSimilarDandisetIds = (
   embeddings: { [dandisetId: string]: number[] },
-  embedding: number[]
+  embedding: number[],
+  restrictToDandisets?: string[]
 ) => {
   const similarities: { dandisetId: string; similarity: number }[] = [];
   for (const dandisetId in embeddings) {
+    if (restrictToDandisets && !restrictToDandisets.includes(dandisetId)) continue;
     const embedding2 = embeddings[dandisetId];
     if (!embedding2) continue;
     const similarity = cosineSimilarity(embedding, embedding2);
@@ -118,7 +120,8 @@ interface DandiSearchResult {
 
 export async function dandiSemanticSearch(
   query: string,
-  limit: number = 10
+  limit: number = 10,
+  dandisets?: string[]
 ): Promise<DandiSearchResult[]> {
   const embeddings = await loadEmbeddings();
   if (!embeddings) {
@@ -131,7 +134,8 @@ export async function dandiSemanticSearch(
   }
   let similarDandisetIds = findSimilarDandisetIds(
     embeddings,
-    queryEmbedding.embedding
+    queryEmbedding.embedding,
+    dandisets
   );
   similarDandisetIds = similarDandisetIds.slice(0, limit);