Jodinho commited on
Commit
0dfff2c
Β·
1 Parent(s): 7841ad6

perf: reduce token consumption by ~80% & fix synthesis truncation

Browse files

1. Tool-Set Pruning:
- Replaced full 18-tool array (~3,130 tokens) on every call with intent-based subsets (5-7 tools = ~700-900 tokens).
- Added select_tools(query) covering MARKET_TOOLS, ANOMALY_TOOLS, PROPERTY_TOOLS, and GEO_TOOLS.
- All subsets include generate_data_export.

2. Synthesis Tool Stripping & Dynamic Token Budget:
- Tool-call rounds send selected tools with max_tokens=300 (tool JSONs are small).
- Synthesis round strips all tools completely (saving 3,130 tokens on final generation) and raises max_tokens to 2,000 to prevent report truncation (resolves finish_reason='length').

3. Session History Fact Compression:
- Implemented _compress_for_history() to store factual text summaries in session history while delivering full rich markdown tables to users.
- Saves ~85% history tokens per turn.

4. Hard Cap on Tool Result Payloads:
- Lowered max returned rows from 50 to 30 and enforced a 2,000-character ceiling per tool result in compress_tool_output().

services/agent_loop.py CHANGED
@@ -44,7 +44,7 @@ from services.tool_executor import (
44
  )
45
  from services.tool_compressor import compress_tool_output
46
  from services.supabase_service import search_methodology_rag
47
- from services.tools import REAL_ESTATE_TOOLS, COMMERCIAL_TOOLS
48
 
49
  logger = setup_logger(__name__)
50
 
@@ -175,14 +175,24 @@ async def run_agent_loop(
175
  response_message = None
176
 
177
  # ── Brain call (primary + fallback) ───────────────────────────────────
 
 
 
 
 
 
 
 
 
 
178
  for model in [GROQ_SYNTHESIS_MODEL, GROQ_FALLBACK_SYNTHESIS_MODEL]:
179
  try:
180
  brain_res = await groq_call(
181
  model=model,
182
  messages=messages,
183
- tools=active_tools if active_tools else None,
184
- tool_choice="auto" if active_tools else "none",
185
- max_tokens=600 if active_tools else 1000,
186
  temperature=0.2,
187
  )
188
  response_message = brain_res.choices[0].message
@@ -268,7 +278,7 @@ async def run_agent_loop(
268
  messages=synth_messages,
269
  tools=None,
270
  tool_choice="none",
271
- max_tokens=1000,
272
  temperature=0.2,
273
  )
274
  return final_res.choices[0].message.content
@@ -315,14 +325,19 @@ async def run_agent_loop_streaming(
315
  for round_num in range(1, MAX_TOOL_ROUNDS + 1):
316
  response_message = None
317
 
 
 
 
 
 
318
  for model in [GROQ_SYNTHESIS_MODEL, GROQ_FALLBACK_SYNTHESIS_MODEL]:
319
  try:
320
  brain_res = await groq_call(
321
  model=model,
322
  messages=messages,
323
- tools=active_tools if active_tools else None,
324
- tool_choice="auto" if active_tools else "none",
325
- max_tokens=600,
326
  temperature=0.2,
327
  )
328
  response_message = brain_res.choices[0].message
@@ -453,6 +468,31 @@ def extract_clarification_options(reply: str) -> tuple[str, list[str]]:
453
  return reply, suggested_actions
454
 
455
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
456
  # ─── PUBLIC ENTRY POINT ───────────────────────────────────────────────────────
457
 
458
 
@@ -513,10 +553,12 @@ async def process_chat_message(
513
  }
514
  )
515
 
516
- # STEP 4: Select tools
 
517
  active_tools = None
518
  if classification in ("PATH_A", "BOTH"):
519
- active_tools = REAL_ESTATE_TOOLS
 
520
  elif classification == "COMMERCIAL_HANDOFF":
521
  active_tools = COMMERCIAL_TOOLS
522
 
@@ -531,8 +573,12 @@ async def process_chat_message(
531
  ):
532
  final_reply = _REPLY_FALLBACK
533
 
534
- # STEP 6: Store assistant reply in session
535
- append_message(session_id, {"role": "assistant", "content": final_reply})
 
 
 
 
536
 
537
  # STEP 7: Strip clarification options from reply text
538
  final_reply, suggested_actions = extract_clarification_options(final_reply)
@@ -605,10 +651,11 @@ async def stream_chat_message(
605
  }
606
  )
607
 
608
- # STEP 4: Tools
609
  active_tools = None
610
  if classification in ("PATH_A", "BOTH"):
611
- active_tools = REAL_ESTATE_TOOLS
 
