allow prioritizing prefixes
This commit is contained in:
@@ -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:
|
||||||
|
|||||||
@@ -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."""
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user