allow prioritizing prefixes

This commit is contained in:
2025-12-25 03:18:12 +01:00
parent a91fdb315c
commit 383e452322
2 changed files with 26 additions and 8 deletions

View File

@@ -59,7 +59,7 @@ def embed(text, input_type):
assert resp.embeddings.float_ is not None assert resp.embeddings.float_ is not None
return resp.embeddings.float_[0] 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.""" """Search with vector similarity, then rerank with Cohere for better relevance."""
query_embedding = embed(query, 'search_query') 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, query=query,
documents=[row['content'] or '' for row in rows], documents=[row['content'] or '' for row in rows],
model='rerank-v3.5', 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 = [] docs = []
for result in rerank_resp.results: for result in rerank_resp.results:
row = rows[result.index] 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 '' content = row['content'] or ''
if len(content) > max_content_length: if len(content) > max_content_length:
content = content[:max_content_length] + '...[truncated, use get_document_page for full text]' 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)}) docs.append({'key': row['key'], 'content': content, 'relevance': round(score, 3)})
return docs
# Re-sort by boosted score and return top_n
docs.sort(key=lambda x: x['relevance'], reverse=True)
return docs[:top_n]
@mcp.tool @mcp.tool
def get_cave_location(cave: str, state: str, county: str) -> list[dict]: 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) return search(f'{cave} Location, latitude, Longitude. Located in {state} and {county} county.', roles)
@mcp.tool @mcp.tool
def general_caving_information(query: str) -> list[dict]: def general_caving_information(query: str, priority_prefixes: list[str] | None = None) -> list[dict]:
"""General purpose search for any topic related to caves.""" """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() roles = get_user_roles()
return search(query, roles) return search(query, roles, priority_prefixes=priority_prefixes)
@mcp.tool @mcp.tool
def get_document_page(key: str) -> dict: def get_document_page(key: str) -> dict:

View File

@@ -85,7 +85,8 @@ Rules:
5. Be direct—no sycophantic phrases. 5. Be direct—no sycophantic phrases.
6. Keep responses concise. 6. Keep responses concise.
7. Use tools sparingly—one search usually suffices. 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.""" 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."""