612
  elif classification == "COMMERCIAL_HANDOFF":
613
  active_tools = COMMERCIAL_TOOLS
614
 
@@ -623,7 +670,10 @@ async def stream_chat_message(
623
  full_reply_parts.append(event["token"])
624
  yield event
625
 
626
- # STEP 6: Persist assembled reply to session
627
  full_reply = "".join(full_reply_parts)
628
  if full_reply:
629
- append_message(session_id, {"role": "assistant", "content": full_reply})
 
 
 
 
44
  )
45
  from services.tool_compressor import compress_tool_output
46
  from services.supabase_service import search_methodology_rag
47
+ from services.tools import REAL_ESTATE_TOOLS, COMMERCIAL_TOOLS, select_tools
48
 
49
  logger = setup_logger(__name__)
50
 
 
175
  response_message = None
176
 
177
  # ── Brain call (primary + fallback) ───────────────────────────────────
178
+ # Synthesis detection: strip tools + raise max_tokens when we already
179
+ # have tool results and the last message is a tool result. The model
180
+ # has all data it needs and should write the reply, not call more tools.
181
+ last_is_tool = bool(messages) and messages[-1].get('role') == 'tool'
182
+ has_data = bool(tool_results) and last_is_tool
183
+ tools_now = None if has_data else active_tools
184
+ choice_now = 'none' if has_data else ('auto' if active_tools else 'none')
185
+ tokens_now = 2000 if has_data else 300 # tool-call JSON is <30 tokens
186
+ if has_data:
187
+ logger.info('[agent_loop] synthesis mode: tools stripped, max_tokens=2000')
188
  for model in [GROQ_SYNTHESIS_MODEL, GROQ_FALLBACK_SYNTHESIS_MODEL]:
189
  try:
190
  brain_res = await groq_call(
191
  model=model,
192
  messages=messages,
193
+ tools=tools_now,
194
+ tool_choice=choice_now,
195
+ max_tokens=tokens_now,
196
  temperature=0.2,
197
  )
198
  response_message = brain_res.choices[0].message
 
278
  messages=synth_messages,
279
  tools=None,
280
  tool_choice="none",
281
+ max_tokens=2000,
282
  temperature=0.2,
283
  )
284
  return final_res.choices[0].message.content
 
325
  for round_num in range(1, MAX_TOOL_ROUNDS + 1):
326
  response_message = None
327
 
328
+ last_is_tool = bool(messages) and messages[-1].get('role') == 'tool'
329
+ has_data = bool(tool_results) and last_is_tool
330
+ tools_now = None if has_data else active_tools
331
+ choice_now = 'none' if has_data else ('auto' if active_tools else 'none')
332
+ tokens_now = 2000 if has_data else 300
333
  for model in [GROQ_SYNTHESIS_MODEL, GROQ_FALLBACK_SYNTHESIS_MODEL]:
334
  try:
335
  brain_res = await groq_call(
336
  model=model,
337
  messages=messages,
338
+ tools=tools_now,
339
+ tool_choice=choice_now,
340
+ max_tokens=tokens_now,
341
  temperature=0.2,
342
  )
343
  response_message = brain_res.choices[0].message
 
468
  return reply, suggested_actions
469
 
470
 
471
+
472
+ # ─── HISTORY COMPRESSOR ───────────────────────────────────────────────────────
473
+
474
+ MAX_HISTORY_CHARS = 400 # Per-turn cap for stored assistant replies
475
+
476
+
477
+ def _compress_for_history(reply: str) -> str:
478
+ """
479
+ Strip heavy markdown formatting from an assistant reply before storing in session.
480
+ Preserves table cell contents and key factual values while removing table borders,
481
+ formatting markup, code fences, and redundant whitespace.
482
+ """
483
+ # Remove table divider lines like |---|---| or |:---:|
484
+ text = re.sub(r"\|[\s\-:]+\|[\s\-:|]*", " ", reply)
485
+ # Replace table column separator pipes with commas/spaces
486
+ text = re.sub(r"\|", " , ", text)
487
+ # Remove markdown headers, bold, italic, code fences, blockquotes, bullets
488
+ text = re.sub(r"[*#`>\-~_]+", " ", text)
489
+ # Collapse whitespace
490
+ text = re.sub(r"\s+", " ", text).strip()
491
+ if len(text) > MAX_HISTORY_CHARS:
492
+ return text[:MAX_HISTORY_CHARS] + "\u2026"
493
+ return text
494
+
495
+
496
  # ─── PUBLIC ENTRY POINT ───────────────────────────────────────────────────────
497
 
498
 
 
553
  }
554
  )
