Spaces:
Running
Running
File size: 10,514 Bytes
79d4fd5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 | import logging
import re
from functools import lru_cache
from typing import Any, Dict, List, Optional, Tuple
NL2SQL_MAX_CONTEXT_TOKENS = 4000
CHAT_HISTORY_MAX_TOKENS = 2000
SCHEMA_CONTEXT_MAX_TOKENS = 2000
FEW_SHOT_MAX_TOKENS = 1500
logger = logging.getLogger(__name__)
class TokenOptimizer:
def __init__(self, max_tokens: int = 4000):
self._max_tokens = max_tokens
self._encoder = None
@property
def encoder(self):
if self._encoder is None:
try:
import tiktoken
self._encoder = tiktoken.get_encoding("cl100k_base")
except ImportError:
logger.warning("tiktoken not available, using approximate token counting")
self._encoder = _ApproximateEncoder()
return self._encoder
def count_tokens(self, text: str) -> int:
if not text:
return 0
return self._count_tokens_cached(text)
@lru_cache(maxsize=1000)
def _count_tokens_cached(self, text: str) -> int:
return len(self.encoder.encode(text))
def truncate_chat_history(
self,
messages: List[Dict[str, str]],
max_tokens: Optional[int] = None,
) -> Tuple[List[Dict[str, str]], int]:
if not messages:
return [], 0
budget = max_tokens or CHAT_HISTORY_MAX_TOKENS
token_counts = [(msg, self.count_tokens(msg.get("content", ""))) for msg in messages]
kept = []
total_tokens = 0
for msg, tokens in reversed(token_counts):
if total_tokens + tokens > budget and kept:
break
kept.insert(0, msg)
total_tokens += tokens
kept = self._preserve_message_pairs(kept, messages)
final_tokens = sum(self.count_tokens(m.get("content", "")) for m in kept)
return kept, final_tokens
def _preserve_message_pairs(
self,
kept: List[Dict[str, str]],
original: List[Dict[str, str]],
) -> List[Dict[str, str]]:
if not kept or len(kept) < 2:
return kept
if kept[0].get("role") == "assistant" and len(original) > 1:
idx = original.index(kept[0]) if kept[0] in original else -1
if idx > 0 and original[idx - 1].get("role") == "user":
kept.insert(0, original[idx - 1])
return kept
def truncate_context(
self,
context: str,
max_tokens: Optional[int] = None,
preserve_start: bool = True,
) -> Tuple[str, int]:
if not context:
return "", 0
budget = max_tokens or self._max_tokens
current_tokens = self.count_tokens(context)
if current_tokens <= budget:
return context, current_tokens
sentences = re.split(r'(?<=[.!?])\s+', context)
if len(sentences) <= 1:
return self._hard_truncate(context, budget, preserve_start)
if preserve_start:
return self._truncate_from_end(sentences, budget)
return self._truncate_from_start(sentences, budget)
def _truncate_from_end(self, sentences: List[str], budget: int) -> Tuple[str, int]:
result = []
total = 0
for sentence in sentences:
tokens = self.count_tokens(sentence)
if total + tokens > budget and result:
break
result.append(sentence)
total += tokens
text = " ".join(result) + "..."
return text, self.count_tokens(text)
def _truncate_from_start(self, sentences: List[str], budget: int) -> Tuple[str, int]:
result = []
total = 0
for sentence in reversed(sentences):
tokens = self.count_tokens(sentence)
if total + tokens > budget and result:
break
result.insert(0, sentence)
total += tokens
text = "..." + " ".join(result)
return text, self.count_tokens(text)
def _hard_truncate(self, text: str, budget: int, preserve_start: bool) -> Tuple[str, int]:
tokens = self.encoder.encode(text)
truncated_tokens = tokens[:budget - 1] if preserve_start else tokens[-(budget - 1):]
truncated = self.encoder.decode(truncated_tokens)
suffix = "..." if preserve_start else ""
prefix = "..." if not preserve_start else ""
result = prefix + truncated + suffix
return result, self.count_tokens(result)
def optimize_schema_context(
self,
schema: str,
relevant_tables: Optional[List[str]] = None,
max_tokens: Optional[int] = None,
) -> str:
if not schema:
return ""
budget = max_tokens or SCHEMA_CONTEXT_MAX_TOKENS
current = self.count_tokens(schema)
if current <= budget:
return schema
table_blocks = self._parse_table_blocks(schema)
if not table_blocks:
truncated, _ = self.truncate_context(schema, budget)
return truncated
prioritized = []
remaining = []
relevant_lower = [t.lower() for t in (relevant_tables or [])]
for name, block in table_blocks:
if relevant_lower and name.lower() in relevant_lower:
prioritized.append(block)
else:
remaining.append(block)
result_parts = []
total = 0
for block in prioritized + remaining:
tokens = self.count_tokens(block)
if total + tokens > budget and result_parts:
break
result_parts.append(block)
total += tokens
return "\n\n".join(result_parts)
def _parse_table_blocks(self, schema: str) -> List[Tuple[str, str]]:
blocks = []
current_name = ""
current_lines: List[str] = []
for line in schema.split("\n"):
table_match = re.match(r'^Table:\s*(\w+)', line, re.IGNORECASE)
if table_match:
if current_name and current_lines:
blocks.append((current_name, "\n".join(current_lines)))
current_name = table_match.group(1)
current_lines = [line]
elif current_name:
current_lines.append(line)
if current_name and current_lines:
blocks.append((current_name, "\n".join(current_lines)))
return blocks
def select_few_shot_examples(
self,
examples: List[Dict[str, Any]],
query: str,
max_tokens: Optional[int] = None,
) -> List[Dict[str, Any]]:
if not examples:
return []
budget = max_tokens or FEW_SHOT_MAX_TOKENS
scored = [(ex, self._keyword_similarity(query, self._example_text(ex))) for ex in examples]
scored.sort(key=lambda x: x[1], reverse=True)
selected = []
total = 0
for ex, score in scored:
ex_text = self._example_text(ex)
tokens = self.count_tokens(ex_text)
if total + tokens > budget and selected:
break
selected.append(ex)
total += tokens
return selected
def _example_text(self, example: Dict[str, Any]) -> str:
parts = []
for key in ("query", "question", "input", "sql", "output", "answer", "context"):
if key in example:
parts.append(str(example[key]))
return " ".join(parts) if parts else str(example)
def _keyword_similarity(self, query: str, text: str) -> float:
query_words = set(query.lower().split())
text_words = set(text.lower().split())
if not query_words:
return 0.0
intersection = query_words & text_words
return len(intersection) / len(query_words)
def optimize_for_llm_call(
self,
components: Dict[str, str],
total_budget: Optional[int] = None,
) -> Dict[str, str]:
budget = total_budget or NL2SQL_MAX_CONTEXT_TOKENS
result = {}
query = components.get("query", "")
system_prompt = components.get("system_prompt", "")
result["query"] = query
result["system_prompt"] = system_prompt
reserved = self.count_tokens(query) + self.count_tokens(system_prompt)
remaining = max(0, budget - reserved)
schema_budget = int(remaining * 0.40)
few_shots_budget = int(remaining * 0.30)
history_budget = remaining - schema_budget - few_shots_budget
schema = components.get("schema", "")
if schema:
truncated, _ = self.truncate_context(schema, schema_budget)
result["schema"] = truncated
else:
result["schema"] = ""
few_shots = components.get("few_shots", "")
if few_shots:
truncated, _ = self.truncate_context(few_shots, few_shots_budget)
result["few_shots"] = truncated
else:
result["few_shots"] = ""
chat_history = components.get("chat_history", "")
if chat_history:
truncated, _ = self.truncate_context(chat_history, history_budget, preserve_start=False)
result["chat_history"] = truncated
else:
result["chat_history"] = ""
return result
def get_optimization_stats(self, original: Dict[str, str], optimized: Dict[str, str]) -> Dict[str, Any]:
original_tokens = sum(self.count_tokens(v) for v in original.values())
optimized_tokens = sum(self.count_tokens(v) for v in optimized.values())
reduction = original_tokens - optimized_tokens
ratio = reduction / original_tokens if original_tokens > 0 else 0.0
return {
"original_tokens": original_tokens,
"optimized_tokens": optimized_tokens,
"tokens_saved": reduction,
"reduction_ratio": round(ratio, 4),
}
class _ApproximateEncoder:
def encode(self, text: str) -> List[int]:
return list(range(len(text) // 4))
def decode(self, tokens: List[int]) -> str:
return " " * (len(tokens) * 4)
_token_optimizer_instance: Optional[TokenOptimizer] = None
def get_token_optimizer(max_tokens: int = 4000) -> TokenOptimizer:
global _token_optimizer_instance
if _token_optimizer_instance is None:
_token_optimizer_instance = TokenOptimizer(max_tokens=max_tokens)
return _token_optimizer_instance
def reset_token_optimizer() -> None:
global _token_optimizer_instance
_token_optimizer_instance = None
|