simplify mcp
This commit is contained in:
@@ -59,14 +59,22 @@ 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, priority_prefixes: list[str] | None = None) -> list[dict]:
|
||||
"""Search with vector similarity, then rerank with Cohere for better relevance."""
|
||||
@mcp.tool
|
||||
def search_caving_documents(query: str, priority_prefixes: list[str] | None = None) -> dict:
|
||||
"""Search caving documents for information about caves, techniques, safety, accidents, history, and more.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
priority_prefixes: Optional list of key prefixes to prioritize (e.g., ['nss/aca'] for rescue topics)
|
||||
"""
|
||||
roles = get_user_roles()
|
||||
if not roles:
|
||||
return {"results": [], "note": "No results. Answer based on your knowledge."}
|
||||
|
||||
query_embedding = embed(query, 'search_query')
|
||||
|
||||
if not roles:
|
||||
return []
|
||||
|
||||
# Fetch more candidates for reranking
|
||||
top_n = 3
|
||||
candidate_limit = top_n * 4
|
||||
rows = conn.execute(
|
||||
'SELECT * FROM embeddings WHERE embedding IS NOT NULL AND role = ANY(%s) ORDER BY embedding <=> %s::vector LIMIT %s',
|
||||
@@ -74,14 +82,14 @@ def search(query, roles: list[str], top_n: int = 3, max_content_length: int = 15
|
||||
).fetchall()
|
||||
|
||||
if not rows:
|
||||
return []
|
||||
return {"results": [], "note": "No results found. Answer based on your knowledge."}
|
||||
|
||||
# Rerank with Cohere for better relevance
|
||||
rerank_resp = co.rerank(
|
||||
query=query,
|
||||
documents=[row['content'] or '' for row in rows],
|
||||
model='rerank-v3.5',
|
||||
top_n=min(top_n * 2, len(rows)), # Get more for re-sorting after boost
|
||||
top_n=min(top_n * 2, len(rows)),
|
||||
)
|
||||
|
||||
# Build results with optional priority boost
|
||||
@@ -94,49 +102,17 @@ def search(query, roles: list[str], top_n: int = 3, max_content_length: int = 15
|
||||
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
|
||||
score = min(1.0, score * 1.3)
|
||||
|
||||
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(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]:
|
||||
"""Lookup cave location as coordinates."""
|
||||
roles = get_user_roles()
|
||||
return search(f'{cave} Location, latitude, Longitude. Located in {state} and {county} county.', roles)
|
||||
|
||||
@mcp.tool
|
||||
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, priority_prefixes=priority_prefixes)
|
||||
|
||||
@mcp.tool
|
||||
def get_document_page(key: str) -> dict:
|
||||
"""Fetch full content for a document page. Pass the exact 'key' value from search results."""
|
||||
roles = get_user_roles()
|
||||
if not roles:
|
||||
return {"error": "No roles assigned"}
|
||||
|
||||
row = conn.execute(
|
||||
'SELECT key, content FROM embeddings WHERE key = %s AND role = ANY(%s)',
|
||||
(key, roles)
|
||||
).fetchone()
|
||||
|
||||
if row:
|
||||
return {"key": row["key"], "content": row["content"]}
|
||||
return {"error": f"Page not found: {key}"}
|
||||
return {
|
||||
"results": docs[:top_n],
|
||||
"note": "These are ALL available results. Do NOT search again - answer using these results now."
|
||||
}
|
||||
|
||||
@mcp.tool
|
||||
def get_user_info() -> dict:
|
||||
|
||||
Reference in New Issue
Block a user