555
 
556
+ # STEP 4: Select tools β€” use keyword-based subset to minimize tool definition tokens
557
+ # select_tools() picks 4-6 relevant tools instead of sending all 18 (~3,130 tokens)
558
  active_tools = None
559
  if classification in ("PATH_A", "BOTH"):
560
+ active_tools = select_tools(user_query)
561
+ logger.info(f"[process_chat] selected {len(active_tools)} tools for classification={classification}")
562
  elif classification == "COMMERCIAL_HANDOFF":
563
  active_tools = COMMERCIAL_TOOLS
564
 
 
573
  ):
574
  final_reply = _REPLY_FALLBACK
575
 
576
+ # STEP 6: Store compressed reply in session (full reply returned to user)
577
+ # Storing full markdown in history triples token cost per turn.
578
+ append_message(session_id, {
579
+ "role": "assistant",
580
+ "content": _compress_for_history(final_reply),
581
+ })
582
 
583
  # STEP 7: Strip clarification options from reply text
584
  final_reply, suggested_actions = extract_clarification_options(final_reply)
 
651
  }
652
  )
653
 
654
+ # STEP 4: Tools β€” keyword-based subset selection
655
  active_tools = None
656
  if classification in ("PATH_A", "BOTH"):
657
+ active_tools = select_tools(user_query)
658
+ logger.info(f"[stream_chat] selected {len(active_tools)} tools for classification={classification}")
659
  elif classification == "COMMERCIAL_HANDOFF":
660
  active_tools = COMMERCIAL_TOOLS
661
 
 
670
  full_reply_parts.append(event["token"])
671
  yield event
672
 
673
+ # STEP 6: Persist compressed reply to session (full text was streamed to user)
674
  full_reply = "".join(full_reply_parts)
675
  if full_reply:
676
+ append_message(session_id, {
677
+ "role": "assistant",
678
+ "content": _compress_for_history(full_reply),
679
+ })
services/tool_compressor.py CHANGED
@@ -19,7 +19,8 @@ from config import MAPBOX_ACCESS_TOKEN
19
 
20
  logger = setup_logger(__name__)
21
 
22
- _MAX_ROWS = 50 # Maximum rows delivered to LLM from any single tool call
 
23
 
24
 
25
  def compress_tool_output(func_name: str, db_result: dict) -> str:
@@ -97,7 +98,20 @@ def compress_tool_output(func_name: str, db_result: dict) -> str:
97
  f"\n[Truncated to {_MAX_ROWS} rows. Advise user to narrow date range or add filters.]"
98
  if truncated else ""
99
  )
100
- return f"{header_str}{buf.getvalue().strip()}{truncation_notice}"
 
 
 
 
 
 
 
 
 
 
 
 
 
101
 
102
 
103
  # ─── GEOCODE HANDLER (Mapbox API) ─────────────────────────────────────────────
 
19
 
20
  logger = setup_logger(__name__)
21
 
22
+ _MAX_ROWS = 30 # Maximum rows per tool result (was 50)
23
+ _MAX_RESULT_CHARS = 2000 # Hard char cap per compressed result (~500 tokens)
24
 
25
 
26
  def compress_tool_output(func_name: str, db_result: dict) -> str:
 
98
  f"\n[Truncated to {_MAX_ROWS} rows. Advise user to narrow date range or add filters.]"
99
  if truncated else ""
100
  )
101
+ result_str = f"{header_str}{buf.getvalue().strip()}{truncation_notice}"
102
+ # Hard cap: prevent any single tool result from exceeding budget
103
+ if len(result_str) > _MAX_RESULT_CHARS:
104
+ lines = result_str.split("\n")
105
+ capped = []
106
+ total = 0
107
+ for line in lines:
108
+ if total + len(line) > _MAX_RESULT_CHARS:
109
+ capped.append("[Result capped at 2000 chars. Advise user to narrow filters.]")
110
+ break
111
+ capped.append(line)
112
+ total += len(line)
113
+ result_str = "\n".join(capped)
114
+ return result_str
115
 
116
 
117
  # ─── GEOCODE HANDLER (Mapbox API) ─────────────────────────────────────────────
services/tools.py CHANGED
@@ -470,3 +470,80 @@ COMMERCIAL_TOOLS = [
470
  },
471
  }
472
  ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
470
  },
471
  }
472
  ]
