Natural language filtering

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 <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
+  - 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]

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
+Given a user's search request:
-2. For each query, call the run_search tool
+1. ALWAYS run the original query first as-is
-3. After all searches complete, respond with "Search complete"
+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