Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |
| 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 | |