473
+
474
+
475
+ # ─── STATIC TOOL GROUPINGS (Fix #4 β€” Tool-Set Pruning) ───────────────────────
476
+ #
477
+ # Instead of sending all 18 tool definitions on every call (~3,130 tokens),
478
+ # the orchestrator picks the smallest group that covers the query.
479
+ #
480
+ # Grouping logic (chosen in agent_loop.py based on query keywords):
481
+ #
482
+ # PATH_A queries break into 3 intent categories:
483
+ # MARKET_TOOLS β€” "how does the market look", "averages", "trends"
484
+ # ANOMALY_TOOLS β€” "spikes", "crashes", "anomalies", "volatile"
485
+ # PROPERTY_TOOLS β€” "property", "listing", "compare", "availability"
486
+ # GEO_TOOLS β€” "near", "address", "distance", "location"
487
+ #
488
+ # Each group includes generate_data_export so the model can always offer downloads.
489
+ # Any group = ~4-6 tools = ~700-900 tokens (vs 3,130 for all 18).
490
+ #
491
+ # For BOTH / ambiguous queries, use FULL_REAL_ESTATE_TOOLS as fallback.
492
+
493
+ _EXPORT_TOOL = [t for t in REAL_ESTATE_TOOLS if t["function"]["name"] == "generate_data_export"]
494
+ _MARKET_NAMES = {"get_market_averages", "get_market_snapshot", "get_market_trend",
495
+ "get_tracked_markets", "get_dashboard_kpis"}
496
+ _ANOMALY_NAMES = {"get_spike_alerts", "get_rate_anomaly_report", "get_most_volatile_properties",
497
+ "get_dashboard_kpis", "get_recently_changed_tracking"}
498
+ _PROPERTY_NAMES = {"get_property_snapshot", "get_property_rate_changes", "compare_properties",
499
+ "search_properties", "get_availability_rate", "get_dashboard_kpis"}
500
+ _GEO_NAMES = {"geocode_address", "get_nearby_properties", "get_distance_km",
501
+ "search_properties"}
502
+
503
+ # Build subsets once at import time
504
+ MARKET_TOOLS = [t for t in REAL_ESTATE_TOOLS if t["function"]["name"] in _MARKET_NAMES] + _EXPORT_TOOL
505
+ ANOMALY_TOOLS = [t for t in REAL_ESTATE_TOOLS if t["function"]["name"] in _ANOMALY_NAMES] + _EXPORT_TOOL
506
+ PROPERTY_TOOLS = [t for t in REAL_ESTATE_TOOLS if t["function"]["name"] in _PROPERTY_NAMES] + _EXPORT_TOOL
507
+ GEO_TOOLS = [t for t in REAL_ESTATE_TOOLS if t["function"]["name"] in _GEO_NAMES] + _EXPORT_TOOL
508
+
509
+ # Keywords used by agent_loop.select_tools() to pick the right subset
510
+ _ANOMALY_KEYWORDS = {"spike", "spikes", "crash", "anomaly", "anomalies", "volatile",
511
+ "changes", "rate change", "increase", "decrease", "shift", "spike report",
512
+ "threshold", "exceed", "25%", "unusual"}
513
+ _PROPERTY_KEYWORDS = {"property", "properties", "listing", "listings", "unit", "units",
514
+ "compare", "comparison", "availability", "available", "booked", "specific"}
515
+ _GEO_KEYWORDS = {"near", "nearby", "close to", "distance", "address", "location",
516
+ "neighborhood", "area", "km", "mile", "minutes", "downtown"}
517
+
518
+
519
+ def select_tools(user_query: str) -> list:
520
+ """
521
+ Pick the smallest tool subset that covers the user's intent.
522
+
523
+ Falls back to the full REAL_ESTATE_TOOLS list only when the query is
524
+ genuinely ambiguous across multiple intent categories (e.g. Geo + Property).
525
+
526
+ Typical savings: 2,000-2,500 tokens vs sending all 18 tools.
527
+ """
528
+ q = user_query.lower()
529
+
530
+ geo_hit = any(kw in q for kw in _GEO_KEYWORDS)
531
+ anomaly_hit = any(kw in q for kw in _ANOMALY_KEYWORDS)
532
+ property_hit = any(kw in q for kw in _PROPERTY_KEYWORDS)
533
+
534
+ # Geo queries that also mention properties (e.g. "find properties near Brickell")
535
+ if geo_hit and property_hit:
536
+ return REAL_ESTATE_TOOLS
537
+
538
+ # Anomaly queries take precedence when anomaly terms like "spike", "volatile" are present
539
+ if anomaly_hit:
540
+ return ANOMALY_TOOLS
541
+
542
+ if geo_hit:
543
+ return GEO_TOOLS
544
+
545
+ if property_hit:
546
+ return PROPERTY_TOOLS
547
+
548
+ # Default: market overview (covers broad questions like "how is the market?")
549
+ return MARKET_TOOLS