11import json
2- from typing import AsyncGenerator
2+ from typing import AsyncGenerator , List
33
44from fastapi import APIRouter , HTTPException
55from fastapi .responses import StreamingResponse
66from loguru import logger
77
88from api .models import AskRequest , AskResponse , SearchRequest , SearchResponse
99from open_notebook .ai .models import Model , model_manager
10- from open_notebook .domain .notebook import text_search , vector_search
10+ from open_notebook .domain .notebook import (
11+ resolve_notebook_scope ,
12+ text_search ,
13+ vector_search ,
14+ )
1115from open_notebook .exceptions import (
1216 DatabaseOperationError ,
1317 InvalidInputError ,
2226async def search_knowledge_base (search_request : SearchRequest ):
2327 """Search the knowledge base using text or vector search."""
2428 try :
29+ notebook_ids = await resolve_notebook_scope (search_request .scope_notebook_ids )
30+
2531 if search_request .type == "vector" :
2632 # Check if embedding model is available for vector search
2733 if not await model_manager .get_embedding_model ():
@@ -36,6 +42,7 @@ async def search_knowledge_base(search_request: SearchRequest):
3642 source = search_request .search_sources ,
3743 note = search_request .search_notes ,
3844 minimum_score = search_request .minimum_score ,
45+ notebook_ids = notebook_ids ,
3946 )
4047 else :
4148 # Text search
@@ -44,6 +51,7 @@ async def search_knowledge_base(search_request: SearchRequest):
4451 results = search_request .limit ,
4552 source = search_request .search_sources ,
4653 note = search_request .search_notes ,
54+ notebook_ids = notebook_ids ,
4755 )
4856
4957 return SearchResponse (
@@ -67,7 +75,11 @@ async def search_knowledge_base(search_request: SearchRequest):
6775
6876
6977async def stream_ask_response (
70- question : str , strategy_model : Model , answer_model : Model , final_answer_model : Model
78+ question : str ,
79+ strategy_model : Model ,
80+ answer_model : Model ,
81+ final_answer_model : Model ,
82+ notebook_ids : List [str ],
7183) -> AsyncGenerator [str , None ]:
7284 """Stream the ask response as Server-Sent Events."""
7385 try :
@@ -76,7 +88,7 @@ async def stream_ask_response(
7688 # LangGraph accepts a partial state dict at runtime, but its typed
7789 # overloads require the full state type (langgraph typing limitation).
7890 async for chunk in ask_graph .astream ( # type: ignore[call-overload]
79- input = dict (question = question ),
91+ input = dict (question = question , notebook_ids = notebook_ids ),
8092 config = dict (
8193 configurable = dict (
8294 strategy_model = strategy_model .id ,
@@ -124,6 +136,10 @@ async def stream_ask_response(
124136async def ask_knowledge_base (ask_request : AskRequest ):
125137 """Ask the knowledge base a question using AI models."""
126138 try :
139+ # Cheapest check first: a malformed or unknown scope fails before any
140+ # model lookup or embedding check can mask it.
141+ notebook_ids = await resolve_notebook_scope (ask_request .scope_notebook_ids )
142+
127143 # Validate models exist
128144 strategy_model = await Model .get (ask_request .strategy_model )
129145 answer_model = await Model .get (ask_request .answer_model )
@@ -155,7 +171,11 @@ async def ask_knowledge_base(ask_request: AskRequest):
155171 # For streaming response
156172 return StreamingResponse (
157173 stream_ask_response (
158- ask_request .question , strategy_model , answer_model , final_answer_model
174+ ask_request .question ,
175+ strategy_model ,
176+ answer_model ,
177+ final_answer_model ,
178+ notebook_ids ,
159179 ),
160180 media_type = "text/event-stream" ,
161181 headers = {
@@ -178,6 +198,10 @@ async def ask_knowledge_base(ask_request: AskRequest):
178198async def ask_knowledge_base_simple (ask_request : AskRequest ):
179199 """Ask the knowledge base a question and return a simple response (non-streaming)."""
180200 try :
201+ # Cheapest check first: a malformed or unknown scope fails before any
202+ # model lookup or embedding check can mask it.
203+ notebook_ids = await resolve_notebook_scope (ask_request .scope_notebook_ids )
204+
181205 # Validate models exist
182206 strategy_model = await Model .get (ask_request .strategy_model )
183207 answer_model = await Model .get (ask_request .answer_model )
@@ -211,7 +235,7 @@ async def ask_knowledge_base_simple(ask_request: AskRequest):
211235 # LangGraph accepts a partial state dict at runtime, but its typed
212236 # overloads require the full state type (langgraph typing limitation).
213237 async for chunk in ask_graph .astream ( # type: ignore[call-overload]
214- input = dict (question = ask_request .question ),
238+ input = dict (question = ask_request .question , notebook_ids = notebook_ids ),
215239 config = dict (
216240 configurable = dict (
217241 strategy_model = strategy_model .id ,
0 commit comments