From 383e452322413954b1e068e126c19b5dad0a4945 Mon Sep 17 00:00:00 2001
From: Paul Walko <paul@bigcavemaps.com>
Date: Thu, 25 Dec 2025 03:18:12 +0100
Subject: [PATCH] allow prioritizing prefixes

---
 mcp/server.py          | 31 ++++++++++++++++++++++++-------
 web/agent/src/agent.py |  3 ++-
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/mcp/server.py b/mcp/server.py
index 41bb9ec..e1a8e7d 100644
--- a/mcp/server.py
+++ b/mcp/server.py
@@ -59,7 +59,7 @@ def embed(text, input_type):
     assert resp.embeddings.float_ is not None
     return resp.embeddings.float_[0]
 
-def search(query, roles: list[str], top_n: int = 3, max_content_length: int = 1500) -> list[dict]:
+def search(query, roles: list[str], top_n: int = 3, max_content_length: int = 1500, priority_prefixes: list[str] | None = None) -> list[dict]:
     """Search with vector similarity, then rerank with Cohere for better relevance."""
     query_embedding = embed(query, 'search_query')
 
@@ -81,17 +81,29 @@ def search(query, roles: list[str], top_n: int = 3, max_content_length: int = 15
         query=query,
         documents=[row['content'] or '' for row in rows],
         model='rerank-v3.5',
-        top_n=top_n,
+        top_n=min(top_n * 2, len(rows)),  # Get more for re-sorting after boost
     )
 
+    # Build results with optional priority boost
     docs = []
     for result in rerank_resp.results:
         row = rows[result.index]
+        score = result.relevance_score
+
+        # Boost score if key starts with any priority prefix (e.g., 'nss/aca')
+        if priority_prefixes:
+            key = row['key'] or ''
+            if any(key.startswith(prefix) for prefix in priority_prefixes):
+                score = min(1.0, score * 1.3)  # 30% boost, capped at 1.0
+
         content = row['content'] or ''
         if len(content) > max_content_length:
             content = content[:max_content_length] + '...[truncated, use get_document_page for full text]'
-        docs.append({'key': row['key'], 'content': content, 'relevance': round(result.relevance_score, 3)})
-    return docs
+        docs.append({'key': row['key'], 'content': content, 'relevance': round(score, 3)})
+
+    # Re-sort by boosted score and return top_n
+    docs.sort(key=lambda x: x['relevance'], reverse=True)
+    return docs[:top_n]
 
 @mcp.tool
 def get_cave_location(cave: str, state: str, county: str) -> list[dict]:
@@ -100,10 +112,15 @@ def get_cave_location(cave: str, state: str, county: str) -> list[dict]:
     return search(f'{cave} Location, latitude, Longitude. Located in {state} and {county} county.', roles)
 
 @mcp.tool
-def general_caving_information(query: str) -> list[dict]:
-    """General purpose search for any topic related to caves."""
+def general_caving_information(query: str, priority_prefixes: list[str] | None = None) -> list[dict]:
+    """General purpose search for any topic related to caves.
+
+    Args:
+        query: Search query
+        priority_prefixes: Optional list of key prefixes to prioritize in results (e.g., ['nss/aca'] for rescue topics)
+    """
     roles = get_user_roles()
-    return search(query, roles)
+    return search(query, roles, priority_prefixes=priority_prefixes)
 
 @mcp.tool
 def get_document_page(key: str) -> dict:
diff --git a/web/agent/src/agent.py b/web/agent/src/agent.py
index ab7f2e6..973349a 100644
--- a/web/agent/src/agent.py
+++ b/web/agent/src/agent.py
@@ -85,7 +85,8 @@ Rules:
 5. Be direct—no sycophantic phrases.
 6. Keep responses concise.
 7. Use tools sparingly—one search usually suffices.
-8. If you hit the search limit, end your reply with an italicized note: *Your question may be too broad. Try asking something more specific.* Do NOT mention "tools" or "tool limits"—the user doesn't know what those are."""
+8. If you hit the search limit, end your reply with an italicized note: *Your question may be too broad. Try asking something more specific.* Do NOT mention "tools" or "tool limits"—the user doesn't know what those are.
+9. For rescue, accident, or emergency-related queries, use priority_prefixes=['nss/aca'] when searching to prioritize official accident reports."""
 
 SOURCES_ONLY_INSTRUCTIONS = """SOURCES ONLY MODE: Give exactly ONE sentence summary. Then list sources with specific page numbers (e.g., "- The Trog 2021, page 19"). No explanations."""