diff --git a/app/backend/agent.py b/app/backend/agent.py index 55e488f7..c839d120 100644 --- a/app/backend/agent.py +++ b/app/backend/agent.py @@ -87,6 +87,17 @@ How to decide which tool to use: - "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations. - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, just output the list of results returned by the tool verbatim. Do NOT summarize or add commentary. +IMPORTANT - When user mentions a document in search/ask: +- If user says "search in ", "find in ", "answer from ", or " in ": + - Extract the TOPIC as `query`/`question` + - Extract the DOCUMENT NAME as `document_name` +- Examples for search: + - "search for latrines in TB MED 593" → query="latrines", document_name="TB MED 593" + - "find waste disposal in the army manual" → query="waste disposal", document_name="army manual" +- Examples for ask: + - "what does TB MED 593 say about latrines?" → question="what are the guidelines for latrines?", document_name="TB MED 593" + - "answer from the army manual about sanitation" → question="what are the sanitation guidelines?", document_name="army manual" + Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user.""" @@ -105,6 +116,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: async def search( ctx: RunContext[ChatDeps], query: str, + document_name: str | None = None, ) -> str: """Search the knowledge base for relevant documents. @@ -112,21 +124,36 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: Results are displayed to the user - just list the titles found. Args: - query: The search query + query: The search query (what to search for) + document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual") """ from search_agent import SearchAgent if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.log(f"Searching: {query}") + msg = f"Searching: {query}" + if document_name: + msg += f" (in {document_name})" + ctx.deps.agui_emitter.log(msg) # Build context from conversation history context = None if ctx.deps.session_state and ctx.deps.session_state.qa_history: context = format_conversation_context(ctx.deps.session_state.qa_history) + # Build filter from document_name + doc_filter = None + if document_name: + escaped = document_name.replace("'", "''") + # Also try without spaces for matching "TB MED 593" to "tbmed593" + no_spaces = escaped.replace(" ", "") + doc_filter = ( + f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') " + f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')" + ) + # Use search agent for query expansion and deduplication search_agent = SearchAgent(ctx.deps.client, ctx.deps.config) - results = await search_agent.search(query, context=context) + results = await search_agent.search(query, context=context, filter=doc_filter) # Store for potential citation resolution ctx.deps.search_results = results @@ -178,7 +205,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: async def ask( ctx: RunContext[ChatDeps], question: str, - document_filter: str | None = None, + document_name: str | None = None, ) -> str: """Answer a specific question using the knowledge base. @@ -186,10 +213,24 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: Args: question: The question to answer - document_filter: Optional SQL WHERE clause to filter documents (e.g. "id IN ('doc1', 'doc2')") + document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual") """ if ctx.deps.agui_emitter: - ctx.deps.agui_emitter.log(f"Answering: {question}") + msg = f"Answering: {question}" + if document_name: + msg += f" (in {document_name})" + ctx.deps.agui_emitter.log(msg) + + # Build filter from document_name + doc_filter = None + if document_name: + escaped = document_name.replace("'", "''") + # Also try without spaces for matching "TB MED 593" to "tbmed593" + no_spaces = escaped.replace(" ", "") + doc_filter = ( + f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') " + f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')" + ) # Build context-aware system prompt if we have history system_prompt = None @@ -205,7 +246,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: ) answer, citations = await ctx.deps.client.ask( - question, system_prompt=system_prompt, filter=document_filter + question, system_prompt=system_prompt, filter=doc_filter ) # Accumulate Q&A in session state @@ -283,19 +324,23 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: doc = await ctx.deps.client.get_document_by_uri(query) escaped_query = query.replace("'", "''") + # Also try without spaces for matching "TB MED 593" to "tbmed593" + no_spaces = escaped_query.replace(" ", "") - # If not found, try partial URI match + # If not found, try partial URI match (with and without spaces) if doc is None: docs = await ctx.deps.client.list_documents( - limit=1, filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%')" + limit=1, + filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%') OR LOWER(uri) LIKE LOWER('%{no_spaces}%')", ) if docs: doc = docs[0] - # If still not found, try partial title match + # If still not found, try partial title match (with and without spaces) if doc is None: docs = await ctx.deps.client.list_documents( - limit=1, filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%')" + limit=1, + filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')", ) if docs: doc = docs[0] diff --git a/app/backend/search_agent.py b/app/backend/search_agent.py index abbef6d0..cd52cc13 100644 --- a/app/backend/search_agent.py +++ b/app/backend/search_agent.py @@ -14,16 +14,22 @@ class SearchDeps: client: HaikuRAG config: AppConfig + filter: str | None = None search_results: list[SearchResult] = field(default_factory=list) -SEARCH_SYSTEM_PROMPT = """You are a search query optimizer. Given a user's search request: +SEARCH_SYSTEM_PROMPT = """You are a search query optimizer for a document knowledge base. -1. Generate 2-4 diverse search queries that cover different aspects/phrasings of the request -2. For each query, call the run_search tool -3. After all searches complete, respond with "Search complete" +Given a user's search request: +1. ALWAYS run the original query first as-is +2. Then generate 1-2 alternative queries using different keywords or phrasings +3. Keep queries SHORT (2-5 words) - use keywords, not full sentences +4. After all searches, respond with "Search complete" -Be thorough but focused. Generate queries that will find relevant results without being redundant.""" +Example: User asks "latrines" → queries: "latrines", "latrine sanitation", "field toilet" +Example: User asks "waste disposal" → queries: "waste disposal", "garbage management", "refuse handling" + +Do NOT generate long verbose queries like "environmental impact of waste disposal methods" - keep it simple.""" class SearchAgent: @@ -52,7 +58,9 @@ class SearchAgent: query: The search query """ limit = ctx.deps.config.search.limit - results = await ctx.deps.client.search(query, limit=limit) + results = await ctx.deps.client.search( + query, limit=limit, filter=ctx.deps.filter + ) results = await ctx.deps.client.expand_context(results) ctx.deps.search_results.extend(results) @@ -64,12 +72,14 @@ class SearchAgent: self, query: str, context: str | None = None, + filter: str | None = None, ) -> list[SearchResult]: """Execute search with query expansion and deduplication. Args: query: The user's search request context: Optional conversation context + filter: Optional SQL WHERE clause to filter documents Returns: Deduplicated list of SearchResult sorted by score @@ -78,7 +88,7 @@ class SearchAgent: if context: prompt = f"Context: {context}\n\nSearch request: {query}" - deps = SearchDeps(client=self._client, config=self._config) + deps = SearchDeps(client=self._client, config=self._config, filter=filter) await self._agent.run(prompt, deps=deps) # Deduplicate by chunk_id, keeping highest score