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
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:

View File

@@ -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."""