feat: enhance session search tool with parent session resolution and parallel summarization
- Added a new function to resolve child sessions to their parent, improving session grouping and deduplication. - Refactored session summarization to run in parallel, enhancing performance and responsiveness. - Updated search syntax documentation to clarify usage of keywords and phrases for better search results.
This commit is contained in:
parent
c1d9e9a285
commit
cc6bea8b90
1 changed files with 61 additions and 36 deletions
|
|
@ -224,54 +224,77 @@ def session_search(
|
||||||
"message": "No matching sessions found.",
|
"message": "No matching sessions found.",
|
||||||
}, ensure_ascii=False)
|
}, ensure_ascii=False)
|
||||||
|
|
||||||
# Group by session_id, keep order (highest ranked first)
|
# Resolve child sessions to their parent — delegation stores detailed
|
||||||
|
# content in child sessions, but the user's conversation is the parent.
|
||||||
|
def _resolve_to_parent(session_id):
|
||||||
|
visited = set()
|
||||||
|
sid = session_id
|
||||||
|
while sid and sid not in visited:
|
||||||
|
visited.add(sid)
|
||||||
|
session = db.get_session(sid)
|
||||||
|
if not session:
|
||||||
|
break
|
||||||
|
parent = session.get("parent_session_id")
|
||||||
|
if parent:
|
||||||
|
sid = parent
|
||||||
|
else:
|
||||||
|
break
|
||||||
|
return sid
|
||||||
|
|
||||||
|
# Group by resolved (parent) session_id, dedup
|
||||||
seen_sessions = {}
|
seen_sessions = {}
|
||||||
for result in raw_results:
|
for result in raw_results:
|
||||||
sid = result["session_id"]
|
raw_sid = result["session_id"]
|
||||||
if sid not in seen_sessions:
|
resolved_sid = _resolve_to_parent(raw_sid)
|
||||||
seen_sessions[sid] = result
|
if resolved_sid not in seen_sessions:
|
||||||
|
result = dict(result)
|
||||||
|
result["session_id"] = resolved_sid
|
||||||
|
seen_sessions[resolved_sid] = result
|
||||||
if len(seen_sessions) >= limit:
|
if len(seen_sessions) >= limit:
|
||||||
break
|
break
|
||||||
|
|
||||||
# Summarize each matching session
|
# Prepare all sessions for parallel summarization
|
||||||
summaries = []
|
tasks = []
|
||||||
for session_id, match_info in seen_sessions.items():
|
for session_id, match_info in seen_sessions.items():
|
||||||
try:
|
try:
|
||||||
# Load full conversation
|
|
||||||
messages = db.get_messages_as_conversation(session_id)
|
messages = db.get_messages_as_conversation(session_id)
|
||||||
if not messages:
|
if not messages:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Get session metadata
|
|
||||||
session_meta = db.get_session(session_id) or {}
|
session_meta = db.get_session(session_id) or {}
|
||||||
|
|
||||||
# Format and truncate
|
|
||||||
conversation_text = _format_conversation(messages)
|
conversation_text = _format_conversation(messages)
|
||||||
conversation_text = _truncate_around_matches(conversation_text, query)
|
conversation_text = _truncate_around_matches(conversation_text, query)
|
||||||
|
tasks.append((session_id, match_info, conversation_text, session_meta))
|
||||||
# Summarize with Gemini Flash (handle both async and sync contexts)
|
|
||||||
coro = _summarize_session(conversation_text, query, session_meta)
|
|
||||||
try:
|
|
||||||
asyncio.get_running_loop()
|
|
||||||
# Already in an async context (gateway) -- run in a thread
|
|
||||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
|
||||||
summary = pool.submit(lambda: asyncio.run(coro)).result(timeout=30)
|
|
||||||
except RuntimeError:
|
|
||||||
# No running loop (normal CLI) -- use asyncio.run directly
|
|
||||||
summary = asyncio.run(coro)
|
|
||||||
|
|
||||||
if summary:
|
|
||||||
summaries.append({
|
|
||||||
"session_id": session_id,
|
|
||||||
"when": _format_timestamp(match_info.get("session_started")),
|
|
||||||
"source": match_info.get("source", "unknown"),
|
|
||||||
"model": match_info.get("model"),
|
|
||||||
"summary": summary,
|
|
||||||
})
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logging.warning(f"Failed to summarize session {session_id}: {e}")
|
logging.warning(f"Failed to prepare session {session_id}: {e}")
|
||||||
|
|
||||||
|
# Summarize all sessions in parallel
|
||||||
|
async def _summarize_all():
|
||||||
|
coros = [
|
||||||
|
_summarize_session(text, query, meta)
|
||||||
|
for _, _, text, meta in tasks
|
||||||
|
]
|
||||||
|
return await asyncio.gather(*coros, return_exceptions=True)
|
||||||
|
|
||||||
|
try:
|
||||||
|
asyncio.get_running_loop()
|
||||||
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
||||||
|
results = pool.submit(lambda: asyncio.run(_summarize_all())).result(timeout=60)
|
||||||
|
except RuntimeError:
|
||||||
|
results = asyncio.run(_summarize_all())
|
||||||
|
|
||||||
|
summaries = []
|
||||||
|
for (session_id, match_info, _, _), result in zip(tasks, results):
|
||||||
|
if isinstance(result, Exception):
|
||||||
|
logging.warning(f"Failed to summarize session {session_id}: {result}")
|
||||||
continue
|
continue
|
||||||
|
if result:
|
||||||
|
summaries.append({
|
||||||
|
"session_id": session_id,
|
||||||
|
"when": _format_timestamp(match_info.get("session_started")),
|
||||||
|
"source": match_info.get("source", "unknown"),
|
||||||
|
"model": match_info.get("model"),
|
||||||
|
"summary": result,
|
||||||
|
})
|
||||||
|
|
||||||
return json.dumps({
|
return json.dumps({
|
||||||
"success": True,
|
"success": True,
|
||||||
|
|
@ -309,9 +332,11 @@ SESSION_SEARCH_SCHEMA = {
|
||||||
"- The user asks 'what did we do about X?' or 'how did we fix Y?'\n\n"
|
"- The user asks 'what did we do about X?' or 'how did we fix Y?'\n\n"
|
||||||
"Don't hesitate to search -- it's fast and cheap. Better to search and confirm "
|
"Don't hesitate to search -- it's fast and cheap. Better to search and confirm "
|
||||||
"than to guess or ask the user to repeat themselves.\n\n"
|
"than to guess or ask the user to repeat themselves.\n\n"
|
||||||
"Search syntax: keywords (docker deployment), phrases (\"exact match\"), "
|
"Search syntax: keywords joined with OR for broad recall (elevenlabs OR baseten OR funding), "
|
||||||
"boolean (python NOT java), prefix (deploy*). Returns summaries of the "
|
"phrases for exact match (\"docker networking\"), boolean (python NOT java), prefix (deploy*). "
|
||||||
"top matching sessions focused on your query."
|
"IMPORTANT: Use OR between keywords for best results — FTS5 defaults to AND which misses "
|
||||||
|
"sessions that only mention some terms. If a broad OR query returns nothing, try individual "
|
||||||
|
"keyword searches in parallel. Returns summaries of the top matching sessions."
|
||||||
),
|
),
|
||||||
"parameters": {
|
"parameters": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue