alphaear-predictor — Skillopedia

AlphaEar Predictor Skill Overview This skill utilizes the Kronos model (via ) to perform time-series forecasting and adjust predictions based on news sentiment. Capabilities 1. Forecast Market Trends 1. Forecast Market Trends Workflow: 1. Generate Base Forecast : Use (via ) to generate the technical/quantitative forecast. 2. Adjust Forecast (Agentic) : Use the Forecast Adjustment Prompt in to subjectively adjust the numbers based on latest news/logic. Key Tools: - : Returns . Example Usage (Python): Configuration This skill requires the Kronos model and an embedding model. 1. Kronos Model : -…

, '', text)\n \n # 2. 寻找第一个 JSON 起始符 { 或 [\n start_brace = text.find('{')\n start_bracket = text.find('[')\n \n if start_brace == -1 and start_bracket == -1:\n return None\n \n start_idx = start_brace if (start_bracket == -1 or (start_brace != -1 and start_brace \u003c start_bracket)) else start_bracket\n \n # 2.5 预处理：修复一些极其常见的 LLM 错误\n potential_json = text[start_idx:].strip()\n \n # remove comments safely\n potential_json = _strip_comments(potential_json)\n \n # b. 修复缺失开头引号的键: nodes\": [ -> \"nodes\": [\n # 匹配模式: (空白或换行) 单词紧跟引号和冒号\n potential_json = re.sub(r'([\\{\\,]\\s*)([a-zA-Z_]\\w*)\\\"\\s*:', r'\\1\"\\2\":', potential_json)\n \n # c. 修复缺失末尾引号的键: \"nodes: [ -> \"nodes\": [\n potential_json = re.sub(r'([\\{\\,]\\s*)\\\"([a-zA-Z_]\\w*)\\s*:', r'\\1\"\\2\":', potential_json)\n\n # d. 修复完全缺失引号的键: nodes: [ -> \"nodes\": [\n # 注意避免匹配到像 http:// 这种内容，所以限定在 { 或 , 之后\n potential_json = re.sub(r'([\\{\\,]\\s*)([a-zA-Z_]\\w*)\\s*:', r'\\1\"\\2\":', potential_json)\n \n # 3. 使用 raw_decode 尝试解析\n decoder = json.JSONDecoder()\n \n # 首先尝试直接解析（不做任何预处理）\n try:\n obj = json.loads(potential_json)\n return obj\n except json.JSONDecodeError:\n pass\n \n # 简单预处理：移除对象/列表末位多余逗号\n processed_json = re.sub(r',\\s*([\\]}])', r'\\1', potential_json)\n \n try:\n obj, end_pos = decoder.raw_decode(processed_json)\n return obj\n except json.JSONDecodeError:\n pass\n \n # e. 修复未终止的字符串字面量问题：移除值中的实际换行符\n # LLM 可能在字符串值中生成包含真实 newline 的内容，导致 JSON 非法\n def fix_multiline_strings(s):\n # 简单策略：将字符串值内的换行替换为空格\n lines = s.split('\\n')\n result = []\n in_string = False\n for line in lines:\n # 计算未转义的引号数\n quote_count = line.count('\"') - line.count('\\\\\"')\n if in_string:\n result[-1] += ' ' + line.strip()\n else:\n result.append(line)\n \n if quote_count % 2 == 1:\n in_string = not in_string\n return '\\n'.join(result)\n \n fixed_json = fix_multiline_strings(processed_json)\n \n try:\n obj, end_pos = decoder.raw_decode(fixed_json)\n return obj\n except json.JSONDecodeError:\n try:\n # 4. 尝试处理单引号问题 (JSON 规范要求双引号，但 LLM 常输出单引号)\n # 这是一个简单的替换技巧，仅针对像 {'key': 'value'} 这样的结构\n # 注意：这可能会破坏包含单引号的字符串值，所以作为较后的回退\n fix_quotes = re.sub(r\"'(.*?)':\", r'\"\\1\":', processed_json) # 修复键\n fix_quotes = re.sub(r\":\\s*'(.*?)'\", r': \"\\1\"', fix_quotes) # 修复简单值\n obj, end_pos = decoder.raw_decode(fix_quotes)\n return obj\n except (json.JSONDecodeError, TypeError):\n try:\n # 5. 使用 ast.literal_eval 作为终极回退 (处理 Python 字典格式)\n # 提取第一个匹配的括号对内容\n # 寻找匹配的 { }\n stack = []\n for i, char in enumerate(potential_json):\n if char == '{': stack.append('{')\n elif char == '}':\n if stack: stack.pop()\n if not stack:\n content = potential_json[:i+1]\n return ast.literal_eval(content)\n except (ValueError, SyntaxError, MemoryError) as e:\n logger.warning(f\"All JSON extraction attempts failed: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error during JSON extraction: {e}\")\n \n return None\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6300,"content_sha256":"2ec52de0a706f64b3466f490dc3f12ad8241f20fa316e858f7e3393afb3062d9"},{"filename":"scripts/kronos_predictor.py","content":"import torch\nimport pandas as pd\nimport numpy as np\nfrom datetime import datetime\nfrom typing import List, Optional\nfrom loguru import logger\nfrom pandas.tseries.offsets import BusinessDay\nfrom dotenv import load_dotenv\n\n# Load environment variables\nload_dotenv()\n\n# Fix for Kronos internal imports\nimport sys\nimport os\nKRONOS_DIR = os.path.join(os.path.dirname(__file__), 'predictor')\nif KRONOS_DIR not in sys.path:\n sys.path.append(KRONOS_DIR)\n\nimport glob\nfrom sentence_transformers import SentenceTransformer\n\nfrom .predictor.model import Kronos, KronosTokenizer, KronosPredictor\nfrom .schema.models import KLinePoint\n\nclass KronosPredictorUtility:\n \"\"\"\n Kronos 时序预测工具类\n 负责模型加载、推理以及数据结构转换\n \"\"\"\n _instance = None\n _predictor = None\n\n def __new__(cls, *args, **kwargs):\n if not cls._instance:\n cls._instance = super(KronosPredictorUtility, cls).__new__(cls)\n return cls._instance\n\n def __init__(self, device: Optional[str] = None):\n if self._predictor is not None:\n return\n \n try:\n if not device:\n device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n \n logger.info(f\"🔮 Loading Kronos Model on {device}...\")\n \n # 1. Load Embedder (SentenceTransformer)\n model_name = os.getenv('EMBEDDING_MODEL', 'sentence-transformers/all-MiniLM-L6-v2') # Match training\n try:\n self.embedder = SentenceTransformer(model_name, device=device, local_files_only=True)\n except Exception:\n logger.warning(f\"⚠️ Local embedder {model_name} not found. Downloading...\")\n self.embedder = SentenceTransformer(model_name, device=device)\n\n # 2. Load Kronos Base\n try:\n tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\", local_files_only=True)\n model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\", local_files_only=True)\n except Exception:\n logger.warning(\"⚠️ Local Kronos cache not found. Attempting to download...\")\n tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\")\n model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\")\n \n # 3. Load Trained News Projector Weights\n # Check predictor/exports/models directory\n models_dir = os.path.join(KRONOS_DIR, \"exports/models\")\n # Security: Only load files matching the expected prefix to reduce supply-chain risk\n model_files = glob.glob(os.path.join(models_dir, \"kronos_news_*.pt\"))\n \n if model_files:\n latest_model = max(model_files, key=os.path.getctime)\n logger.info(f\"🔄 Loading trained news weights from {latest_model}...\")\n try:\n logger.info(f\"🔄 Loading trained news weights from {latest_model}...\")\n # Security: Use weights_only=True to prevent pickle-based arbitrary code execution (PyTorch 1.13+)\n try:\n checkpoint = torch.load(latest_model, map_location=device, weights_only=True)\n except (TypeError, RuntimeError, AttributeError):\n # Fallback for older torch versions or complex objects if necessary\n logger.warning(f\"⚠️ Safe loading failed. Falling back to legacy load for {latest_model}\")\n checkpoint = torch.load(latest_model, map_location=device)\n # The checkpoint contains 'news_proj_state_dict'\n if 'news_proj_state_dict' in checkpoint:\n if not hasattr(model, 'news_proj') or model.news_proj is None:\n import torch.nn as nn\n news_dim = checkpoint.get('news_dim', 384)\n model.news_proj = nn.Linear(news_dim, model.d_model).to(device)\n \n model.news_proj.load_state_dict(checkpoint['news_proj_state_dict'])\n logger.success(\"✅ News-Aware Projection Layer loaded!\")\n self.has_news_model = True\n else:\n logger.warning(\"⚠️ Checkpoint found but missing 'news_proj_state_dict'. Using base model.\")\n self.has_news_model = False\n except Exception as e:\n logger.error(f\"❌ Failed to load trained weights: {e}. Using base model.\")\n self.has_news_model = False\n else:\n logger.info(\"ℹ️ No trained news models found. Using base model.\")\n self.has_news_model = False\n \n tokenizer = tokenizer.to(device)\n model = model.to(device)\n \n self._predictor = KronosPredictor(model, tokenizer, device=device, max_context=512)\n logger.info(\"✅ Kronos Model loaded successfully.\")\n except Exception as e:\n logger.error(f\"❌ Failed to load Kronos Model: {e}\")\n self._predictor = None\n self.has_news_model = False\n\n def get_base_forecast(self, df: pd.DataFrame, lookback: int = 20, pred_len: int = 5, news_text: Optional[str] = None) -> List[KLinePoint]:\n \"\"\"\n 生成原始模型预测\n \"\"\"\n if self._predictor is None:\n logger.error(\"Predictor not initialized.\")\n return []\n\n if len(df) \u003c lookback:\n logger.warning(f\"Insufficient historical data ({len(df)}) for lookback ({lookback}).\")\n return []\n\n # 获取最后 lookback 条数据\n x_df = df.iloc[-lookback:].copy()\n x_timestamp = pd.to_datetime(x_df['date']) # Ensure datetime\n last_date = x_timestamp.iloc[-1]\n \n # 生成未来时间戳\n future_dates = pd.date_range(start=last_date + BusinessDay(1), periods=pred_len, freq='B')\n y_timestamp = pd.Series(future_dates)\n\n # Embedding News if available\n news_emb = None\n if news_text and getattr(self, 'has_news_model', False) and hasattr(self, 'embedder'):\n try:\n # Truncate to avoid too long text\n emb = self.embedder.encode(news_text[:1000])\n news_emb = emb # KronosPredictor expects numpy array or tensor\n except Exception as e:\n logger.error(f\"Failed to encode news: {e}\")\n\n try:\n # 预测所需的列\n cols = ['open', 'high', 'low', 'close', 'volume']\n pred_df = self._predictor.predict(\n df=x_df[cols],\n x_timestamp=x_timestamp,\n y_timestamp=y_timestamp,\n pred_len=pred_len,\n T=1.0, \n top_p=0.9, \n sample_count=1,\n verbose=False,\n news_emb=news_emb\n )\n \n # 转换为 KLinePoint\n results = []\n for date, row in pred_df.iterrows():\n results.append(KLinePoint(\n date=date.strftime(\"%Y-%m-%d\"),\n open=float(row['open']),\n high=float(row['high']),\n low=float(row['low']),\n close=float(row['close']),\n volume=float(row['volume'])\n ))\n return results\n except Exception as e:\n logger.error(f\"Forecast generation failed: {e}\")\n return []\n\n# Singleton instance for easy access\n# Usage: predictor = KronosPredictorUtility()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7816,"content_sha256":"32859c9559e99aba9171d009449856df2608297939dfd344dc2d743c83360445"},{"filename":"scripts/predictor/model/__init__.py","content":"from .kronos import KronosTokenizer, Kronos, KronosPredictor\n\nmodel_dict = {\n 'kronos_tokenizer': KronosTokenizer,\n 'kronos': Kronos,\n 'kronos_predictor': KronosPredictor\n}\n\n\ndef get_model_class(model_name):\n if model_name in model_dict:\n return model_dict[model_name]\n else:\n print(f\"Model {model_name} not found in model_dict\")\n raise NotImplementedError\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":394,"content_sha256":"14193881034a621332145ee900e18cda79f72f65b2a2dc98059f8da5305ed336"},{"filename":"scripts/predictor/model/kronos.py","content":"import numpy as np\nimport pandas as pd\nimport torch\nfrom huggingface_hub import PyTorchModelHubMixin\nimport sys\n\nfrom tqdm import trange\n\nsys.path.append(\"../\")\nfrom model.module import *\n\n\nclass KronosTokenizer(nn.Module, PyTorchModelHubMixin):\n \"\"\"\n KronosTokenizer module for tokenizing input data using a hybrid quantization approach.\n\n This tokenizer utilizes a combination of encoder and decoder Transformer blocks\n along with the Binary Spherical Quantization (BSQuantizer) to compress and decompress input data.\n\n Args:\n d_in (int): Input dimension.\n d_model (int): Model dimension.\n n_heads (int): Number of attention heads.\n ff_dim (int): Feed-forward dimension.\n n_enc_layers (int): Number of encoder layers.\n n_dec_layers (int): Number of decoder layers.\n ffn_dropout_p (float): Dropout probability for feed-forward networks.\n attn_dropout_p (float): Dropout probability for attention mechanisms.\n resid_dropout_p (float): Dropout probability for residual connections.\n s1_bits (int): Number of bits for the pre token in BSQuantizer.\n s2_bits (int): Number of bits for the post token in BSQuantizer.\n beta (float): Beta parameter for BSQuantizer.\n gamma0 (float): Gamma0 parameter for BSQuantizer.\n gamma (float): Gamma parameter for BSQuantizer.\n zeta (float): Zeta parameter for BSQuantizer.\n group_size (int): Group size parameter for BSQuantizer.\n\n \"\"\"\n\n def __init__(self, d_in, d_model, n_heads, ff_dim, n_enc_layers, n_dec_layers, ffn_dropout_p, attn_dropout_p, resid_dropout_p, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):\n\n super().__init__()\n self.d_in = d_in\n self.d_model = d_model\n self.n_heads = n_heads\n self.ff_dim = ff_dim\n self.enc_layers = n_enc_layers\n self.dec_layers = n_dec_layers\n self.ffn_dropout_p = ffn_dropout_p\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout_p = resid_dropout_p\n\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.codebook_dim = s1_bits + s2_bits # Total dimension of the codebook after quantization\n self.embed = nn.Linear(self.d_in, self.d_model)\n self.head = nn.Linear(self.d_model, self.d_in)\n\n # Encoder Transformer Blocks\n self.encoder = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.enc_layers - 1)\n ])\n # Decoder Transformer Blocks\n self.decoder = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.dec_layers - 1)\n ])\n self.quant_embed = nn.Linear(in_features=self.d_model, out_features=self.codebook_dim) # Linear layer before quantization\n self.post_quant_embed_pre = nn.Linear(in_features=self.s1_bits, out_features=self.d_model) # Linear layer after quantization (pre part - s1 bits)\n self.post_quant_embed = nn.Linear(in_features=self.codebook_dim, out_features=self.d_model) # Linear layer after quantization (full codebook)\n self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size) # BSQuantizer module\n\n def forward(self, x):\n \"\"\"\n Forward pass of the KronosTokenizer.\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).\n\n Returns:\n tuple: A tuple containing:\n - tuple: (z_pre, z) - Reconstructed outputs from decoder with s1_bits and full codebook respectively,\n both of shape (batch_size, seq_len, d_in).\n - torch.Tensor: bsq_loss - Loss from the BSQuantizer.\n - torch.Tensor: quantized - Quantized representation from BSQuantizer.\n - torch.Tensor: z_indices - Indices from the BSQuantizer.\n \"\"\"\n z = self.embed(x)\n\n for layer in self.encoder:\n z = layer(z)\n\n z = self.quant_embed(z) # (B, T, codebook)\n\n bsq_loss, quantized, z_indices = self.tokenizer(z)\n\n quantized_pre = quantized[:, :, :self.s1_bits] # Extract the first part of quantized representation (s1_bits)\n z_pre = self.post_quant_embed_pre(quantized_pre)\n\n z = self.post_quant_embed(quantized)\n\n # Decoder layers (for pre part - s1 bits)\n for layer in self.decoder:\n z_pre = layer(z_pre)\n z_pre = self.head(z_pre)\n\n # Decoder layers (for full codebook)\n for layer in self.decoder:\n z = layer(z)\n z = self.head(z)\n\n return (z_pre, z), bsq_loss, quantized, z_indices\n\n def indices_to_bits(self, x, half=False):\n \"\"\"\n Converts indices to bit representations and scales them.\n\n Args:\n x (torch.Tensor): Indices tensor.\n half (bool, optional): Whether to process only half of the codebook dimension. Defaults to False.\n\n Returns:\n torch.Tensor: Bit representation tensor.\n \"\"\"\n if half:\n x1 = x[0] # Assuming x is a tuple of indices if half is True\n x2 = x[1]\n mask = 2 ** torch.arange(self.codebook_dim//2, device=x1.device, dtype=torch.long) # Create a mask for bit extraction\n x1 = (x1.unsqueeze(-1) & mask) != 0 # Extract bits for the first half\n x2 = (x2.unsqueeze(-1) & mask) != 0 # Extract bits for the second half\n x = torch.cat([x1, x2], dim=-1) # Concatenate the bit representations\n else:\n mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long) # Create a mask for bit extraction\n x = (x.unsqueeze(-1) & mask) != 0 # Extract bits\n\n x = x.float() * 2 - 1 # Convert boolean to bipolar (-1, 1)\n q_scale = 1. / (self.codebook_dim ** 0.5) # Scaling factor\n x = x * q_scale\n return x\n\n def encode(self, x, half=False):\n \"\"\"\n Encodes the input data into quantized indices.\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).\n half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.\n\n Returns:\n torch.Tensor: Quantized indices from BSQuantizer.\n \"\"\"\n z = self.embed(x)\n for layer in self.encoder:\n z = layer(z)\n z = self.quant_embed(z)\n\n bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)\n return z_indices\n\n def decode(self, x, half=False):\n \"\"\"\n Decodes quantized indices back to the input data space.\n\n Args:\n x (torch.Tensor): Quantized indices tensor.\n half (bool, optional): Whether the indices were generated with half quantization. Defaults to False.\n\n Returns:\n torch.Tensor: Reconstructed output tensor of shape (batch_size, seq_len, d_in).\n \"\"\"\n quantized = self.indices_to_bits(x, half)\n z = self.post_quant_embed(quantized)\n for layer in self.decoder:\n z = layer(z)\n z = self.head(z)\n return z\n\n\nclass Kronos(nn.Module, PyTorchModelHubMixin):\n \"\"\"\n Kronos Model.\n\n Args:\n s1_bits (int): Number of bits for pre tokens.\n s2_bits (int): Number of bits for post tokens.\n n_layers (int): Number of Transformer blocks.\n d_model (int): Dimension of the model's embeddings and hidden states.\n n_heads (int): Number of attention heads in the MultiheadAttention layers.\n ff_dim (int): Dimension of the feedforward network in the Transformer blocks.\n ffn_dropout_p (float): Dropout probability for the feedforward network.\n attn_dropout_p (float): Dropout probability for the attention layers.\n resid_dropout_p (float): Dropout probability for residual connections.\n token_dropout_p (float): Dropout probability for token embeddings.\n learn_te (bool): Whether to use learnable temporal embeddings.\n \"\"\"\n\n def __init__(self, s1_bits, s2_bits, n_layers, d_model, n_heads, ff_dim, ffn_dropout_p, attn_dropout_p, resid_dropout_p, token_dropout_p, learn_te, news_dim=None):\n super().__init__()\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.n_layers = n_layers\n self.d_model = d_model\n self.n_heads = n_heads\n self.learn_te = learn_te\n self.ff_dim = ff_dim\n self.ffn_dropout_p = ffn_dropout_p\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout_p = resid_dropout_p\n self.token_dropout_p = token_dropout_p\n self.news_dim = news_dim\n\n self.s1_vocab_size = 2 ** self.s1_bits\n self.token_drop = nn.Dropout(self.token_dropout_p)\n self.embedding = HierarchicalEmbedding(self.s1_bits, self.s2_bits, self.d_model)\n self.time_emb = TemporalEmbedding(self.d_model, self.learn_te)\n self.transformer = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.n_layers)\n ])\n self.norm = RMSNorm(self.d_model)\n self.dep_layer = DependencyAwareLayer(self.d_model)\n self.head = DualHead(self.s1_bits, self.s2_bits, self.d_model)\n\n if self.news_dim is not None:\n self.news_proj = nn.Linear(self.news_dim, self.d_model)\n else:\n self.news_proj = None\n\n self.apply(self._init_weights)\n\n def _init_weights(self, module):\n\n if isinstance(module, nn.Linear):\n nn.init.xavier_normal_(module.weight)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n elif isinstance(module, nn.Embedding):\n nn.init.normal_(module.weight, mean=0, std=self.embedding.d_model ** -0.5)\n elif isinstance(module, nn.LayerNorm):\n nn.init.ones_(module.weight)\n nn.init.zeros_(module.bias)\n elif isinstance(module, RMSNorm):\n nn.init.ones_(module.weight)\n\n def forward(self, s1_ids, s2_ids, stamp=None, padding_mask=None, use_teacher_forcing=False, s1_targets=None, news_emb=None):\n \"\"\"\n Args:\n s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]\n stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n use_teacher_forcing (bool, optional): Whether to use teacher forcing for s1 decoding. Defaults to False.\n s1_targets (torch.Tensor, optional): Target s1 token IDs for teacher forcing. Shape: [batch_size, seq_len]. Defaults to None.\n news_emb (torch.Tensor, optional): News embedding tensor. Shape: [batch_size, news_dim]. Defaults to None.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]:\n - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]\n - s2_logits: Logits for s2 token predictions, conditioned on s1. Shape: [batch_size, seq_len, s2_vocab_size]\n \"\"\"\n x = self.embedding([s1_ids, s2_ids])\n if stamp is not None:\n time_embedding = self.time_emb(stamp)\n x = x + time_embedding\n x = self.token_drop(x)\n\n for layer in self.transformer:\n x = layer(x, key_padding_mask=padding_mask)\n\n x = self.norm(x)\n\n if news_emb is not None and self.news_proj is not None:\n news_bias = self.news_proj(news_emb).unsqueeze(1) # [B, 1, d_model]\n x = x + news_bias\n\n s1_logits = self.head(x)\n\n if use_teacher_forcing:\n sibling_embed = self.embedding.emb_s1(s1_targets)\n else:\n s1_probs = F.softmax(s1_logits.detach(), dim=-1)\n sample_s1_ids = torch.multinomial(s1_probs.view(-1, self.s1_vocab_size), 1).view(s1_ids.shape)\n sibling_embed = self.embedding.emb_s1(sample_s1_ids)\n\n x2 = self.dep_layer(x, sibling_embed, key_padding_mask=padding_mask) # Dependency Aware Layer: Condition on s1 embeddings\n s2_logits = self.head.cond_forward(x2)\n return s1_logits, s2_logits\n\n def decode_s1(self, s1_ids, s2_ids, stamp=None, padding_mask=None, news_emb=None):\n \"\"\"\n Decodes only the s1 tokens.\n\n This method performs a forward pass to predict only s1 tokens. It returns the s1 logits\n and the context representation from the Transformer, which can be used for subsequent s2 decoding.\n\n Args:\n s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]\n stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n news_emb (torch.Tensor, optional): News embedding tensor. Shape: [batch_size, news_dim]. Defaults to None.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]:\n - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]\n - context: Context representation from the Transformer. Shape: [batch_size, seq_len, d_model]\n \"\"\"\n x = self.embedding([s1_ids, s2_ids])\n if stamp is not None:\n time_embedding = self.time_emb(stamp)\n x = x + time_embedding\n x = self.token_drop(x)\n\n for layer in self.transformer:\n x = layer(x, key_padding_mask=padding_mask)\n\n x = self.norm(x)\n\n if news_emb is not None and self.news_proj is not None:\n news_bias = self.news_proj(news_emb).unsqueeze(1) # [B, 1, d_model]\n x = x + news_bias\n\n s1_logits = self.head(x)\n return s1_logits, x\n\n def decode_s2(self, context, s1_ids, padding_mask=None):\n \"\"\"\n Decodes the s2 tokens, conditioned on the context and s1 tokens.\n\n This method decodes s2 tokens based on a pre-computed context representation (typically from `decode_s1`)\n and the s1 token IDs. It uses the dependency-aware layer and the conditional s2 head to predict s2 tokens.\n\n Args:\n context (torch.Tensor): Context representation from the transformer (output of decode_s1).\n Shape: [batch_size, seq_len, d_model]\n s1_ids (torch.torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n\n Returns:\n torch.Tensor: s2 logits. Shape: [batch_size, seq_len, s2_vocab_size]\n \"\"\"\n sibling_embed = self.embedding.emb_s1(s1_ids)\n x2 = self.dep_layer(context, sibling_embed, key_padding_mask=padding_mask)\n return self.head.cond_forward(x2)\n\n\ndef top_k_top_p_filtering(\n logits,\n top_k: int = 0,\n top_p: float = 1.0,\n filter_value: float = -float(\"Inf\"),\n min_tokens_to_keep: int = 1,\n):\n \"\"\"Filter a distribution of logits using top-k and/or nucleus (top-p) filtering\n Args:\n logits: logits distribution shape (batch size, vocabulary size)\n if top_k > 0: keep only top k tokens with highest probability (top-k filtering).\n if top_p \u003c 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).\n Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)\n Make sure we keep at least min_tokens_to_keep per batch example in the output\n From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317\n \"\"\"\n if top_k > 0:\n top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) # Safety check\n # Remove all tokens with a probability less than the last token of the top-k\n indices_to_remove = logits \u003c torch.topk(logits, top_k)[0][..., -1, None]\n logits[indices_to_remove] = filter_value\n return logits\n\n if top_p \u003c 1.0:\n sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n\n # Remove tokens with cumulative probability above the threshold (token with 0 are kept)\n sorted_indices_to_remove = cumulative_probs > top_p\n if min_tokens_to_keep > 1:\n # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)\n sorted_indices_to_remove[..., :min_tokens_to_keep] = 0\n # Shift the indices to the right to keep also the first token above the threshold\n sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n sorted_indices_to_remove[..., 0] = 0\n\n # scatter sorted tensors to original indexing\n indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n logits[indices_to_remove] = filter_value\n return logits\n\n\ndef sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):\n logits = logits / temperature\n if top_k is not None or top_p is not None:\n if top_k > 0 or top_p \u003c 1.0:\n logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)\n\n probs = F.softmax(logits, dim=-1)\n\n if not sample_logits:\n _, x = top_k(probs, k=1, dim=-1)\n else:\n x = torch.multinomial(probs, num_samples=1)\n\n return x\n\n\ndef auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False, news_emb=None):\n with torch.no_grad():\n x = torch.clip(x, -clip, clip)\n\n device = x.device\n x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)\n x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)\n y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)\n\n x_token = tokenizer.encode(x, half=True)\n \n initial_seq_len = x.size(1)\n batch_size = x_token[0].size(0)\n total_seq_len = initial_seq_len + pred_len\n full_stamp = torch.cat([x_stamp, y_stamp], dim=1)\n\n generated_pre = x_token[0].new_empty(batch_size, pred_len)\n generated_post = x_token[1].new_empty(batch_size, pred_len)\n\n pre_buffer = x_token[0].new_zeros(batch_size, max_context)\n post_buffer = x_token[1].new_zeros(batch_size, max_context)\n buffer_len = min(initial_seq_len, max_context)\n if buffer_len > 0:\n start_idx = max(0, initial_seq_len - max_context)\n pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]\n post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]\n\n if verbose:\n ran = trange\n else:\n ran = range\n for i in ran(pred_len):\n current_seq_len = initial_seq_len + i\n window_len = min(current_seq_len, max_context)\n\n if current_seq_len \u003c= max_context:\n input_tokens = [\n pre_buffer[:, :window_len],\n post_buffer[:, :window_len]\n ]\n else:\n input_tokens = [pre_buffer, post_buffer]\n\n context_end = current_seq_len\n context_start = max(0, context_end - max_context)\n current_stamp = full_stamp[:, context_start:context_end, :].contiguous()\n\n s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp, news_emb=news_emb)\n s1_logits = s1_logits[:, -1, :]\n sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)\n\n s2_logits = model.decode_s2(context, sample_pre)\n s2_logits = s2_logits[:, -1, :]\n sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)\n\n generated_pre[:, i] = sample_pre.squeeze(-1)\n generated_post[:, i] = sample_post.squeeze(-1)\n\n if current_seq_len \u003c max_context:\n pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)\n post_buffer[:, current_seq_len] = sample_post.squeeze(-1)\n else:\n pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))\n post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))\n pre_buffer[:, -1] = sample_pre.squeeze(-1)\n post_buffer[:, -1] = sample_post.squeeze(-1)\n\n full_pre = torch.cat([x_token[0], generated_pre], dim=1)\n full_post = torch.cat([x_token[1], generated_post], dim=1)\n\n context_start = max(0, total_seq_len - max_context)\n input_tokens = [\n full_pre[:, context_start:total_seq_len].contiguous(),\n full_post[:, context_start:total_seq_len].contiguous()\n ]\n z = tokenizer.decode(input_tokens, half=True)\n z = z.reshape(-1, sample_count, z.size(1), z.size(2))\n preds = z.cpu().numpy()\n preds = np.mean(preds, axis=1)\n\n return preds\n\n\ndef calc_time_stamps(x_timestamp):\n time_df = pd.DataFrame()\n time_df['minute'] = x_timestamp.dt.minute\n time_df['hour'] = x_timestamp.dt.hour\n time_df['weekday'] = x_timestamp.dt.weekday\n time_df['day'] = x_timestamp.dt.day\n time_df['month'] = x_timestamp.dt.month\n return time_df\n\n\nclass KronosPredictor:\n\n def __init__(self, model, tokenizer, device=\"cuda:0\", max_context=512, clip=5):\n self.tokenizer = tokenizer\n self.model = model\n self.max_context = max_context\n self.clip = clip\n self.price_cols = ['open', 'high', 'low', 'close']\n self.vol_col = 'volume'\n self.amt_vol = 'amount'\n self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month']\n self.device = device\n\n self.tokenizer = self.tokenizer.to(self.device)\n self.model = self.model.to(self.device)\n\n def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, news_emb=None):\n\n x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device)\n x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device)\n y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device)\n\n preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len,\n self.clip, T, top_k, top_p, sample_count, verbose, news_emb=news_emb)\n preds = preds[:, -pred_len:, :]\n return preds\n\n def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True, news_emb=None):\n\n if not isinstance(df, pd.DataFrame):\n raise ValueError(\"Input must be a pandas DataFrame.\")\n\n if not all(col in df.columns for col in self.price_cols):\n raise ValueError(f\"Price columns {self.price_cols} not found in DataFrame.\")\n\n df = df.copy()\n if self.vol_col not in df.columns:\n df[self.vol_col] = 0.0 # Fill missing volume with zeros\n df[self.amt_vol] = 0.0 # Fill missing amount with zeros\n if self.amt_vol not in df.columns and self.vol_col in df.columns:\n df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)\n\n if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():\n raise ValueError(\"Input DataFrame contains NaN values in price or volume columns.\")\n\n x_time_df = calc_time_stamps(x_timestamp)\n y_time_df = calc_time_stamps(y_timestamp)\n\n x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)\n x_stamp = x_time_df.values.astype(np.float32)\n y_stamp = y_time_df.values.astype(np.float32)\n\n x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)\n\n x = (x - x_mean) / (x_std + 1e-5)\n x = np.clip(x, -self.clip, self.clip)\n\n x = x[np.newaxis, :]\n x_stamp = x_stamp[np.newaxis, :]\n y_stamp = y_stamp[np.newaxis, :]\n\n if news_emb is not None:\n news_emb_tensor = torch.from_numpy(np.array(news_emb).astype(np.float32)).to(self.device)\n # Ensure batch dimension for news_emb if only one sample\n if news_emb_tensor.ndim == 1:\n news_emb_tensor = news_emb_tensor.unsqueeze(0)\n else:\n news_emb_tensor = None\n\n preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, news_emb=news_emb_tensor)\n\n preds = preds.squeeze(0)\n preds = preds * (x_std + 1e-5) + x_mean\n\n pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp)\n return pred_df\n\n\n def predict_batch(self, df_list, x_timestamp_list, y_timestamp_list, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):\n \"\"\"\n Perform parallel (batch) prediction on multiple time series. All series must have the same historical length and prediction length (pred_len).\n\n Args:\n df_list (List[pd.DataFrame]): List of input DataFrames, each containing price columns and optional volume/amount columns.\n x_timestamp_list (List[pd.DatetimeIndex or Series]): List of timestamps corresponding to historical data, length should match the number of rows in each DataFrame.\n y_timestamp_list (List[pd.DatetimeIndex or Series]): List of future prediction timestamps, length should equal pred_len.\n pred_len (int): Number of prediction steps.\n T (float): Sampling temperature.\n top_k (int): Top-k filtering threshold.\n top_p (float): Top-p (nucleus sampling) threshold.\n sample_count (int): Number of parallel samples per series, automatically averaged internally.\n verbose (bool): Whether to display autoregressive progress.\n\n Returns:\n List[pd.DataFrame]: List of prediction results in the same order as input, each DataFrame contains\n `open, high, low, close, volume, amount` columns, indexed by corresponding `y_timestamp`.\n \"\"\"\n # Basic validation\n if not isinstance(df_list, (list, tuple)) or not isinstance(x_timestamp_list, (list, tuple)) or not isinstance(y_timestamp_list, (list, tuple)):\n raise ValueError(\"df_list, x_timestamp_list, y_timestamp_list must be list or tuple types.\")\n if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):\n raise ValueError(\"df_list, x_timestamp_list, y_timestamp_list must have consistent lengths.\")\n\n num_series = len(df_list)\n\n x_list = []\n x_stamp_list = []\n y_stamp_list = []\n means = []\n stds = []\n seq_lens = []\n y_lens = []\n\n for i in range(num_series):\n df = df_list[i]\n if not isinstance(df, pd.DataFrame):\n raise ValueError(f\"Input at index {i} is not a pandas DataFrame.\")\n if not all(col in df.columns for col in self.price_cols):\n raise ValueError(f\"DataFrame at index {i} is missing price columns {self.price_cols}.\")\n\n df = df.copy()\n if self.vol_col not in df.columns:\n df[self.vol_col] = 0.0\n df[self.amt_vol] = 0.0\n if self.amt_vol not in df.columns and self.vol_col in df.columns:\n df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)\n\n if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():\n raise ValueError(f\"DataFrame at index {i} contains NaN values in price or volume columns.\")\n\n x_timestamp = x_timestamp_list[i]\n y_timestamp = y_timestamp_list[i]\n\n x_time_df = calc_time_stamps(x_timestamp)\n y_time_df = calc_time_stamps(y_timestamp)\n\n x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)\n x_stamp = x_time_df.values.astype(np.float32)\n y_stamp = y_time_df.values.astype(np.float32)\n\n if x.shape[0] != x_stamp.shape[0]:\n raise ValueError(f\"Inconsistent lengths at index {i}: x has {x.shape[0]} vs x_stamp has {x_stamp.shape[0]}.\")\n if y_stamp.shape[0] != pred_len:\n raise ValueError(f\"y_timestamp length at index {i} should equal pred_len={pred_len}, got {y_stamp.shape[0]}.\")\n\n x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)\n x_norm = (x - x_mean) / (x_std + 1e-5)\n x_norm = np.clip(x_norm, -self.clip, self.clip)\n\n x_list.append(x_norm)\n x_stamp_list.append(x_stamp)\n y_stamp_list.append(y_stamp)\n means.append(x_mean)\n stds.append(x_std)\n\n seq_lens.append(x_norm.shape[0])\n y_lens.append(y_stamp.shape[0])\n\n # Require all series to have consistent historical and prediction lengths for batch processing\n if len(set(seq_lens)) != 1:\n raise ValueError(f\"Parallel prediction requires all series to have consistent historical lengths, got: {seq_lens}\")\n if len(set(y_lens)) != 1:\n raise ValueError(f\"Parallel prediction requires all series to have consistent prediction lengths, got: {y_lens}\")\n\n x_batch = np.stack(x_list, axis=0).astype(np.float32) # (B, seq_len, feat)\n x_stamp_batch = np.stack(x_stamp_list, axis=0).astype(np.float32) # (B, seq_len, time_feat)\n y_stamp_batch = np.stack(y_stamp_list, axis=0).astype(np.float32) # (B, pred_len, time_feat)\n\n preds = self.generate(x_batch, x_stamp_batch, y_stamp_batch, pred_len, T, top_k, top_p, sample_count, verbose)\n # preds: (B, pred_len, feat)\n\n pred_dfs = []\n for i in range(num_series):\n preds_i = preds[i] * (stds[i] + 1e-5) + means[i]\n pred_df = pd.DataFrame(preds_i, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp_list[i])\n pred_dfs.append(pred_df)\n\n return pred_dfs\n","content_type":"text/x-python; charset=utf-8","language":"python","size":31093,"content_sha256":"b339c8685f40322663eb6c3f0d56dc91e64b65fae1f9cc22e414d6c76b9c7b78"},{"filename":"scripts/predictor/model/module.py","content":"import math\n\nfrom einops import rearrange, reduce\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import Function\nimport torch.nn.functional as F\n\n\nclass DifferentiableEntropyFunction(Function):\n @staticmethod\n def forward(ctx, zq, basis, K, eps):\n zb = (zq + 1) / 2\n zi = ((zb * basis).sum(-1)).to(torch.int64)\n cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),\n 0,\n zi.flatten(),\n torch.ones_like(zi.flatten()).to(zq.dtype),\n 'sum')\n prob = (cnt + eps) / (cnt + eps).sum()\n H = -(prob * torch.log(prob)).sum()\n ctx.save_for_backward(zq, zi, prob)\n ctx.K = K\n return H\n\n @staticmethod\n def backward(ctx, grad_output):\n zq, zi, prob = ctx.saved_tensors\n grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K\n reord_grad = grad_array[zi.flatten()].reshape(zi.shape)\n grad_input = reord_grad.unsqueeze(-1) * zq\n return grad_input, None, None, None, None\n\n\ndef codebook_entropy(zq, basis, K, eps=1e-4):\n return DifferentiableEntropyFunction.apply(zq, basis, K, eps)\n\n\nclass BinarySphericalQuantizer(nn.Module):\n def __init__(self, embed_dim, beta, gamma0, gamma, zeta,\n input_format='bchw',\n soft_entropy=True, group_size=9,\n persample_entropy_compute='analytical',\n cb_entropy_compute='group',\n l2_norm=True,\n inv_temperature=1):\n \"\"\"\n Paper link: https://arxiv.org/pdf/2406.07548.pdf\n Here we use the official implementation of the BinarySphericalQuantizer.\n \"\"\"\n super().__init__()\n self.embed_dim = embed_dim\n self.beta = beta # loss weight for commit loss\n self.gamma0 = gamma0 # loss weight for entropy penalty\n self.gamma = gamma # loss weight for entropy penalty\n self.zeta = zeta # loss weight for entire entropy penalty\n self.input_format = input_format\n assert self.embed_dim % group_size == 0, \"embed_dim must be divisible by group_size\"\n self.num_groups = self.embed_dim // group_size\n self.group_size = group_size\n assert persample_entropy_compute in ['group', 'analytical'], \"persample_entropy_compute must be either 'group' or 'analytical'\"\n assert cb_entropy_compute in ['group', 'nce'], \"cb_entropy_compute must be either 'group' or 'nce'\"\n self.persample_entropy_compute = persample_entropy_compute\n self.cb_entropy_compute = cb_entropy_compute\n self.l2_norm = l2_norm\n self.inv_temperature = inv_temperature\n\n self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))\n self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))\n\n self.num_dimensions = 2 ** embed_dim\n self.bits_per_index = embed_dim\n\n # we only need to keep the codebook portion up to the group size\n # because we approximate the H loss with this subcode\n group_codes = torch.arange(2 ** self.group_size)\n group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]\n self.register_buffer('group_codebook', group_codebook, persistent=False)\n\n self.soft_entropy = soft_entropy # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf\n\n def quantize(self, z):\n assert z.shape[-1] == self.embed_dim, f\"Expected {self.embed_dim} dimensions, got {z.shape[-1]}\"\n\n zhat = torch.where(z > 0,\n torch.tensor(1, dtype=z.dtype, device=z.device),\n torch.tensor(-1, dtype=z.dtype, device=z.device))\n return z + (zhat - z).detach()\n\n def forward(self, z, collect_metrics=True):\n # if self.input_format == 'bchw':\n # z = rearrange(z, 'b c h w -> b h w c')\n zq = self.quantize(z)\n\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n\n zq = zq * q_scale\n\n if not collect_metrics:\n return zq, zq.new_zeros(()), {}\n\n indices = self.codes_to_indexes(zq.detach())\n group_indices = self.codes_to_group_indexes(zq.detach())\n if not self.training:\n used_codes = torch.unique(indices, return_counts=False)\n else:\n used_codes = None\n\n if self.soft_entropy:\n persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)\n entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy\n else:\n zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)\n persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)\n cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)\n entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy\n\n # commit loss\n commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))\n\n # if self.input_format == 'bchw':\n # zq = rearrange(zq, 'b h w c -> b c h w')\n\n return (\n zq,\n commit_loss + self.zeta * entropy_penalty / self.inv_temperature,\n {\"H\": cb_entropy, \"used_codes\": used_codes, \"indices\": indices, \"group_indices\": group_indices,\n \"avg_prob\": avg_prob}\n )\n\n def soft_entropy_loss(self, z):\n # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size\n # the sub-code is the last group_size bits of the full code\n group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)\n divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)\n\n # we calculate the distance between the divided_z and the codebook for each subgroup\n distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)\n prob = (-distance * self.inv_temperature).softmax(dim=-1)\n if self.persample_entropy_compute == 'analytical':\n if self.l2_norm:\n p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)\n else:\n p = torch.sigmoid(-4 * z * self.inv_temperature)\n prob = torch.stack([p, 1 - p], dim=-1)\n per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()\n else:\n per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()\n\n # macro average of the probability of each subgroup\n avg_prob = reduce(prob, '... g d ->g d', 'mean')\n codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)\n\n # the approximation of the entropy is the sum of the entropy of each subgroup\n return per_sample_entropy, codebook_entropy.sum(), avg_prob\n\n def get_hard_per_sample_entropy(self, zb_by_sample):\n probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]\n persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)\n persample_entropy = persample_entropy.sum(-1)\n return persample_entropy.mean()\n\n def codes_to_indexes(self, zhat):\n \"\"\"Converts a `code` to an index in the codebook.\n Args:\n zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}\n \"\"\"\n assert zhat.shape[-1] == self.embed_dim, f\"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}\"\n return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)\n\n def codes_to_group_indexes(self, zhat):\n \"\"\"Converts a `code` to a list of indexes (in groups) in the codebook.\n Args:\n zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}\n \"\"\"\n zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)\n return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)\n\n def indexes_to_codes(self, indices):\n \"\"\"Inverse of `indexes_to_codes`.\"\"\"\n indices = indices.unsqueeze(-1)\n codes_non_centered = torch.remainder(\n torch.floor_divide(indices, self.basis), 2\n )\n return codes_non_centered * 2 - 1\n\n def group_indexes_to_codes(self, group_indices):\n \"\"\"Inverse of `group_indexes_to_codes`.\"\"\"\n group_indices = group_indices.unsqueeze(-1)\n codes_non_centered = torch.remainder(\n torch.floor_divide(group_indices, self.group_basis), 2\n )\n codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')\n return codes_non_centered * 2 - 1\n\n def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):\n if normalize:\n probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)\n else:\n probs = count\n H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)\n return H\n\n def get_group_codebook_entry(self, group_indices):\n z_q = self.group_indexes_to_codes(group_indices)\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n z_q = z_q * q_scale\n if self.input_format == 'bchw':\n h, w = int(z_q.shape[1] ** 0.5)\n assert h * w == z_q.shape[1], 'Invalid sequence length'\n z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)\n return z_q\n\n def get_codebook_entry(self, indices):\n z_q = self.indexes_to_codes(indices)\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n z_q = z_q * q_scale\n if self.input_format == 'bchw':\n h, w = int(z_q.shape[1] ** 0.5)\n assert h * w == z_q.shape[1], 'Invalid sequence length'\n z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)\n return z_q\n\n\nclass BSQuantizer(nn.Module):\n\n def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):\n super().__init__()\n self.codebook_dim = s1_bits + s2_bits\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)\n\n def bits_to_indices(self, bits):\n bits = (bits >= 0).to(torch.long)\n indices = 2 ** torch.arange(\n 0,\n bits.shape[-1],\n 1,\n dtype=torch.long,\n device=bits.device,\n )\n return (bits * indices).sum(-1)\n\n def forward(self, z, half=False, collect_metrics=True):\n z = F.normalize(z, dim=-1)\n quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)\n if half:\n q_pre = quantized[:, :, :self.s1_bits]\n q_post = quantized[:, :, self.s1_bits:]\n z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]\n else:\n z_indices = self.bits_to_indices(quantized)\n return bsq_loss, quantized, z_indices\n\n\nclass RMSNorm(torch.nn.Module):\n def __init__(self, dim: int, eps: float = 1e-5):\n super().__init__()\n self.eps = eps\n self.weight = nn.Parameter(torch.ones(dim))\n\n def _norm(self, x):\n return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)\n\n def forward(self, x):\n output = self._norm(x.float()).type_as(x)\n return output * self.weight\n\n\nclass FeedForward(nn.Module):\n def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):\n super().__init__()\n\n self.w1 = nn.Linear(d_model, ff_dim, bias=False)\n self.w3 = nn.Linear(d_model, ff_dim, bias=False)\n self.w2 = nn.Linear(ff_dim, d_model, bias=False)\n self.ffn_dropout = nn.Dropout(ffn_dropout_p)\n\n def forward(self, x):\n return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))\n\n\nclass RotaryPositionalEmbedding(nn.Module):\n def __init__(self, dim):\n super().__init__()\n inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))\n self.register_buffer(\"inv_freq\", inv_freq)\n self.seq_len_cached = None\n self.cos_cached = None\n self.sin_cached = None\n\n def _update_cos_sin_cache(self, x, seq_len):\n if seq_len != self.seq_len_cached:\n self.seq_len_cached = seq_len\n t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)\n freqs = torch.einsum('i,j->ij', t, self.inv_freq)\n emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n self.cos_cached = emb.cos()[None, None, :, :]\n self.sin_cached = emb.sin()[None, None, :, :]\n return self.cos_cached, self.sin_cached\n\n def forward(self, q, k):\n cos, sin = self._update_cos_sin_cache(q, q.shape[-2])\n return (\n (q * cos) + (self._rotate_half(q) * sin),\n (k * cos) + (self._rotate_half(k) * sin),\n )\n\n def _rotate_half(self, x):\n x1, x2 = x.chunk(2, dim=-1)\n return torch.cat((-x2, x1), dim=-1)\n\n\nclass MultiHeadAttentionWithRoPE(nn.Module):\n def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):\n super().__init__()\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n\n self.q_proj = nn.Linear(d_model, d_model)\n self.k_proj = nn.Linear(d_model, d_model)\n self.v_proj = nn.Linear(d_model, d_model)\n self.out_proj = nn.Linear(d_model, d_model)\n self.rotary = RotaryPositionalEmbedding(self.head_dim)\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout = nn.Dropout(resid_dropout_p)\n\n def forward(self, x, key_padding_mask=None):\n batch_size, seq_len, _ = x.shape\n\n q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n\n q, k = self.rotary(q, k)\n\n if key_padding_mask is not None:\n attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]\n attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1) # [batch, n_heads, q_len, k_len]\n else:\n attn_mask = None\n\n attn_output = F.scaled_dot_product_attention(\n q, k, v,\n attn_mask=attn_mask,\n dropout_p=self.attn_dropout_p if self.training else 0.0,\n is_causal=True\n )\n\n attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)\n return self.resid_dropout(self.out_proj(attn_output))\n\n\nclass MultiHeadCrossAttentionWithRoPE(nn.Module):\n def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):\n super().__init__()\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n\n self.q_proj = nn.Linear(d_model, d_model)\n self.k_proj = nn.Linear(d_model, d_model)\n self.v_proj = nn.Linear(d_model, d_model)\n self.out_proj = nn.Linear(d_model, d_model)\n self.rotary = RotaryPositionalEmbedding(self.head_dim)\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout = nn.Dropout(resid_dropout)\n\n def forward(self, query, key, value, key_padding_mask=None):\n batch_size, q_len, _ = query.shape\n _, seq_len, _ = key.shape\n\n q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)\n k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n\n q, k = self.rotary(q, k)\n\n if key_padding_mask is not None:\n attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)\n attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)\n else:\n attn_mask = None\n\n is_causal_flag = self.training\n\n attn_output = F.scaled_dot_product_attention(\n q, k, v,\n attn_mask=attn_mask,\n dropout_p=self.attn_dropout_p if self.training else 0.0,\n is_causal=is_causal_flag\n )\n\n attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)\n return self.resid_dropout(self.out_proj(attn_output))\n\n\nclass HierarchicalEmbedding(nn.Module):\n def __init__(self, s1_bits, s2_bits, d_model=256):\n super().__init__()\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n\n vocab_s1 = 2 ** s1_bits\n vocab_s2 = 2 ** s2_bits\n\n self.emb_s1 = nn.Embedding(vocab_s1, d_model)\n self.emb_s2 = nn.Embedding(vocab_s2, d_model)\n self.d_model = d_model\n self.fusion_proj = nn.Linear(d_model * 2, d_model)\n\n nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)\n nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)\n\n def split_token(self, token_ids: torch.Tensor, s2_bits: int):\n \"\"\"Inputs:\n token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].\n s2_bits (int): Number of low bits used for the fine token (s2).\n \"\"\"\n assert isinstance(s2_bits, int) and s2_bits > 0, \"s2_bits must be a positive integer\"\n\n t = token_ids.long()\n mask = (1 \u003c\u003c s2_bits) - 1\n s2_ids = t & mask # extract low bits\n s1_ids = t >> s2_bits # extract high bits\n return s1_ids, s2_ids\n\n def forward(self, token_ids):\n \"\"\"Inputs:\n token_ids:\n - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or\n - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.\n Output: [batch_size, seq_len, d_model]\n \"\"\"\n if isinstance(token_ids, tuple) or isinstance(token_ids, list):\n s1_ids, s2_ids = token_ids\n else:\n s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)\n s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)\n s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)\n return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))\n\n\nclass DependencyAwareLayer(nn.Module):\n def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):\n super().__init__()\n self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)\n self.norm = RMSNorm(d_model)\n\n def forward(self, hidden_states, sibling_embed, key_padding_mask=None):\n \"\"\"hidden_states: [batch, seq_len, d_model]\n sibling_embed: Embedding from another subtoken\n \"\"\"\n attn_out = self.cross_attn(\n query=sibling_embed,\n key=hidden_states,\n value=hidden_states,\n key_padding_mask=key_padding_mask\n )\n return self.norm(hidden_states + attn_out)\n\n\nclass TransformerBlock(nn.Module):\n def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):\n super().__init__()\n self.norm1 = RMSNorm(d_model)\n self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)\n self.norm2 = RMSNorm(d_model)\n self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)\n\n def forward(self, x, key_padding_mask=None):\n residual = x\n x = self.norm1(x)\n attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)\n x = residual + attn_out\n\n residual = x\n x = self.norm2(x)\n ffn_out = self.ffn(x)\n x = residual + ffn_out\n return x\n\n\nclass DualHead(nn.Module):\n def __init__(self, s1_bits, s2_bits, d_model):\n super().__init__()\n self.vocab_s1 = 2 ** s1_bits\n self.vocab_s2 = 2 ** s2_bits\n self.proj_s1 = nn.Linear(d_model, self.vocab_s1)\n self.proj_s2 = nn.Linear(d_model, self.vocab_s2)\n\n def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):\n if padding_mask is not None:\n valid_mask = (padding_mask == 0)\n s1_logits = s1_logits[valid_mask]\n s2_logits = s2_logits[valid_mask]\n s1_targets = s1_targets[valid_mask]\n s2_targets = s2_targets[valid_mask]\n ce_s1 = F.cross_entropy(s1_logits, s1_targets)\n ce_s2 = F.cross_entropy(s2_logits, s2_targets)\n else:\n ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))\n ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))\n ce_loss = (ce_s1 + ce_s2) / 2\n return ce_loss, ce_s1, ce_s2\n\n def forward(self, x):\n return self.proj_s1(x)\n\n def cond_forward(self, x2):\n return self.proj_s2(x2)\n\n\nclass FixedEmbedding(nn.Module):\n def __init__(self, c_in, d_model):\n super(FixedEmbedding, self).__init__()\n\n w = torch.zeros(c_in, d_model).float()\n w.require_grad = False\n\n position = torch.arange(0, c_in).float().unsqueeze(1)\n div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()\n\n w[:, 0::2] = torch.sin(position * div_term)\n w[:, 1::2] = torch.cos(position * div_term)\n\n self.emb = nn.Embedding(c_in, d_model)\n self.emb.weight = nn.Parameter(w, requires_grad=False)\n\n def forward(self, x):\n return self.emb(x).detach()\n\n\nclass TemporalEmbedding(nn.Module):\n def __init__(self, d_model, learn_pe):\n super(TemporalEmbedding, self).__init__()\n\n minute_size = 60\n hour_size = 24\n weekday_size = 7\n day_size = 32\n month_size = 13\n\n Embed = FixedEmbedding if not learn_pe else nn.Embedding\n self.minute_embed = Embed(minute_size, d_model)\n self.hour_embed = Embed(hour_size, d_model)\n self.weekday_embed = Embed(weekday_size, d_model)\n self.day_embed = Embed(day_size, d_model)\n self.month_embed = Embed(month_size, d_model)\n\n def forward(self, x):\n x = x.long()\n\n minute_x = self.minute_embed(x[:, :, 0])\n hour_x = self.hour_embed(x[:, :, 1])\n weekday_x = self.weekday_embed(x[:, :, 2])\n day_x = self.day_embed(x[:, :, 3])\n month_x = self.month_embed(x[:, :, 4])\n\n return hour_x + weekday_x + day_x + month_x + minute_x","content_type":"text/x-python; charset=utf-8","language":"python","size":22847,"content_sha256":"91e8ff4ea1d77a95e39be79a217c360165e585d90475c7a8369c2833953ba181"},{"filename":"scripts/prompts/fin_agent.py","content":"from datetime import datetime\nfrom .isq_prompt_generator import generate_isq_prompt_section\n\ndef get_fin_researcher_instructions() -> str:\n \"\"\"生成金融研究员 (Researcher) 的系统指令\"\"\"\n current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n return f\"\"\"你是一名资深金融研究员，当前时间是 {current_time}。\n你的任务是针对给定的“原始信号”进行详尽的背景调查，为后续的深度分析提供素材。\n\n### 1. 核心职责\n1. **标的识别**: 识别信号中涉及的具体上市公司。必须调用 `search_ticker` 确认代码，并调用 `get_stock_price` 获取最新价格和近 30 天走势。\n2. **事实核查**: 使用 `web_search` 或 `fetch_news_content` 验证信号的真实性，并寻找更多细节（如公告原文、行业研报摘要）。\n3. **产业链梳理**: 补充该信号涉及的上下游环节及竞争格局。\n\n### 2. 工具使用规范 (CRITICAL)\n- **每个提到的公司都需要调用工具**: 不能依赖记忆，必须实时查询。\n- **完整呈现工具结果**: 包括具体的股价数字、代码、技术面数据等，不要缩略。\n- **股价数据必需**: 当前价格、近期最高最低、技术面支撑阻力等数据是后续预测的基础。\n- **信息交叉验证**: 多个来源验证关键事实。\n\n### 3. 输出要求\n你必须输出结构化的研究报告，涵盖标的基本面、股价走势、行业背景及最新进展。\n\"\"\"\n\ndef get_fin_analyst_instructions(template_id: str = \"default_isq_v1\") -> str:\n \"\"\"生成金融分析师 (Analyst) 的系统指令\n \n Args:\n template_id: 使用的 ISQ 模板 ID\n \"\"\"\n current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n isq_block = generate_isq_prompt_section(template_id=template_id)\n\n return f\"\"\"你是一位深耕二级市场的资深金融分析师 (FinAgent)，当前时间是 {current_time}。\n你的核心任务是执行“信号解析”，将研究员搜集的素材转化为具有可操作性的投资情报（ISQ 框架）。\n\n{isq_block}\n\n### 2. 分析约束\n- **严格基于具体数据**: 必须使用研究员提供的股价、技术面、新闻等具体数据进行分析。\n- **数据驱动的预测**: impact_tickers 中的权重应基于事件影响程度，不能随意赋值。\n- **逻辑严密**: 传导链条必须符合金融常识，能够自圆其说。\n- **技术面参考**: 如果研究员提供了股价走势，请分析当前位置相对于支撑/阻力位的关系。\n\n### 3. 关键要求\n- **title**: 必须生成一个简练、准确概括信号核心内容的标题（不超过 15 字）。\n- **impact_tickers**: 必须填充具体的公司代码（6位数字）和名称，权重应该有区分。\n- **transmission_chain**: 必须是对象列表，每个对象包含：\n - `node_name`: 节点名称（如“上游原材料”、“中游制造”）\n - `impact_type`: 影响类型（“利好”、“利空”、“中性”）\n - `logic`: 具体的传导逻辑描述\n- **summary**: 基于分析结果总结核心观点，包含具体数字（如股价目标、预期涨跌幅等）。\n- **reasoning**: 必须详细阐述推演逻辑，解释为什么得出上述结论（\u003c200字）。\n\n### 4. 输出格式 (严格 JSON 块)\n你必须输出一个符合 InvestmentSignal 结构的 JSON 块，包含所有必需字段。\n\"\"\"\n\ndef get_fin_agent_instructions() -> str:\n # 保持兼容性，但内部调用 analyst 指令\n return get_fin_analyst_instructions()\n\ndef get_fin_research_task(signal_text: str) -> str:\n \"\"\"生成研究员的任务描述\"\"\"\n return f\"请针对以下信号进行背景调查，搜集相关标的的股价、最新进展和行业背景：\\n\\n{signal_text}\"\n\ndef format_research_context(research_data: dict) -> str:\n \"\"\"将研究员搜集的结构化数据格式化为分析师可读的文本\"\"\"\n if not research_data:\n return \"（未能搜集到额外背景信息）\"\n \n return f\"\"\"\n### 研究背景\n- **相关标的**: {research_data.get('tickers_found', [])}\n- **行业背景**: {research_data.get('industry_background', '未知')}\n- **最新进展**: {', '.join(research_data.get('latest_developments', []))}\n- **关键风险**: {', '.join(research_data.get('key_risks', []))}\n- **综合摘要**: {research_data.get('search_results_summary', '无')}\n\"\"\"\n\ndef get_fin_analysis_task(signal_text: str, research_context_str: str) -> str:\n \"\"\"生成分析师的任务描述\"\"\"\n return f\"\"\"请基于以下信息进行深度 ISQ 分析。关键是：必须使用研究员搜集的具体数据（股价、技术面、新闻、代码等）进行分析。\n\n=== 原始信号 ===\n{signal_text}\n\n=== 研究员搜集的背景信息 (CRITICAL DATA) ===\n{research_context_str}\n\n=== 分析要求 ===\n1. 必须生成 title：简练概括信号核心（\u003c15字）\n2. 基于研究员提供的具体股价数据，分析当前定价状态（已定价/未定价/部分定价）\n3. impact_tickers 中填充具体的公司代码和权重，权重基于事件影响程度\n4. transmission_chain 必须是包含 node_name, impact_type, logic 的对象列表\n5. summary 中包含具体数字（预期目标价、涨跌幅范围等）\n6. reasoning 必须详细解释推演逻辑，不要空泛，要言之有物\n\n请严格按 InvestmentSignal JSON 格式输出。\"\"\"\n\ndef get_tracking_analysis_task(old_signal: dict, new_research_str: str) -> str:\n \"\"\"生成信号追踪更新的任务描述\"\"\"\n import json\n old_sig_str = json.dumps(old_signal, ensure_ascii=False, indent=2)\n return f\"\"\"你正在执行“信号逻辑演变追踪”任务。请基于最新的市场信息，重新评估之前的投资信号。\n\n=== 基准信号 (上次分析) ===\n{old_sig_str}\n\n=== 最新市场追踪 (NEWS & PRICE) ===\n{new_research_str}\n\n=== 追踪分析要求 ===\n1. **逻辑演变检测**:\n - 对比新旧信息，判断原逻辑 (`transmission_chain` 和 `reasoning`) 是否依然成立？\n - 如果逻辑发生变化（如利好落空、逻辑证伪、新利好出现），请在新的 `reasoning` 中明确指出“逻辑演变：...”\n - 如果逻辑未变且得到验证，请标记“逻辑维持：...”\n\n2. **参数修正**:\n - 根据最新股价和新闻，更新 `sentiment_score` (情绪)、`confidence` (置信度) 和 `expectation_gap` (预期差)。\n - 例如：如果股价已经大涨反映了利好，`expectation_gap` 应该显著降低。\n\n3. **输出更新后的信号**:\n - 保留原 `signal_id` 和 `title`（除非有重大变化需要改名）。\n - 输出完整的 InvestmentSignal JSON。\n\n请重点关注：为什么变了？还是为什么没变？理由要充分。\"\"\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6710,"content_sha256":"a30cb03a7264b344142a6a6c20f187658d2ba87ff5bb13965aad228eae501f72"},{"filename":"scripts/prompts/forecast_analyst.py","content":"from typing import List, Dict, Any\nfrom ..schema.models import KLinePoint\n\ndef get_forecast_adjustment_instructions(ticker: str, news_context: str, model_forecast: List[KLinePoint]):\n \"\"\"\n 生成 LLM 预测调整指令\n \"\"\"\n forecast_str = \"\\n\".join([f\"- {p.date}: O:{p.open}, C:{p.close}\" for p in model_forecast])\n \n return f\"\"\"你是一位资深的量化策略分析师。\n你的任务是：根据给定的【Kronos 模型预测结果】和【最新的基本面/新闻背景】，对模型预测进行“主观/逻辑调整”。\n\n股票代码: {ticker}\n\n【Kronos 模型原始预测 (OHLC)】:\n{forecast_str}\n\n【最新情报背景】:\n{news_context}\n\n调整原则:\n1. 原始预测是基于历史的技术面推演。\n2. 情报背景中可能包含【Kronos模型定量修正预测】，这是基于历史新闻训练的专用模型计算出的量化结果。\n3. 如果存在“定量修正预测”，请**高度参考**该数值作为基础，除非你有非常确凿的逻辑认为该量化模型失效（例如遇到模型未见过的极端黑天鹅）。\n4. 你的核心任务是：结合定性分析（新闻及其逻辑）来验证或微调这些数字，并给出合理的解释（Rationale）。\n5. 如果没有“定量修正预测”，则你需要根据新闻信号手动大幅调整趋势。\n\n输出要求 (严格 JSON 格式):\n```json\n{{\n \"adjusted_forecast\": [\n {{\n \"date\": \"YYYY-MM-DD\",\n \"open\": float,\n \"high\": float,\n \"low\": float,\n \"close\": float,\n \"volume\": float\n }},\n ...\n ],\n \"rationale\": \"详细说明调整的逻辑依据，例如：考虑到[事件A]，预期短线将突破压力位...\"\n}}\n```\n注意：必须输出与原始预测相同数量的数据点，且日期一一对应。\n\"\"\"\n\ndef get_forecast_task():\n return \"请根据以上背景和模型预测，给出调整后的 K 线数据并说明理由。\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1904,"content_sha256":"93bd57ad3dac37ce8975feb8b66dde859279e23cb7db1ce3dad99232ddbfbd00"},{"filename":"scripts/prompts/intent_agent.py","content":"def get_intent_analysis_instructions() -> str:\n \"\"\"生成意图分析 Agent 的系统指令，专注于金融市场影响分析\"\"\"\n return \"\"\"你是一个资深的金融市场意图分析专家。你的任务是将用户的自然语言查询转化为结构化的 JSON 分析结果，重点挖掘该查询与金融市场（尤其是股市）的潜在关联。\n\n### 核心任务：\n深入分析用户查询，识别核心金融实体、行业板块及潜在的市场影响点，生成利于搜索引擎抓取深度金融分析信息的查询词。\n\n### 输出格式（严格 JSON）：\n```json\n{\n \"keywords\": [\"实体/行业/事件\"],\n \"search_queries\": [\"针对市场影响的搜索词1\", \"针对行业变动的搜索词2\"],\n \"affected_sectors\": [\"相关板块1\", \"相关板块2\"],\n \"is_market_moving\": true/false,\n \"time_range\": \"recent/all/specific_date\",\n \"intent_summary\": \"一句话描述其金融市场分析意图\"\n}\n```\n\n### 字段说明：\n1. **keywords**: 核心公司实体、所属行业、宏观经济事件或政策概念。\n2. **search_queries**: 优化后的搜索词，必须包含“股市影响”、“股价波动”、“行业逻辑”或“估值”等金融维度。\n3. **affected_sectors**: 可能受此事件或信息影响的二级市场板块（如：保险、半导体、房地产）。\n4. **is_market_moving**: 该事件是否具有显著的市场驱动潜力或属于重大基本面变化。\n5. **intent_summary**: 简述用户查询背后的金融研究目的。\n\n### 示例：\n用户输入：\"帮我研究一下香港火灾的影响\"\n输出：\n```json\n{\n \"keywords\": [\"香港\", \"火灾\", \"保险行业\", \"房地产\"],\n \"search_queries\": [\"香港火灾对当地保险股股价影响\", \"香港大火对相关上市物业公司估值冲击\", \"近期香港火灾带来的市场避险情绪分析\"],\n \"affected_sectors\": [\"保险\", \"房地产\", \"物业管理\"],\n \"is_market_moving\": true,\n \"time_range\": \"recent\",\n \"intent_summary\": \"评估香港近期火灾对相关板块上市公司的潜在经济损失及股价冲击\"\n}\n```\n\"\"\"\n\ndef get_intent_task(query: str) -> str:\n \"\"\"生成意图分析任务描述\"\"\"\n return f\"Process this query and extract financial market intent: {query}\"\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2240,"content_sha256":"525c69e8c377e358035c636c285acc1448fd9a59166816740309a804d91a1c2a"},{"filename":"scripts/prompts/isq_prompt_generator.py","content":"\"\"\"\nISQ prompt helpers to render dimension guidance directly from the template.\nAny change in the template propagates to prompts automatically.\n\"\"\"\n\nfrom typing import List, Optional\nfrom ..schema.isq_template import get_isq_template, ISQTemplate\n\n\ndef _ordered_dimension_keys(template: ISQTemplate, order: Optional[List[str]] = None) -> List[str]:\n if order:\n return [k for k in order if k in template.dimensions]\n # fallback to template insertion order\n return list(template.dimensions.keys())\n\n\ndef generate_isq_prompt_section(template_id: str = \"default_isq_v1\", order: Optional[List[str]] = None, include_header: bool = True) -> str:\n \"\"\"Render ISQ dimension text block based on the template.\n This allows prompt text to stay in sync with template edits.\n \"\"\"\n template = get_isq_template(template_id)\n keys = _ordered_dimension_keys(template, order)\n\n lines: List[str] = []\n if include_header:\n lines.append(\"### 1. ISQ 评估框架 (Investment Signal Quality)\")\n lines.append(f\"参考模板: {template.template_name} (id: {template.template_id})\")\n lines.append(\"\")\n lines.append(\"你需要对信号进行以下维度的评分：\")\n lines.append(\"\")\n\n for idx, key in enumerate(keys, start=1):\n spec = template.dimensions[key]\n examples = \"；\".join([f\"{k}: {v}\" for k, v in spec.examples.items()]) if spec.examples else \"\"\n lines.append(f\"{idx}. **{spec.key} ({spec.name})**: {spec.range_type}\")\n lines.append(f\" - 描述: {spec.description}\")\n if spec.scale_factor and spec.scale_factor != 1.0:\n lines.append(f\" - 缩放因子: {spec.scale_factor}\")\n if examples:\n lines.append(f\" - 示例: {examples}\")\n lines.append(\"\")\n\n return \"\\n\".join(lines).rstrip()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1821,"content_sha256":"38b3f53f2dbada2776fd8ada8cae9add20e7413355a4613ce4055f307bacaa97"},{"filename":"scripts/prompts/report_agent.py","content":"# src/prompts/report_agent.py\nfrom datetime import datetime\nfrom typing import Optional\nfrom .isq_prompt_generator import generate_isq_prompt_section\n\ndef get_report_planner_base_instructions() -> str:\n \"\"\"生成报告策划员 (Planner) 的基础系统指令\"\"\"\n return \"\"\"你是一名资深的金融研报主编。你的任务是规划报告的结构，将零散的信号聚类成有逻辑的主题。\n你拥有 RAG 搜索工具，可以检索已生成的章节内容以确保逻辑连贯性。\n在规划时，应重点关注信号之间的关联性、产业链的完整性以及用户特定的关注点。\"\"\"\n\ndef get_report_writer_base_instructions() -> str:\n \"\"\"生成报告撰写员 (Writer) 的基础系统指令\"\"\"\n return \"\"\"你是一名资深金融分析师。你的任务是根据策划员提供的信号簇撰写深度研报章节。\n你应当运用专业的金融知识，将信号转化为深刻的洞察。\n注意：你没有外部搜索工具，你的分析必须基于提供给你的信号内容和行情数据。\"\"\"\n\ndef get_report_editor_base_instructions() -> str:\n \"\"\"生成报告编辑 (Editor) 的基础系统指令\"\"\"\n return \"\"\"你是一名严谨的金融研报编辑。你的任务是审核和润色撰写员生成的章节。\n你拥有 RAG 搜索工具，可以检索其他章节的内容，以消除重复、修正逻辑冲突并确保术语一致性。\n你应当确保报告符合专业的金融写作规范，且标题层级正确。\"\"\"\n\n# 1. 策划阶段 (Structural Planning)\ndef format_signal_for_report(signal: any, index: int, cite_keys: Optional[list] = None) -> str:\n \"\"\"格式化单个信号供研报生成使用\"\"\"\n # 这里的逻辑从 ReportAgent._format_signal_input 迁移过来\n from ..schema.models import InvestmentSignal\n \n if isinstance(signal, dict):\n try:\n sig_obj = InvestmentSignal(**signal)\n except:\n return f\"--- 信号 [{index}] ---\\n标题: {signal.get('title')}\\n内容: {signal.get('content', '')[:500]}\"\n else:\n sig_obj = signal\n\n chain_str = \" -> \".join([f\"{n.node_name}({n.impact_type})\" for n in sig_obj.transmission_chain])\n \n text = f\"--- 信号 [{index}] ---\\n\"\n text += f\"标题: {sig_obj.title}\\n\"\n text += f\"逻辑摘要: {sig_obj.summary}\\n\"\n text += f\"传导链条: {chain_str}\\n\"\n text += f\"ISQ 评分: 情绪({sig_obj.sentiment_score}), 确定性({sig_obj.confidence}), 强度({sig_obj.intensity})\\n\"\n text += f\"预期博弈: 时窗({sig_obj.expected_horizon}), 预期差({sig_obj.price_in_status})\\n\"\n \n tickers = \", \".join([f\"{t.get('name')}({t.get('ticker')})\" for t in sig_obj.impact_tickers])\n if tickers:\n text += f\"受影响标的: {tickers}\\n\"\n\n # Stable bibliography-style citation keys (LaTeX/BibTeX-like)\n if cite_keys:\n joined = \" \".join([f\"[@{k}]\" for k in cite_keys if k])\n if joined:\n text += f\"引用: {joined}\\n\"\n \n return text\n\ndef get_cluster_planner_instructions(signals_text: str, user_query: str = None) -> str:\n \"\"\"生成信号聚类指令 - 将零散信号组织成逻辑主题\"\"\"\n query_context = f\"用户重点关注：{user_query}\" if user_query else \"\"\n return f\"\"\"你是一位资深的金融研报主编。你的任务是将以下零散的金融信号聚类成 3-5 个核心逻辑主题，以便撰写一份结构清晰的研报。\n \n {query_context}\n\n ### 输入信号列表\n {signals_text}\n\n ### 聚类要求\n 1. **主题聚合**: 将相关性强的信号归为一组（例如：都涉及“建筑安全法规”或“某产业链上下游”）。\n 2. **叙事逻辑**: 只需要生成主题名称和包含的信号 ID。\n 3. **控制数量**: 将所有信号归类到 3-5 个主要主题中，不要遗漏。\n \n ### 输出格式 (JSON)\n 请仅输出以下 JSON 格式，不要包含 Markdown 标记：\n {{\n \"clusters\": [\n {{\n \"theme_title\": \"主题名称（如：建筑安全法规收紧引发的产业链重构）\",\n \"signal_ids\": [1, 3, 5],\n \"rationale\": \"这些信号都指向政府对高层建筑防火标准的政策调整...\"\n }},\n ...\n ]\n }}\n \"\"\"\n\ndef get_report_planner_instructions(toc: str, signal_count: int, user_query: str = None) -> str:\n \"\"\"生成报告规划指令 - 重点在于逻辑关联与分歧识别\"\"\"\n # ... (原有逻辑保持不变，但实际在新的聚类流程后这个可能作为备用或二次优化)\n query_context = f\"用户重点关注：{user_query}\" if user_query else \"\"\n return f\"\"\"你是一位资深的金融研报主编。你的任务是根据现有的草稿章节，规划出一份逻辑严密、穿透力强的终稿结构。\n \n ### 任务核心：\n 1. **识别主线**: 从草稿中识别出贯穿多个章节的“核心逻辑主线”（如：产业链共振、货币政策转向）。\n 2. **分歧评估 (Entropy)**: 识别各章节中观点冲突或确定性不一之处，规划如何在正文中呈现这些“分歧点”。\n 3. **结构蓝图**: \n - 定义一级标题（逻辑主题）。\n - 归类章节：哪些信号应放入同一主题下深度解析？\n - 排序：将 ISQ 强度最高、与{query_context}最相关的信号置前。\n\n ### 现有草稿目录 (TOC)\n {toc}\n\n 请输出你的【终稿修订大纲】（Markdown 格式）。\n \"\"\"\n\n# 2. 撰写阶段 (Section Writing)\ndef get_report_writer_instructions(theme_title: str, signal_cluster_text: str, signal_indices: list, price_context: str = \"\", user_query: str = None) -> str:\n \"\"\"生成 Writer Agent 指令 - 基于主题聚类撰写综合分析\"\"\"\n \n price_info = f\"\\n### 近期价格参考\\n{price_context}\\n\" if price_context else \"\"\n query_context = f\"\\n**用户意图**: \\\"{user_query}\\\"\\n请确保分析内容回应了用户的关注点。\\n\" if user_query else \"\"\n isq_block = generate_isq_prompt_section(include_header=False)\n \n # Keep citation scheme stable across re-ordering / edits.\n # Cite keys are provided in each signal block as: 引用: [@KEY]\n\n return f\"\"\"你是一位资深金融分析师。请针对核心主题 **\"{theme_title}\"** 撰写一篇深度研报章节。\n {query_context}\n\n ### 输入信号集 (本章节需综合的信号)\n {signal_cluster_text}\n {price_info}\n \n ### ISQ 评分说明\n {isq_block}\n \n ### 写作要求\n 1. **叙事逻辑**: 不要罗列信号，要将这些信号编织成一个连贯的故事。先讲宏观/行业背景，再讲具体事件传导，最后落脚到个股/标的影响。\n 2. **量化支撑**: 引用 ISQ 评分（确定性、强度、预期差）来佐证你的观点。关键观点必须关联相应的 ISQ 分值。\n 3. **引用规范（稳定 CiteKey）**: 关键论断必须标注来源引用，使用 `[@CITE_KEY]` 格式。\n - CiteKey 已在输入信号块中以 `引用: [@KEY]` 提供，请直接复制使用。\n - 不要使用 `[[1]]` 这类不稳定编号。\n 4. **关联标的预测**: **必须**在章节末尾明确给出受影响标的的预测分析，包括：\n - 至少列出 1-2 个相关上市公司代码（如 600519.SH）\n - 给出短期（T+3或T+5）的方向性判断\n - 如果可能，给出预期价格区间或涨跌幅预测\n \n ### 【重要】标题层级规范\n \n ❌ **错误示例**（绝对不要这样）：\n ```markdown\n # {theme_title}\n \n ### 宏观背景\n ...\n ```\n \n ✅ **正确示例**（必须这样）：\n ```markdown\n ## {theme_title}\n \n ### 宏观背景\n \n 近期全球经济环境...\n \n ### 具体传导机制分析\n \n ...\n \n ### 核心标的分析\n \n 建议关注：贵州茅台（600519.SH）...\n ```\n \n **关键要求**：\n - 章节主标题使用 `##` (H2)\n - 章节子标题使用 `###` (H3)\n - **绝对禁止**使用 `#` (H1)\n - 第一行必须是 `## {theme_title}` 开头\n\n ### 核心：图表叙事 (Visual Storytelling)\n **必须**在文中插入至少 1-2 个图表，且图表必须与上下文紧密结合（不要堆砌在末尾）。\n \n ### 宏观背景\n ...\n ```\n \n ✅ **正确示例**（必须这样）：\n ```markdown\n ## {theme_title}\n \n ### 宏观背景\n \n 近期全球经济环境...\n \n ### 具体传导机制分析\n \n ...\n \n ### 核心标的分析\n \n 建议关注：贵州茅台（600519.SH）...\n ```\n \n **关键要求**：\n - 章节主标题使用 `##` (H2)\n - 章节子标题使用 `###` (H3)\n - **绝对禁止**使用 `#` (H1)\n - 第一行必须是 `## {theme_title}` 开头\n\n ### 核心：图表叙事 (Visual Storytelling)\n **必须**在文中插入至少 1-2 个图表，且图表必须与上下文紧密结合（不要堆砌在末尾）。\n \n **可选图表类型 (请根据内容选择最合适的 1-2 种):**\n\n **A. AI 预测 + 走势 (Forecast) - 【强烈推荐 / 最新规范】**\n *适用*: 当文中明确提及某上市公司时，**必须**使用此图表展示股价走势与 AI 预测。\n *必填字段*:\n - `ticker`: 股票代码，A股 6 位 / 港股 5 位，允许带后缀（如 \"002371.SZ\"、\"9868.HK\"）\n - `pred_len`: 预测交易日长度（建议 3 或 5）\n *代码示例*:\n ```json-chart\n {{\"type\": \"forecast\", \"ticker\": \"002371.SZ\", \"title\": \"北方华创（002371）T+5 预测\", \"pred_len\": 5}}\n ```\n **重要**：禁止手写 `prediction` 数组（预测由系统自动生成并渲染）。\n *注意*: 如果提及多只股票，应为每只生成独立的 forecast 图表。\n\n **【推荐写法：多情景 → 最终归因 → 产出唯一预测图】**\n 你可以在正文里描述多种情景（如：基准/乐观/悲观），但在插入预测图之前，必须明确给出“本报告最终选择的最可能情景”及其归因，然后用 `forecast` 图表做最终总结。\n 为了让系统把“最终归因”可靠地传递给预测模块，请在 `forecast` JSON 中可选补充以下字段（字段均为可选，越完整越好）：\n - `selected_scenario`: 最可能情景名称（如 \"基准\" / \"乐观\" / \"悲观\"）\n - `selection_reason`: 选择该情景的归因理由（1-3 句）\n - `scenarios`: 情景列表（数组），每个元素可包含 `name`、`description`、`probability`（0-1）\n *示例*:\n ```json-chart\n {{\n \"type\": \"forecast\",\n \"ticker\": \"002371.SZ\",\n \"title\": \"北方华创（002371）T+5 预测（基准情景）\",\n \"pred_len\": 5,\n \"selected_scenario\": \"基准\",\n \"selection_reason\": \"结合订单能见度与行业景气，基准情景概率最高；短期扰动主要来自估值与市场风险偏好。\",\n \"scenarios\": [\n {{\"name\": \"乐观\", \"description\": \"国产替代与资本开支超预期\", \"probability\": 0.25}},\n {{\"name\": \"基准\", \"description\": \"订单稳健、利润率小幅波动\", \"probability\": 0.55}},\n {{\"name\": \"悲观\", \"description\": \"需求回落或交付节奏放缓\", \"probability\": 0.20}}\n ]\n }}\n ```\n\n **B. 历史走势 (Stock) - 仅作为兼容兜底**\n *适用*: 当你无法给出预测时（例如无法确定标的），可仅展示历史走势。\n *代码示例*:\n ```json-chart\n {{\"type\": \"stock\", \"ticker\": \"002371\", \"title\": \"北方华创历史走势\"}}\n ```\n\n **C. 舆情情绪演变 (Sentiment Trend)**\n *适用*: 当讨论行业政策、突发事件（如“火灾”、“新规”）的民意变化时。\n *注意*: `keywords` 必须是事件核心词。\n *代码*:\n ```json-chart\n {{\"type\": \"sentiment\", \"keywords\": [\"建筑安全\", \"防火标准\"], \"title\": \"市场对防火新规的情绪演变\"}}\n ```\n\n **D. 逻辑传导链条 (Transmission Chain)**\n *适用*: 复杂的蝴蝶效应分析（支持分支结构）。\n *代码*:\n ```json-chart\n {{\n \"type\": \"transmission\",\n \"nodes\": [\n {{\"node_name\": \"突发火灾\", \"impact_type\": \"中性\", \"logic\": \"事件发端\"}},\n {{\"node_name\": \"监管收紧\", \"impact_type\": \"利空\", \"logic\": \"合规成本上升\", \"source\": \"突发火灾\"}},\n {{\"node_name\": \"设备升级\", \"impact_type\": \"利好\", \"logic\": \"采购需求释放\", \"source\": \"突发火灾\"}},\n {{\"node_name\": \"龙头受益\", \"impact_type\": \"利好\", \"logic\": \"市占率提升\", \"source\": \"设备升级\"}}\n ],\n \"title\": \"火灾事件的逻辑传导与分支\"\n }}\n ```\n *说明*: 使用 `source` 字段指定父节点名称以创建分支结构。\n \n **E. 信号质量评估 (ISQ Radar)**\n *适用*: 对某个关键信号进行多维度（确定性、预期差等）定性评估时。\n *代码*:\n ```json-chart\n {{\"type\": \"isq\", \"sentiment\": 0.8, \"confidence\": 0.9, \"intensity\": 4, \"expectation_gap\": 0.7, \"timeliness\": 0.9, \"title\": \"核心信号质量评估\"}}\n ```\n \"\"\"\n\n# 3. 整合阶段 (Final Assembly) - 原版，保留用于 fallback\ndef get_report_editor_instructions(draft_sections: str, plan: str, sources_list: str) -> str:\n \"\"\"生成最终编辑指令 - 根据规划蓝图重组内容\"\"\"\n return f\"\"\"你是一位专业的研报编辑。请将以下基于主题撰写的草稿章节整合成最终研报。\n \n ### 原始草稿内容\n {draft_sections}\n\n ### 原始引用来源\n {sources_list}\n\n ### 任务与要求\n 1. **结构化**: 为每个草稿章节添加合适的 Markdown 标题 (## 级别)。\n 2. **连贯性**: 确保章节之间过渡自然。\n 3. **完整性**:\n - 必须保留所有 `json-chart` 代码块（图表配置）。\n - 必须保留引用标注 `[@CITE_KEY]`。\n - 生成 `## 核心观点摘要`、`## 参考文献` 和 `## 风险提示`。\n\n ### 输出\n 只输出最终的 Markdown 研报内容。\n \"\"\"\n\n\n# 4. 单节编辑 (Incremental Section Editing with RAG)\ndef get_section_editor_instructions(section_index: int, total_sections: int, toc: str) -> str:\n \"\"\"生成单节编辑 prompt，支持 RAG 工具调用\"\"\"\n return f\"\"\"你是一位研报编辑。你正在编辑报告的第 {section_index}/{total_sections} 节。\n\n ### 当前目录 (TOC)\n {toc}\n\n ### 你的任务\n 1. 润色当前章节内容，确保逻辑清晰、语言专业。\n 2. 保留所有 `[@CITE_KEY](#ref-CITE_KEY)` 或 `[@CITE_KEY]` 格式的引用。\n 3. 保留所有 `json-chart` 代码块，不做修改。\n 4. 如果需要参考其他章节内容，使用 `search_context` 工具搜索。\n 5. 只输出编辑后的章节内容，不要输出其他章节。\n \n ### 【关键】标题层级规范\n **严格遵守以下规则：**\n - 章节主标题使用 `##` (H2)\n - 章节子标题使用 `###` (H3)\n - **禁止使用** `#` (H1) - 只有报告大标题可以使用 H1\n - 如果原文中有 H1，必须将其降级为 H2\n - 不要输出与 \"参考文献\"、\"风险提示\" 相同的标题\n\n 直接输出编辑后的 Markdown 内容。\n \"\"\"\n\n\n# 5. 摘要生成 (Summary Generation)\ndef get_summary_generator_instructions(toc: str, section_summaries: str) -> str:\n \"\"\"生成报告摘要指令 - 包含市场分歧度分析\"\"\"\n return f\"\"\"你是一位资深研报主笔。请生成今日报告的核心观点摘要的**正文内容**。\n\n ### 章节摘要\n {section_summaries}\n\n ### 任务：\n 1. **核心逻辑提炼**: 用 150 字以内总结今日最核心的投资主线。\n 2. **分歧识别**: 如果不同信号对同一板块有冲突观点，请明确指出\"市场分歧点\"。\n 3. **确定性排序**: 标记出今日确定性最高的前两个机会（需列出具体标的代码）。\n\n ### 【重要】输出格式规范：\n \n ❌ **错误示例**（不要遗漏二级标题）：\n ```markdown\n ### 核心逻辑提炼\n ...\n ```\n \n ✅ **正确示例**（应该这样输出）：\n ```markdown\n ## 核心观点摘要\n\n ### 核心逻辑提炼\n \n 科技自立战略加速半导体设备国产化，叠加AI算力需求爆发...\n \n ### 市场分歧点\n \n 资本市场波动显示医药、新能源等板块估值逻辑受政策敏感性增强...\n \n ### 确定性排序\n \n 1. **网络安全替代需求**（ISQ确定性0.85，推荐标的：深信服 300454.SZ）\n 2. **半导体设备材料**（ISQ确定性0.75，推荐标的：北方华创 002371.SZ）\n ```\n \n ### 关键要求：\n - 第一行必须是 `## 核心观点摘要`\n - 主体部分使用 H3 (`###`) 和 H4 (`####`) 级别标题\n - **必须**包含 `## 核心观点摘要` 这一级标题\n \n 现在请按照正确示例的格式输出摘要内容。\n \"\"\"\n\n\n# 6. 最终组装 (Final Assembly with Sections)\ndef get_final_assembly_instructions(sources_list: str) -> str:\n \"\"\"生成最终报告组装的 prompt\"\"\"\n return f\"\"\"你是一位研报主笔。请完成以下任务：\n\n ### 任务\n 1. 生成 \"## 参考文献\" 章节（需要按照顺序，顺序不对时进行调整）：\n - 原始来源：\n {sources_list}\n - 格式：`\u003ca id=\"ref-CITE_KEY\">\u003c/a>[@CITE_KEY] 标题 (来源), [链接地址]`\n 2. 生成 \"## 风险提示\" (标准免责声明)。\n 3. 生成 \"## 快速扫描\" 表格，汇总各主题的核心观点。\n - 表格列：**主题**, **核心观点**, **强度(Intensity)**, **确定性(Confidence)**。\n - 强度和确定性请参考原章节中的 ISQ 评分。\n\n 只输出上述三个章节的 Markdown 内容。\n \"\"\"\n\ndef get_cluster_task(signals_preview: str) -> str:\n \"\"\"生成聚类任务描述\"\"\"\n return f\"请对以下信号进行主题聚类：\\n\\n{signals_preview}\"\n\ndef get_writer_task(theme_title: str) -> str:\n \"\"\"生成撰写任务描述\"\"\"\n return f\"请依据主题 '{theme_title}' 和输入信号集开始撰写深度分析章节。\"\n\ndef get_planner_task() -> str:\n \"\"\"生成规划任务描述\"\"\"\n return \"请阅读现有草稿并规划终稿大纲，识别核心逻辑主线和市场分歧点。\"\n\ndef get_editor_task() -> str:\n \"\"\"生成编辑任务描述\"\"\"\n return \"请根据规划大纲和草稿内容，生成最终研报。确保逻辑连贯，保留所有图表和引用。\"\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":18529,"content_sha256":"fa0c4ad1d6b0960c77faf75dea8e8498e229007476571c6a575c73560a55ff7e"},{"filename":"scripts/prompts/trend_agent.py","content":"from typing import Any\nfrom datetime import datetime\nfrom .isq_prompt_generator import generate_isq_prompt_section\n\ndef get_trend_scanner_instructions() -> str:\n \"\"\"生成趋势扫描员 (Scanner) 的系统指令\"\"\"\n current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n return f\"\"\"你是一名专业的数据扫描员，当前时间是 {current_time}。\n你的任务是利用各种工具从互联网和数据库中获取最新的金融新闻、热点趋势和市场数据。\n\n### 1. 核心职责\n1. **多源采集**: 使用 `news_toolkit` 获取最新新闻，使用 `stock_toolkit` 获取行情，使用 `polymarket_toolkit` 获取预测市场数据。\n2. **情绪感知**: 使用 `sentiment_toolkit` 对关键新闻进行情绪分析。\n3. **深度搜索**: 针对模糊的热点，使用 `search_toolkit` 进行全网搜索补充细节。\n\n### 2. 工具使用规范\n- **广度优先**: 尽可能覆盖多个数据源。\n- **数据新鲜度**: 优先获取最近 24 小时内的信息。\n- **结构化输出**: 整理搜集到的原始数据，为后续评估提供清晰的素材。\n\"\"\"\n\ndef get_trend_evaluator_instructions() -> str:\n \"\"\"生成趋势评估员 (Evaluator) 的系统指令\"\"\"\n current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n isq_block = generate_isq_prompt_section(include_header=True)\n\n return f\"\"\"\n 你是一名顶级的金融情报专家 (TrendAgent)，擅长从海量信息中识别具有深度价值的\"二级市场投资信号\"。\n 当前时间：{current_time}\n\n ### 核心使命：\n 不仅是发现\"热点\"，更要解析\"信号\"。你需要识别那些能触发**传导链条 (Transmission Chain)** 且具有**高确定性 (Confidence)** 的事件。\n\n {isq_block}\n\n ### 核心能力与标准：\n 1. **信号识别 (Signal Discovery)**: 基于扫描员提供的素材，识别具有投资价值的信号。优先关注政策、产业变革、重大诉求及跨境套利机会。\n 2. **逻辑相干性**: 是否具备清晰的\"原因-结果\"传导？\n 3. **影响力系数**: 是否会引发板块性的联动或财务指标的实质性扰动？\n 4. **市场认知差**: 市场是否已提前消化（Price-in）？寻找尚未被充分交易的\"Alpha\"。\n 5. **实体穿透**: 必须关联到具体的 Ticker 或核心产业链节点。\n\n ### 严禁事项：\n - 严禁编造数据。\n - 严禁仅输出情绪极性（Positive/Negative），必须带有逻辑依据。\n - 严禁将纯娱乐或单纯的社会负面事件（除非具有宏观破坏性）视为金融信号。\n\n ### 输出要求：\n 你发现的每个信号应包含：\n - **核心摘要**: 穿透表象的逻辑总结。\n - **传导节点**: A -> B -> C 的逻辑推导。\n - **推荐关注**: 板块或 Ticker。\n - **ISQ 评估**: 基于模板的 5 个维度进行初步评分（具体评分由后续 FinAgent 完成）。\n \"\"\"\n\ndef get_trend_agent_instructions() -> str:\n # 保持兼容性\n return get_trend_evaluator_instructions()\n\ndef get_trend_scan_task(task_description: str) -> str:\n \"\"\"生成扫描员的任务描述\"\"\"\n return f\"请根据以下任务描述，搜集相关的原始数据和新闻：\\n\\n{task_description}\"\n\ndef format_scan_context(scan_data: dict) -> str:\n \"\"\"将扫描员搜集的结构化数据格式化为评估员可读的文本\"\"\"\n if not scan_data:\n return \"（未能搜集到原始数据）\"\n \n return f\"\"\"\n### 扫描数据概览\n- **热点话题**: {', '.join(scan_data.get('hot_topics', []))}\n- **情绪概览**: {scan_data.get('sentiment_overview', '未知')}\n- **关键新闻**: {len(scan_data.get('news_summaries', []))} 条\n- **数据摘要**: {scan_data.get('raw_data_summary', '无')}\n\"\"\"\n\ndef get_trend_eval_task(task_description: str, raw_data_str: str) -> str:\n \"\"\"生成评估员的任务描述\"\"\"\n return f\"\"\"请基于以下搜集到的原始数据，完成最终的分析任务：\n \n任务描述: {task_description}\n\n原始数据:\n{raw_data_str}\n\n请识别出最具金融价值的信号，并给出评估理由。\"\"\"\n\ndef get_news_filter_instructions(news_count: int, depth: Any, user_query: str = None) -> str:\n \"\"\"生成新闻筛选 prompt，使用 FilterResult schema 加快推理并减少 token 消耗\n \n Args:\n news_count: 输入新闻总数\n depth: 目标筛选数量，若为 auto 则由 LLM 自主判断\n user_query: 用户输入的查询/关注点（可选）\n \"\"\"\n \n # 1. 深度控制逻辑\n if str(depth).lower() == 'auto':\n depth_guide = \"的数量不设固定限制（建议 3-10 条），根据新闻含金量自动判断\"\n limit_instruction = \"宁缺毋滥，如果高价值信息很少，可以只选 1-2 条；如果都很重要，可以多选。\"\n else:\n try:\n d_int = int(depth)\n depth_guide = f\"约 {d_int} 条\"\n limit_instruction = f\"请尽量凑满 {d_int} 条，但如果剩余新闻全是噪音，则不必强行凑数。\"\n except:\n depth_guide = \"适量\"\n limit_instruction = \"根据内容价值判断。\"\n\n target_desc = f\"筛选出最具投资分析价值的新闻（{depth_guide}）。\"\n \n # 2. 用户意图逻辑\n query_instruction = \"\"\n if user_query:\n target_desc = f\"筛选出与用户意图【{user_query}】最相关的新闻。\"\n query_instruction = f\"\"\"\n ### 核心任务（High Priority）：\n 用户明确关注：\"{user_query}\"。\n 1. **第一优先级**：必须包含所有与\"{user_query}\"直接或间接相关的新闻，不要遗漏。\n - 即使这些新闻看起来\"价值不高\"，只要相关都要保留。\n 2. **第二优先级**：在满足第一优先级后，如果名额未满，再补充其他重大的市场热点。\n \"\"\"\n\n return f\"\"\"你是一名专业的金融情报精排师。你需要从给定的 {news_count} 条原始新闻流中，{target_desc}\n\n {query_instruction}\n\n ### FSD (Financial Signal Density) 筛选准则：\n 1. **逻辑传导性 (Transmission)**: 该新闻是否预示着一个明确的产业链传导逻辑？（如：上游涨价 -> 中游成本压力 -> 下游提价预期）\n 2. **预期差 (Alpha Potential)**: 是否包含尚未被市场充分Price-in的新突发情况？\n 3. **确定性 (Confidence)**: 信息来源是否权威？是否包含具体的财务数据、订单金额或明确的政策日期？\n 4. **排除噪音**: 坚决剔除明星八卦、鸡汤文、以及无实质增量的\"口号式\"新闻。\n\n ### {limit_instruction}\n\n ### 快速有效性检查（TOKEN 优化）：\n 在开始详细筛选前，先快速判断：这 {news_count} 条新闻中是否至少包含 1 条有效的金融信号？\n - 如果全是无关内容（如体育、娱乐、纯生活信息），直接返回 \"has_valid_signals\": false\n - 如果有至少 1 条金融相关的新闻，再进行详细 FSD 筛选\n\n ### 输出格式（必须为 JSON，使用 FilterResult schema）：\n ```json\n {{\n \"has_valid_signals\": true/false,\n \"selected_ids\": [\"id_1\", \"id_2\", ...],\n \"themes\": [\n {{\n \"name\": \"高概括性主题\",\n \"news_ids\": [\"相关id_1\", ...],\n \"fsd_reason\": \"基于 FSD 准则的筛选理由，重点描述传导逻辑和预期差。\"\n }}\n ],\n \"reason\": \"如果 has_valid_signals=false，简要说明原因。否则可为空。\"\n }}\n ```\n \"\"\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":7534,"content_sha256":"73c12f7198857ba0f779bb1045b73ab9653d143806d37ff2801ef36b2151a4a9"},{"filename":"scripts/prompts/visualizer.py","content":"def get_drawio_system_prompt():\n return \"\"\"You are an expert at creating Draw.io (MxGraph) diagrams in XML format.\nYour task is to generate a valid MXGraphModel XML based on the user's description.\n\n### Rules:\n1. Output ONLY the XML code. Start with \u003cmxGraphModel> and end with \u003c/mxGraphModel>.\n2. Do not use compressed XML. Use plain XML.\n3. Use standard shapes: 'rounded=1;whiteSpace=wrap;html=1;' for boxes.\n4. Auto-layout Strategy:\n - Identify \"layers\" or \"stages\" in the logic.\n - Assign X coordinates based on layers (e.g., 0, 200, 400).\n - Assign Y coordinates to distribute nodes vertically (e.g., 0, 100, 200).\n - Ensure nodes do not overlap.\n5. Edges: Connect nodes logically using \u003cmxCell edge=\"1\" ...>.\n\n### Template:\n\u003cmxGraphModel dx=\"1000\" dy=\"1000\" grid=\"1\" gridSize=\"10\" guides=\"1\" tooltips=\"1\" connect=\"1\" arrows=\"1\" fold=\"1\" page=\"1\" pageScale=\"1\" pageWidth=\"827\" pageHeight=\"1169\" math=\"0\" shadow=\"0\">\n \u003croot>\n \u003cmxCell id=\"0\"/>\n \u003cmxCell id=\"1\" parent=\"0\"/>\n \n \u003c!-- Node -->\n \u003cmxCell id=\"n1\" value=\"Node Label\" style=\"rounded=1;whiteSpace=wrap;html=1;fillColor=#dae8fc;strokeColor=#6c8ebf;\" vertex=\"1\" parent=\"1\">\n \u003cmxGeometry x=\"100\" y=\"100\" width=\"120\" height=\"60\" as=\"geometry\"/>\n \u003c/mxCell>\n \n \u003c!-- Edge -->\n \u003cmxCell id=\"e1\" value=\"Connection\" style=\"edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;\" edge=\"1\" parent=\"1\" source=\"n1\" target=\"n2\">\n \u003cmxGeometry relative=\"1\" as=\"geometry\"/>\n \u003c/mxCell>\n \u003c/root>\n\u003c/mxGraphModel>\n\"\"\"\n\ndef get_drawio_task(nodes_data: list, title: str) -> str:\n import json\n nodes_json = json.dumps(nodes_data, ensure_ascii=False, indent=2)\n return f\"\"\"Please generate a Draw.io XML diagram for the following logic flow:\n\n**Title**: {title}\n\n**Nodes and Logic**:\n{nodes_json}\n\nEnsure the layout flows logically from Left to Right (or Top to Bottom for hierarchies).\nUse different colors for 'Positive' (Greenish), 'Negative' (Reddish), and 'Neutral' (Grey/Blue) impacts if described.\n\"\"\"\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2028,"content_sha256":"3770fb21c1a8d1dd34c1ff5d614061e4194ad18d020ef89c7be116ccab7012c3"},{"filename":"scripts/schema/isq_template.py","content":"\"\"\"\nISQ (Investment Signal Quality) 评估框架 Template\n\n统一定义 ISQ 的各个维度、评分标准、和使用方法。\n支持默认 template 和自定义 template。\n\"\"\"\n\nfrom typing import Dict, List, Any, Optional\nfrom pydantic import BaseModel, Field\nfrom enum import Enum\nfrom pathlib import Path\nimport json\n\n\nclass ISQDimension(str, Enum):\n \"\"\"ISQ 评估维度\"\"\"\n SENTIMENT = \"sentiment\" # 情绪/走势方向\n CONFIDENCE = \"confidence\" # 确定性/可信度\n INTENSITY = \"intensity\" # 强度/影响量级\n EXPECTATION_GAP = \"expectation_gap\" # 预期差/市场认知差\n TIMELINESS = \"timeliness\" # 时效性/窗口紧迫度\n TRANSMISSION = \"transmission\" # 逻辑传导清晰度\n\n\nclass ISQDimensionSpec(BaseModel):\n \"\"\"ISQ 单个维度的定义规范\"\"\"\n name: str = Field(..., description=\"维度名称\")\n key: str = Field(..., description=\"维度键名\")\n description: str = Field(..., description=\"维度描述\")\n range_type: str = Field(default=\"0-1\", description=\"取值范围 (0-1 或 1-5 等)\")\n scale_factor: float = Field(default=1.0, description=\"显示时的缩放因子\")\n examples: Dict[str, str] = Field(default_factory=dict, description=\"不同分值的示例解释\")\n visualization_color: Optional[str] = Field(default=None, description=\"可视化颜色\")\n\n\nclass ISQTemplate(BaseModel):\n \"\"\"ISQ 评估框架 Template\"\"\"\n template_id: str = Field(..., description=\"模板 ID\")\n template_name: str = Field(..., description=\"模板名称\")\n description: str = Field(..., description=\"模板描述\")\n \n # 核心维度定义\n dimensions: Dict[str, ISQDimensionSpec] = Field(..., description=\"维度定义字典\")\n \n # 评分指导\n scoring_guide: str = Field(..., description=\"评分指导说明\")\n \n # 应用场景\n applicable_scenarios: List[str] = Field(default_factory=list, description=\"适用场景\")\n \n # 聚合算法\n aggregation_method: str = Field(default=\"weighted_average\", description=\"聚合方法 (weighted_average, product 等)\")\n dimension_weights: Dict[str, float] = Field(default_factory=dict, description=\"维度权重\")\n\n\nclass ISQScore(BaseModel):\n \"\"\"单个信号的 ISQ 评分结果\"\"\"\n signal_id: str = Field(..., description=\"信号 ID\")\n template_id: str = Field(..., description=\"使用的模板 ID\")\n \n # 各维度评分\n scores: Dict[str, float] = Field(..., description=\"各维度评分\")\n \n # 总分\n overall_score: float = Field(..., description=\"综合评分\")\n \n # 评分理由\n rationale: Dict[str, str] = Field(default_factory=dict, description=\"各维度评分理由\")\n \n # 时间戳\n timestamp: str = Field(..., description=\"评分时间\")\n\n\n# =====================================================\n# 默认 Template\n# =====================================================\n\nDEFAULT_ISQ_TEMPLATE = ISQTemplate(\n template_id=\"default_isq_v1\",\n template_name=\"标准投资信号质量评估框架 (ISQ v1.0)\",\n description=\"AlphaEar 默认的 ISQ 评估框架，用于标准化评估投资信号的质量维度\",\n \n dimensions={\n \"sentiment\": ISQDimensionSpec(\n name=\"情绪/走势\",\n key=\"sentiment\",\n description=\"基础情绪偏向和市场走势判断\",\n range_type=\"-1.0 到 1.0\",\n scale_factor=1.0,\n examples={\n \"-1.0\": \"极度悲观/极度看空\",\n \"-0.5\": \"明显看空\",\n \"0.0\": \"中性/没有明确方向\",\n \"0.5\": \"明显看多\",\n \"1.0\": \"极度乐观/极度看多\"\n },\n visualization_color=\"#ef4444\" # 红色表示负面，绿色表示正面\n ),\n \n \"confidence\": ISQDimensionSpec(\n name=\"确定性\",\n key=\"confidence\",\n description=\"信号的可信度和确定性程度\",\n range_type=\"0.0 到 1.0\",\n scale_factor=1.0,\n examples={\n \"0.0-0.3\": \"信息来源不可靠/传言多/逻辑推导牵强\",\n \"0.3-0.6\": \"信息相对可靠/有一定逻辑/但仍有不确定性\",\n \"0.6-0.8\": \"信息来源权威/逻辑清晰/高度可信\",\n \"0.8-1.0\": \"官方确认/数据明确/完全确定\"\n },\n visualization_color=\"#3b82f6\" # 蓝色\n ),\n \n \"intensity\": ISQDimensionSpec(\n name=\"强度/影响量级\",\n key=\"intensity\",\n description=\"信号对相关板块/个股的潜在影响程度\",\n range_type=\"1 到 5\",\n scale_factor=20.0, # 用于雷达图缩放 (5 -> 100)\n examples={\n \"1\": \"影响微弱，可能被市场忽略\",\n \"2\": \"小幅影响，短期可能有波动\",\n \"3\": \"中等影响，值得重点关注\",\n \"4\": \"强烈影响，可能成为市场焦点\",\n \"5\": \"极强影响，市场预期明显变化\"\n },\n visualization_color=\"#f97316\" # 橙色\n ),\n \n \"expectation_gap\": ISQDimensionSpec(\n name=\"预期差\",\n key=\"expectation_gap\",\n description=\"市场预期与现实之间的差距\",\n range_type=\"0.0 到 1.0\",\n scale_factor=1.0,\n examples={\n \"0.0-0.2\": \"市场充分认知，预期差小\",\n \"0.2-0.5\": \"市场部分认知，存在一定预期差\",\n \"0.5-0.8\": \"市场认知不足，预期差较大，存在博弈空间\",\n \"0.8-1.0\": \"市场严重低估/高估，巨大预期差\"\n },\n visualization_color=\"#22c55e\" # 绿色\n ),\n \n \"timeliness\": ISQDimensionSpec(\n name=\"时效性\",\n key=\"timeliness\",\n description=\"信号的时间窗口紧迫度\",\n range_type=\"0.0 到 1.0\",\n scale_factor=1.0,\n examples={\n \"0.0-0.2\": \"长期信号，反应窗口 > 3 月\",\n \"0.2-0.5\": \"中期信号，反应窗口 1-3 月\",\n \"0.5-0.8\": \"短期信号，反应窗口 1 周 - 1 月\",\n \"0.8-1.0\": \"超短期信号，反应窗口 \u003c 1 周（需立即行动）\"\n },\n visualization_color=\"#a855f7\" # 紫色\n ),\n },\n \n scoring_guide=\"\"\"\n ### ISQ 评分指导 (Investment Signal Quality)\n \n ISQ 框架用于多维度评估投资信号的质量。每个信号由 5 个维度组成：\n \n 1. **情绪 (Sentiment)**: -1.0 到 1.0，表示看空(-)/中性(0)/看多(+)\n 2. **确定性 (Confidence)**: 0.0 到 1.0，数值越高越确定\n 3. **强度 (Intensity)**: 1 到 5，数值越高影响越大\n 4. **预期差 (Expectation Gap)**: 0.0 到 1.0，市场预期与现实的差距\n 5. **时效性 (Timeliness)**: 0.0 到 1.0，反应窗口的紧迫程度\n \n ### 综合评分算法\n \n 综合评分 = 确定性 × 0.35 + 强度/5 × 0.30 + 预期差 × 0.20 + 时效性 × 0.15\n \n 范围: 0.0 到 1.0\n - 0.0-0.3: 信号质量较差，不建议跟进\n - 0.3-0.6: 信号质量一般，可作参考\n - 0.6-0.8: 信号质量良好，值得跟进\n - 0.8-1.0: 信号质量优异，强烈推荐\n \n ### 评分时的注意事项\n \n - **不要混淆方向和强度**：情绪可以是看空，但确定性和强度仍可能很高\n - **预期差往往是 Alpha 来源**：高预期差 + 高确定性 = 最佳博弈机会\n - **考虑时间成本**：长期信号需要更高的确定性才值得跟进\n - **数据为王**：所有评分必须有具体数据支撑\n \"\"\",\n \n applicable_scenarios=[\n \"上市公司基本面变化分析\",\n \"产业政策与监管事件评估\",\n \"地缘政治与宏观经济影响\",\n \"技术进步与产业升级\",\n \"突发事件与应急响应\"\n ],\n \n aggregation_method=\"weighted_average\",\n dimension_weights={\n \"confidence\": 0.35,\n \"intensity\": 0.30,\n \"expectation_gap\": 0.20,\n \"timeliness\": 0.15\n }\n)\n\n\n# =====================================================\n# ISQ Template 管理系统\n# =====================================================\n\nclass ISQTemplateManager:\n \"\"\"ISQ Template 管理器\"\"\"\n \n def __init__(self):\n self.templates: Dict[str, ISQTemplate] = {\n DEFAULT_ISQ_TEMPLATE.template_id: DEFAULT_ISQ_TEMPLATE\n }\n \n def register_template(self, template: ISQTemplate) -> None:\n \"\"\"注册新的 template\"\"\"\n self.templates[template.template_id] = template\n\n def register_template_dict(self, template_dict: Dict[str, Any]) -> ISQTemplate:\n \"\"\"从 dict 注册模板，返回实例。\"\"\"\n tpl = ISQTemplate(**template_dict)\n self.register_template(tpl)\n return tpl\n \n def get_template(self, template_id: str) -> ISQTemplate:\n \"\"\"获取指定 template\"\"\"\n if template_id not in self.templates:\n return DEFAULT_ISQ_TEMPLATE\n return self.templates[template_id]\n \n def list_templates(self) -> List[Dict[str, str]]:\n \"\"\"列出所有可用 template\"\"\"\n return [\n {\n \"id\": t.template_id,\n \"name\": t.template_name,\n \"description\": t.description,\n \"dimensions\": list(t.dimensions.keys())\n }\n for t in self.templates.values()\n ]\n \n def get_dimension(self, template_id: str, dimension_key: str) -> ISQDimensionSpec:\n \"\"\"获取指定 template 的某个维度定义\"\"\"\n template = self.get_template(template_id)\n return template.dimensions.get(dimension_key)\n \n def get_scoring_prompt(self, template_id: str) -> str:\n \"\"\"获取用于 LLM 的评分 prompt\"\"\"\n template = self.get_template(template_id)\n \n dimensions_desc = \"\\n\".join([\n f\"- **{d.name} ({d.key})**\\n\"\n f\" 范围: {d.range_type}\\n\"\n f\" 说明: {d.description}\\n\"\n f\" 示例: {', '.join(f'{k}={v}' for k, v in list(d.examples.items())[:3])}\"\n for d in template.dimensions.values()\n ])\n \n return f\"\"\"\n### ISQ 评估指导 ({template.template_name})\n\n使用以下 {len(template.dimensions)} 个维度评估信号质量：\n\n{dimensions_desc}\n\n### 评分标准\n{template.scoring_guide}\n\n### 输出格式 (JSON)\n请输出以下 JSON 格式的评分结果：\n{{\n \"sentiment\": \u003cfloat>,\n \"confidence\": \u003cfloat>,\n \"intensity\": \u003cint>,\n \"expectation_gap\": \u003cfloat>,\n \"timeliness\": \u003cfloat>,\n \"rationale\": {{\n \"sentiment\": \"评分理由\",\n \"confidence\": \"评分理由\",\n \"intensity\": \"评分理由\",\n \"expectation_gap\": \"评分理由\",\n \"timeliness\": \"评分理由\"\n }}\n}}\n\"\"\"\n\n\n# 全局 template 管理器实例\nisq_template_manager = ISQTemplateManager()\n\n\n# =====================================================\n# 配置加载\n# =====================================================\n\ndef load_templates_from_config(config_path: Optional[str] = None) -> None:\n \"\"\"从配置目录加载所有 JSON 模板文件，未找到则跳过，不影响默认模板。\n 支持单个 JSON 文件或目录（目录下的所有 .json 文件）。\n \"\"\"\n if config_path:\n path = Path(config_path)\n else:\n # 默认目录：config/isq_templates/\n # __file__ = src/schema/isq_template.py\n # parent = src/schema, parent.parent = src, parent.parent.parent = 项目根目录\n path = Path(__file__).resolve().parent.parent.parent / \"config\"\n \n if not path.exists():\n return\n \n # 如果是目录，扫描所有 .json 文件\n if path.is_dir():\n json_files = list(path.glob(\"*.json\"))\n else:\n json_files = [path]\n \n for json_file in json_files:\n try:\n data = json.loads(json_file.read_text(encoding=\"utf-8\"))\n \n # 如果是单个模板对象，转为列表\n if isinstance(data, dict):\n templates = [data]\n elif isinstance(data, list):\n templates = data\n else:\n continue\n \n # 注册所有模板\n for tpl_dict in templates:\n if not isinstance(tpl_dict, dict):\n continue\n try:\n isq_template_manager.register_template_dict(tpl_dict)\n except Exception:\n # 忽略单个模板的加载错误，继续其他模板\n continue\n except Exception:\n # JSON 解析失败，跳过该文件\n continue\n\n\n# 在模块加载时自动尝试加载配置模板\nload_templates_from_config()\n\n\n# =====================================================\n# 便利函数\n# =====================================================\n\ndef get_isq_template(template_id: str = \"default_isq_v1\") -> ISQTemplate:\n \"\"\"获取 ISQ template\"\"\"\n return isq_template_manager.get_template(template_id)\n\n\ndef get_isq_scoring_prompt(template_id: str = \"default_isq_v1\") -> str:\n \"\"\"获取用于 LLM 的 ISQ 评分 prompt\"\"\"\n return isq_template_manager.get_scoring_prompt(template_id)\n\n\ndef calculate_isq_overall_score(scores: Dict[str, float], template_id: str = \"default_isq_v1\") -> float:\n \"\"\"计算 ISQ 综合评分\"\"\"\n template = get_isq_template(template_id)\n \n overall = 0.0\n for dim_key, weight in template.dimension_weights.items():\n if dim_key in scores:\n score = scores[dim_key]\n # 处理强度维度的特殊缩放 (1-5 -> 0-1)\n if dim_key == \"intensity\":\n score = score / 5.0\n overall += score * weight\n \n return min(1.0, max(0.0, overall)) # 限制在 0-1 之间\n","content_type":"text/x-python; charset=utf-8","language":"python","size":13967,"content_sha256":"2bf63b07a6b53600301a933535fa803a851e535c4686ee838b7dcfaae079dabc"},{"filename":"scripts/schema/models.py","content":"from pydantic import BaseModel, Field\nfrom typing import List, Optional, Dict, Any\nfrom datetime import datetime\n\nclass TransmissionNode(BaseModel):\n node_name: str = Field(..., description=\"产业链节点名称\")\n impact_type: str = Field(..., description=\"利好/利空/中性\")\n logic: str = Field(..., description=\"该节点的传导逻辑\")\n\nclass IntentAnalysis(BaseModel):\n keywords: List[str] = Field(..., description=\"核心实体、事件或概念关键词\")\n search_queries: List[str] = Field(..., description=\"优化后的搜索引擎查询词\")\n is_specific_event: bool = Field(..., description=\"是否查询特定突发事件\")\n time_range: str = Field(..., description=\"时间范围 (recent/all/specific_date)\")\n intent_summary: str = Field(..., description=\"一句话意图描述\")\n\nclass FilterResult(BaseModel):\n \"\"\"LLM 筛选结果 - 快速判断是否有有效信号\"\"\"\n has_valid_signals: bool = Field(..., description=\"列表中是否包含有效的金融信号\")\n selected_ids: List[int] = Field(default_factory=list, description=\"筛选出的有效信号 ID 列表\")\n themes: List[str] = Field(default_factory=list, description=\"信号涉及的主题\")\n reason: Optional[str] = Field(default=None, description=\"如果无有效信号，说明原因\")\n\nclass InvestmentSignal(BaseModel):\n # 核心元数据\n signal_id: str = Field(default=\"unknown_sig\", description=\"唯一信号 ID\")\n title: str = Field(..., description=\"信号标题\")\n summary: str = Field(default=\"暂无摘要分析\", description=\"100 字核心观点快报\")\n reasoning: str = Field(default=\"\", description=\"详细的推演逻辑和理由\")\n \n # 逻辑传导 (ISQ Key 1)\n transmission_chain: List[TransmissionNode] = Field(default_factory=list, description=\"产业链传导逻辑链条\")\n \n # 信号质量 (ISQ Key 2) - 来自 isq_template.DEFAULT_ISQ_TEMPLATE\n # 参考: src/schema/isq_template.py 的 DEFAULT_ISQ_TEMPLATE 定义\n sentiment_score: float = Field(default=0.0, description=\"[ISQ] 情绪/走势 (-1.0=极度看空 ~ 0.0=中性 ~ 1.0=极度看多)\")\n confidence: float = Field(default=0.5, description=\"[ISQ] 确定性 (0.0=不可信 ~ 1.0=完全确定)\")\n intensity: int = Field(default=3, description=\"[ISQ] 强度/影响量级 (1=微弱 ~ 5=极强)\")\n expectation_gap: float = Field(default=0.5, description=\"[ISQ] 预期差/博弈空间 (0.0=充分定价 ~ 1.0=巨大预期差)\")\n timeliness: float = Field(default=0.8, description=\"[ISQ] 时效性 (0.0=长期 ~ 1.0=超短期)\")\n \n # 预测与博弈 (ISQ Key 3)\n expected_horizon: str = Field(default=\"T+N\", description=\"预期的反应时窗 (如: T+0, T+3, Long-term)\")\n price_in_status: str = Field(default=\"未知\", description=\"市场预期消化程度 (未定价/部分定价/充分定价)\")\n \n # 关联实体\n impact_tickers: List[Dict[str, Any]] = Field(default_factory=list, description=\"受影响的代码列表及其权重\")\n industry_tags: List[str] = Field(default_factory=list, description=\"关联行业标签\")\n \n # 溯源\n sources: List[Dict[str, str]] = Field(default_factory=list, description=\"来源详情 (包含 title, url, source_name)\")\n\nclass ResearchContext(BaseModel):\n \"\"\"研究员搜集的背景信息结构\"\"\"\n raw_signal: str = Field(..., description=\"原始信号内容\")\n tickers_found: List[Dict[str, Any]] = Field(default_factory=list, description=\"找到的相关标的及其基本面/股价信息\")\n industry_background: str = Field(..., description=\"行业背景及产业链现状\")\n latest_developments: List[str] = Field(default_factory=list, description=\"相关事件的最新进展\")\n key_risks: List[str] = Field(default_factory=list, description=\"潜在风险点\")\n search_results_summary: str = Field(..., description=\"搜索结果的综合摘要\")\n\nclass ScanContext(BaseModel):\n \"\"\"扫描员搜集的原始数据结构\"\"\"\n hot_topics: List[str] = Field(..., description=\"当前市场热点话题\")\n news_summaries: List[Dict[str, Any]] = Field(..., description=\"关键新闻摘要列表\")\n market_data: Dict[str, Any] = Field(default_factory=dict, description=\"相关的市场行情数据\")\n sentiment_overview: str = Field(..., description=\"整体市场情绪概览\")\n raw_data_summary: str = Field(..., description=\"原始数据的综合摘要\")\n\nclass SignalCluster(BaseModel):\n theme_title: str = Field(..., description=\"主题名称\")\n signal_ids: List[int] = Field(..., description=\"包含的信号 ID 列表\")\n rationale: str = Field(..., description=\"聚类理由\")\n\nclass ClusterContext(BaseModel):\n \"\"\"信号聚类结果结构\"\"\"\n clusters: List[SignalCluster] = Field(..., description=\"聚类列表\")\n\nclass KLinePoint(BaseModel):\n date: str = Field(..., description=\"日期\")\n open: float = Field(..., description=\"开盘价\")\n high: float = Field(..., description=\"最高价\")\n low: float = Field(..., description=\"最低价\")\n close: float = Field(..., description=\"收盘价\")\n volume: float = Field(..., description=\"成交量\")\n\nclass ForecastResult(BaseModel):\n ticker: str = Field(..., description=\"股票代码\")\n base_forecast: List[KLinePoint] = Field(default_factory=list, description=\"Kronos 模型原始预测\")\n adjusted_forecast: List[KLinePoint] = Field(default_factory=list, description=\"LLM 调整后的预测\")\n rationale: str = Field(default=\"\", description=\"预测调整理由及逻辑说明\")\n timestamp: str = Field(default_factory=lambda: datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\"), description=\"生成时间\")\n\nclass InvestmentReport(BaseModel):\n overall_sentiment: str = Field(..., description=\"整体市场情绪评价\")\n market_entropy: float = Field(..., description=\"市场分歧度 (0-1, 1代表极高分歧)\")\n signals: List[InvestmentSignal] = Field(..., description=\"深度解析的投资信号列表\")\n forecasts: List[ForecastResult] = Field(default_factory=list, description=\"相关标的的预测结果\")\n timestamp: str = Field(..., description=\"报告生成时间\")\n meta_info: Optional[Dict[str, Any]] = Field(default_factory=dict, description=\"其他元数据\")\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6252,"content_sha256":"c7c1738a8dfe643db928a93dfd5e92b3276aac25a3faf1618fd38ea45c96e015"},{"filename":"scripts/utils/__init__.py","content":"# AlphaEar utils package\n","content_type":"text/x-python; charset=utf-8","language":"python","size":25,"content_sha256":"75c902eae79c725f91cb7c974d6c7711dc5a412996c67a56660e328b85321a35"},{"filename":"scripts/utils/database_manager.py","content":"import sqlite3\nimport json\nfrom datetime import datetime, date\nfrom pathlib import Path\nfrom typing import List, Dict, Optional, Any, Union\nimport pandas as pd\nfrom loguru import logger\n\nclass DatabaseManager:\n \"\"\"\n AlphaEar 数据库管理器 - 负责存储热点数据、搜索缓存和股价数据\n 使用 SQLite 进行持久化存储\n \"\"\"\n \n def __init__(self, db_path: str = \"data/signal_flux.db\"):\n self.db_path = Path(db_path)\n self.db_path.parent.mkdir(parents=True, exist_ok=True)\n self.conn = sqlite3.connect(str(self.db_path), check_same_thread=False)\n self.conn.row_factory = sqlite3.Row\n self._init_db()\n logger.info(f\"💾 Database initialized at {self.db_path}\")\n\n def _init_db(self):\n \"\"\"初始化表结构\"\"\"\n cursor = self.conn.cursor()\n \n # 1. 每日热点新闻表\n cursor.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS daily_news (\n id TEXT PRIMARY KEY,\n source TEXT,\n rank INTEGER,\n title TEXT,\n url TEXT,\n content TEXT,\n publish_time TEXT,\n crawl_time TEXT,\n sentiment_score REAL,\n analysis TEXT,\n meta_data TEXT\n )\n \"\"\")\n \n # 尝试添加 analysis 列（如果表已存在但没有该列）\n try:\n cursor.execute(\"ALTER TABLE daily_news ADD COLUMN analysis TEXT\")\n except:\n pass # 列已存在\n\n \n # 2. 搜索缓存表 (原有 JSON 缓存)\n cursor.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS search_cache (\n query_hash TEXT PRIMARY KEY,\n query TEXT,\n engine TEXT,\n results TEXT,\n timestamp TEXT\n )\n \"\"\")\n\n # 2.5 搜索详情表 (展开的搜索结果)\n cursor.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS search_detail (\n id TEXT,\n query_hash TEXT,\n rank INTEGER,\n title TEXT,\n url TEXT,\n content TEXT,\n publish_time TEXT,\n crawl_time TEXT,\n sentiment_score REAL,\n source TEXT,\n meta_data TEXT,\n PRIMARY KEY (query_hash, id)\n )\n \"\"\")\n \n # 3. 股价数据表\n cursor.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS stock_prices (\n ticker TEXT,\n date TEXT,\n open REAL,\n close REAL,\n high REAL,\n low REAL,\n volume REAL,\n change_pct REAL,\n PRIMARY KEY (ticker, date)\n )\n \"\"\")\n \n # 4. 股票列表表 (用于检索)\n cursor.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS stock_list (\n code TEXT PRIMARY KEY,\n name TEXT\n )\n \"\"\")\n \n # 5. 投资信号表 (ISQ Framework)\n cursor.execute(\"\"\"\n CREATE TABLE IF NOT EXISTS signals (\n signal_id TEXT PRIMARY KEY,\n title TEXT,\n summary TEXT,\n transmission_chain TEXT,\n sentiment_score REAL,\n confidence REAL,\n intensity INTEGER,\n expected_horizon TEXT,\n price_in_status TEXT,\n impact_tickers TEXT,\n industry_tags TEXT,\n sources TEXT,\n user_id TEXT,\n created_at TEXT\n )\n \"\"\")\n \n\n \n # 6. 创建索引以优化查询性能\n cursor.execute(\"CREATE INDEX IF NOT EXISTS idx_news_crawl_time ON daily_news(crawl_time)\")\n cursor.execute(\"CREATE INDEX IF NOT EXISTS idx_news_source ON daily_news(source)\")\n cursor.execute(\"CREATE INDEX IF NOT EXISTS idx_search_cache_timestamp ON search_cache(timestamp)\")\n cursor.execute(\"CREATE INDEX IF NOT EXISTS idx_stock_prices_ticker_date ON stock_prices(ticker, date)\")\n # 尝试添加 user_id 列到 signals 表\n try:\n cursor.execute(\"ALTER TABLE signals ADD COLUMN user_id TEXT\")\n except:\n pass\n \n cursor.execute(\"CREATE INDEX IF NOT EXISTS idx_signals_user_id ON signals(user_id)\")\n \n self.conn.commit()\n \n # \n # self.conn.commit()\n\n\n # --- 新闻数据操作 ---\n \n def save_daily_news(self, news_list: List[Dict]) -> int:\n \"\"\"保存热点新闻，包含发布时间与抓取时间\"\"\"\n cursor = self.conn.cursor()\n count = 0\n crawl_time = datetime.now().isoformat()\n \n for news in news_list:\n try:\n # 兼容不同来源的 ID 生成逻辑\n news_id = news.get('id') or f\"{news.get('source')}_{news.get('rank')}_{crawl_time[:10]}\"\n cursor.execute(\"\"\"\n INSERT OR REPLACE INTO daily_news \n (id, source, rank, title, url, content, publish_time, crawl_time, sentiment_score, meta_data)\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n \"\"\", (\n news_id,\n news.get('source'),\n news.get('rank'),\n news.get('title'),\n news.get('url'),\n news.get('content', ''),\n news.get('publish_time'), # 新增支持发布时间\n crawl_time,\n news.get('sentiment_score'),\n json.dumps(news.get('meta_data', {}))\n ))\n count += 1\n except sqlite3.Error as e:\n logger.error(f\"Database error saving news item {news.get('title')}: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error saving news item {news.get('title')}: {e}\")\n \n self.conn.commit()\n return count\n\n def get_daily_news(self, source: Optional[str] = None, limit: int = 100, days: int = 1) -> List[Dict]:\n \"\"\"获取最近 N 天的热点新闻\"\"\"\n cursor = self.conn.cursor()\n # 使用 crawl_time 过滤，保证结果的新鲜度\n time_threshold = (datetime.now().timestamp() - days * 86400)\n time_threshold_str = datetime.fromtimestamp(time_threshold).isoformat()\n \n query = \"SELECT * FROM daily_news WHERE crawl_time >= ?\"\n params = [time_threshold_str]\n \n if source:\n query += \" AND source = ?\"\n params.append(source)\n \n query += \" ORDER BY crawl_time DESC, rank LIMIT ?\"\n params.append(limit)\n \n cursor.execute(query, params)\n return [dict(row) for row in cursor.fetchall()]\n\n def lookup_reference_by_url(self, url: str) -> Optional[Dict[str, Any]]:\n \"\"\"Best-effort lookup of a source item by URL.\n\n This is used to render a stable bibliography from DB-backed metadata.\n It searches both `daily_news` and `search_detail`.\n \"\"\"\n url = (url or \"\").strip()\n if not url:\n return None\n\n cursor = self.conn.cursor()\n\n try:\n cursor.execute(\n \"\"\"\n SELECT title, source, publish_time, crawl_time, url\n FROM daily_news\n WHERE url = ?\n ORDER BY crawl_time DESC\n LIMIT 1\n \"\"\",\n (url,),\n )\n row = cursor.fetchone()\n if row:\n return dict(row)\n except Exception:\n pass\n\n try:\n cursor.execute(\n \"\"\"\n SELECT title, source, publish_time, crawl_time, url\n FROM search_detail\n WHERE url = ?\n ORDER BY crawl_time DESC\n LIMIT 1\n \"\"\",\n (url,),\n )\n row = cursor.fetchone()\n if row:\n return dict(row)\n except Exception:\n pass\n\n return None\n\n def delete_news(self, news_id: str) -> bool:\n \"\"\"删除特定新闻\"\"\"\n cursor = self.conn.cursor()\n cursor.execute(\"DELETE FROM daily_news WHERE id = ?\", (news_id,))\n self.conn.commit()\n return cursor.rowcount > 0\n \n def update_news_content(self, news_id: str, content: str = None, analysis: str = None) -> bool:\n \"\"\"更新新闻的内容或分析结果\"\"\"\n cursor = self.conn.cursor()\n updates = []\n params = []\n \n if content is not None:\n updates.append(\"content = ?\")\n params.append(content)\n if analysis is not None:\n updates.append(\"analysis = ?\")\n params.append(analysis)\n \n if not updates:\n return False\n \n params.append(news_id)\n query = f\"UPDATE daily_news SET {', '.join(updates)} WHERE id = ?\"\n cursor.execute(query, params)\n self.conn.commit()\n return cursor.rowcount > 0\n\n # --- 搜索缓存辅助 ---\n \n def get_search_cache(self, query_hash: str, ttl_seconds: Optional[int] = None) -> Optional[Dict]:\n \"\"\"获取搜索缓存 (优先查 search_detail)\"\"\"\n cursor = self.conn.cursor()\n \n # 1. 尝试从 search_detail 获取展开的结构化数据\n cursor.execute(\"\"\"\n SELECT * FROM search_detail \n WHERE query_hash = ? \n ORDER BY rank\n \"\"\", (query_hash,))\n details = [dict(row) for row in cursor.fetchall()]\n \n if details:\n # 检查 TTL (取第一条的时间)\n first_time = datetime.fromisoformat(details[0]['crawl_time'])\n if ttl_seconds and (datetime.now() - first_time).total_seconds() > ttl_seconds:\n logger.info(f\"⌛ Detailed cache expired for hash {query_hash}\")\n pass # Expired, fall through or return None? If Detail expired, Cache likely expired too.\n # But let's check basic cache just in case metadata differs? \n # Actually if details exist, we prefer them. If expired, we return None.\n return None\n \n logger.info(f\"✅ Hit detailed search cache for {query_hash} ({len(details)} items)\")\n # Reconstruct the expected 'results' list format for SearchTools\n # SearchTools expects a list of dicts. \n # We return a dict wrapper to match get_search_cache signature returning Dict usually containing 'results' string.\n # But SearchTools logic: \n # cache = db.get_search_cache(...)\n # cached_data = json.loads(cache['results'])\n \n # To minimize SearchTools changes, we can return a dict mimicking the old structure\n # OR Change SearchTools to handle list return.\n # Let's return a special dict that SearchTools can recognize or just format it as before.\n return {\"results\": json.dumps(details), \"timestamp\": details[0]['crawl_time']}\n\n # 2. Fallback to old table\n cursor.execute(\"SELECT * FROM search_cache WHERE query_hash = ?\", (query_hash,))\n row = cursor.fetchone()\n \n if not row:\n return None\n \n row_dict = dict(row)\n if ttl_seconds:\n cache_time = datetime.fromisoformat(row_dict['timestamp'])\n if (datetime.now() - cache_time).total_seconds() > ttl_seconds:\n logger.info(f\"⌛ Cache expired for hash {query_hash}\")\n return None\n \n return row_dict\n\n def save_search_cache(self, query_hash: str, query: str, engine: str, results: Union[str, List[Dict]]):\n \"\"\"保存搜索结果 (同时保存到 search_cache 和 search_detail)\"\"\"\n cursor = self.conn.cursor()\n current_time = datetime.now().isoformat()\n \n results_str = results if isinstance(results, str) else json.dumps(results)\n \n # 1. Save summary to search_cache\n cursor.execute(\"\"\"\n INSERT OR REPLACE INTO search_cache (query_hash, query, engine, results, timestamp)\n VALUES (?, ?, ?, ?, ?)\n \"\"\", (query_hash, query, engine, results_str, current_time))\n \n # 2. Save details to search_detail if results is a list\n if isinstance(results, list):\n for item in results:\n try:\n item_id = item.get('id') or f\"{hash(item.get('url', ''))}\"\n cursor.execute(\"\"\"\n INSERT OR REPLACE INTO search_detail\n (id, query_hash, rank, title, url, content, publish_time, crawl_time, sentiment_score, source, meta_data)\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n \"\"\", (\n str(item_id),\n query_hash,\n item.get('rank', 0),\n item.get('title'),\n item.get('url'),\n item.get('content', ''),\n item.get('publish_time'),\n item.get('crawl_time') or current_time,\n item.get('sentiment_score'),\n item.get('source'),\n json.dumps(item.get('meta_data', {}))\n ))\n except sqlite3.Error as e:\n logger.error(f\"Database error saving search detail {item.get('title')}: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error saving search detail {item.get('title')}: {e}\")\n \n self.conn.commit()\n\n def find_similar_queries(self, query: str, limit: int = 5) -> List[Dict]:\n \"\"\"模糊搜索相似的已缓存查询\"\"\"\n cursor = self.conn.cursor()\n \n # Simple fuzzy match: query in cached OR cached in query\n q_wild = f\"%{query}%\"\n cursor.execute(\"\"\"\n SELECT query, query_hash, timestamp, results \n FROM search_cache \n WHERE query LIKE ? OR ? LIKE ('%' || query || '%')\n ORDER BY timestamp DESC\n LIMIT ?\n \"\"\", (q_wild, query, limit))\n \n return [dict(row) for row in cursor.fetchall()]\n\n def search_local_news(self, query: str, limit: int = 5) -> List[Dict]:\n \"\"\"从本地 daily_news 搜索相关新闻\"\"\"\n cursor = self.conn.cursor()\n q_wild = f\"%{query}%\"\n # Search title and content\n cursor.execute(\"\"\"\n SELECT * FROM daily_news\n WHERE title LIKE ? OR content LIKE ?\n ORDER BY crawl_time DESC\n LIMIT ?\n \"\"\", (q_wild, q_wild, limit))\n return [dict(row) for row in cursor.fetchall()]\n\n # --- 股票数据操作 ---\n\n def save_stock_list(self, df: pd.DataFrame):\n \"\"\"保存股票列表到 stock_list 表\"\"\"\n cursor = self.conn.cursor()\n try:\n # 清空旧表\n cursor.execute(\"DELETE FROM stock_list\")\n \n # 批量插入\n data = df[['code', 'name']].to_dict('records')\n cursor.executemany(\n \"INSERT INTO stock_list (code, name) VALUES (:code, :name)\",\n data\n )\n self.conn.commit()\n except sqlite3.Error as e:\n logger.error(f\"Database error saving stock list: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error saving stock list: {e}\")\n\n def search_stock(self, query: str, limit: int = 5) -> List[Dict]:\n \"\"\"模糊搜索股票代码或名称\"\"\"\n cursor = self.conn.cursor()\n wild = f\"%{query}%\"\n cursor.execute(\"\"\"\n SELECT code, name FROM stock_list \n WHERE code LIKE ? OR name LIKE ? \n LIMIT ?\n \"\"\", (wild, wild, limit))\n return [dict(row) for row in cursor.fetchall()]\n\n def get_stock_by_code(self, code: str) -> Optional[Dict[str, str]]:\n \"\"\"精确按代码获取股票信息。\n\n Args:\n code: 股票代码（A股6位 / 港股5位），必须为纯数字字符串。\n\n Returns:\n dict: {\"code\": str, \"name\": str} 或 None。\n \"\"\"\n if not code:\n return None\n clean = \"\".join([c for c in str(code).strip() if c.isdigit()])\n if not clean:\n return None\n\n cursor = self.conn.cursor()\n cursor.execute(\"SELECT code, name FROM stock_list WHERE code = ? LIMIT 1\", (clean,))\n row = cursor.fetchone()\n return dict(row) if row else None\n\n def save_stock_prices(self, ticker: str, df: pd.DataFrame):\n \"\"\"保存股价历史数据\"\"\"\n if df.empty:\n return\n \n cursor = self.conn.cursor()\n \n # 确保 DataFrame 有必要的列\n required_cols = ['date', 'open', 'close', 'high', 'low', 'volume', 'change_pct']\n for col in required_cols:\n if col not in df.columns:\n logger.warning(f\"Missing column {col} in stock data for {ticker}\")\n return\n\n try:\n for _, row in df.iterrows():\n cursor.execute(\"\"\"\n INSERT OR REPLACE INTO stock_prices \n (ticker, date, open, close, high, low, volume, change_pct)\n VALUES (?, ?, ?, ?, ?, ?, ?, ?)\n \"\"\", (\n ticker,\n row['date'],\n row['open'],\n row['close'],\n row['high'],\n row['low'],\n row['volume'],\n row['change_pct']\n ))\n self.conn.commit()\n except sqlite3.Error as e:\n logger.error(f\"Database error saving stock prices for {ticker}: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error saving stock prices for {ticker}: {e}\")\n\n def get_stock_prices(self, ticker: str, start_date: str, end_date: str) -> pd.DataFrame:\n \"\"\"获取指定日期范围的股价数据\"\"\"\n cursor = self.conn.cursor()\n \n cursor.execute(\"\"\"\n SELECT * FROM stock_prices \n WHERE ticker = ? AND date >= ? AND date \u003c= ?\n ORDER BY date\n \"\"\", (ticker, start_date, end_date))\n \n rows = cursor.fetchall()\n if not rows:\n return pd.DataFrame()\n \n columns = ['ticker', 'date', 'open', 'close', 'high', 'low', 'volume', 'change_pct']\n return pd.DataFrame([dict(row) for row in rows], columns=columns)\n\n def execute_query(self, query: str, params: tuple = ()) -> List[Any]:\n \"\"\"执行自定义 SQL 查询\"\"\"\n try:\n cursor = self.conn.cursor()\n cursor.execute(query, params)\n if query.strip().upper().startswith(\"SELECT\"):\n return cursor.fetchall()\n else:\n self.conn.commit()\n return []\n except sqlite3.Error as e:\n logger.error(f\"SQL execution failed (Database error): {e}\")\n return []\n except Exception as e:\n logger.error(f\"SQL execution failed (Unexpected error): {e}\")\n return []\n\n # --- 投资信号操作 (ISQ Framework) ---\n\n def save_signal(self, signal: Dict[str, Any]):\n \"\"\"保存投资信号\"\"\"\n cursor = self.conn.cursor()\n created_at = datetime.now().isoformat()\n \n cursor.execute(\"\"\"\n INSERT OR REPLACE INTO signals \n (signal_id, title, summary, transmission_chain, sentiment_score, \n confidence, intensity, expected_horizon, price_in_status, \n impact_tickers, industry_tags, sources, user_id, created_at)\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)\n \"\"\", (\n signal.get('signal_id'),\n signal.get('title'),\n signal.get('summary'),\n json.dumps(signal.get('transmission_chain', [])),\n signal.get('sentiment_score', 0.0),\n signal.get('confidence', 0.0),\n signal.get('intensity', 1),\n signal.get('expected_horizon', 'T+0'),\n signal.get('price_in_status', '未知'),\n json.dumps(signal.get('impact_tickers', [])),\n json.dumps(signal.get('industry_tags', [])),\n json.dumps(signal.get('sources', [])),\n signal.get('user_id'),\n created_at\n ))\n self.conn.commit()\n\n def get_recent_signals(self, limit: int = 20, user_id: Optional[str] = None) -> List[Dict]:\n \"\"\"获取最近的投资信号\"\"\"\n cursor = self.conn.cursor()\n if user_id:\n cursor.execute(\"SELECT * FROM signals WHERE user_id = ? ORDER BY created_at DESC LIMIT ?\", (user_id, limit))\n else:\n cursor.execute(\"SELECT * FROM signals ORDER BY created_at DESC LIMIT ?\", (limit,))\n rows = cursor.fetchall()\n \n signals = []\n for row in rows:\n d = dict(row)\n # 解析 JSON 字段\n for field in ['transmission_chain', 'impact_tickers', 'industry_tags', 'sources']:\n if d.get(field):\n try:\n d[field] = json.loads(d[field])\n except:\n pass\n signals.append(d)\n return signals\n\n def close(self):\n if self.conn:\n self.conn.close()\n logger.info(\"Database connection closed.\")\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":21840,"content_sha256":"5344226bed657c9c2d6915385371004ea1abd143dce243fd727e96dbd854bf11"},{"filename":"scripts/utils/json_utils.py","content":"import ast\nimport json\nimport re\nfrom typing import Optional, Any\nfrom loguru import logger\n\ndef _strip_comments(text: str) -> str:\n \"\"\"\n Safely remove C-style comments (// and /* */) from JSON-like text,\n preserving strings (including URLs like http://).\n \"\"\"\n result = []\n i = 0\n n = len(text)\n in_string = False\n escape = False\n \n while i \u003c n:\n char = text[i]\n \n if in_string:\n if char == '\\\\':\n escape = not escape\n elif char == '\"' and not escape:\n in_string = False\n else:\n escape = False\n result.append(char)\n i += 1\n continue\n \n # Not in string\n if char == '\"':\n in_string = True\n result.append(char)\n i += 1\n continue\n \n # Check for // comment\n if i + 1 \u003c n and text[i:i+2] == '//':\n i += 2\n while i \u003c n and text[i] != '\\n':\n i += 1\n continue\n \n # Check for /* comment\n if i + 1 \u003c n and text[i:i+2] == '/*':\n i += 2\n while i + 1 \u003c n and text[i:i+2] != '*/':\n i += 1\n i += 2\n continue\n \n result.append(char)\n i += 1\n \n return ''.join(result)\n\ndef extract_json(text: str) -> Optional[Any]:\n \"\"\"\n 更加鲁棒的 JSON 提取工具。\n 处理:\n 1. Markdown 代码块 (```json ... ```)\n 2. 首尾多余字符\n 3. 同一个文本中多个 JSON 对象 (仅提取第一个)\n 4. 简单的 JSON 修复 (末尾逗号等)\n 5. C 风格注释 (// 和 /* */)\n \"\"\"\n if not text:\n return None\n \n # 1. 清理明显的 Markdown 包装\n text = text.strip()\n \n # 先尝试精确匹配 ```json ... ``` 或 ```...```\n md_match = re.search(r'```(?:json)?\\s*\\n?(.*?)\\n?```', text, re.DOTALL)\n if md_match:\n text = md_match.group(1).strip()\n elif text.startswith(\"```\"):\n # 回退：如果开头有 ``` 但没完整匹配\n text = re.sub(r'^```[a-z]*\\n?', '', text)\n text = re.sub(r'\\n?```\\s*

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.

, '', text)\n \n # 2. 寻找第一个 JSON 起始符 { 或 [\n start_brace = text.find('{')\n start_bracket = text.find('[')\n \n if start_brace == -1 and start_bracket == -1:\n return None\n \n start_idx = start_brace if (start_bracket == -1 or (start_brace != -1 and start_brace \u003c start_bracket)) else start_bracket\n \n # 2.5 预处理：修复一些极其常见的 LLM 错误\n potential_json = text[start_idx:].strip()\n \n # remove comments safely\n potential_json = _strip_comments(potential_json)\n \n # b. 修复缺失开头引号的键: nodes\": [ -> \"nodes\": [\n # 匹配模式: (空白或换行) 单词紧跟引号和冒号\n potential_json = re.sub(r'([\\{\\,]\\s*)([a-zA-Z_]\\w*)\\\"\\s*:', r'\\1\"\\2\":', potential_json)\n \n # c. 修复缺失末尾引号的键: \"nodes: [ -> \"nodes\": [\n potential_json = re.sub(r'([\\{\\,]\\s*)\\\"([a-zA-Z_]\\w*)\\s*:', r'\\1\"\\2\":', potential_json)\n\n # d. 修复完全缺失引号的键: nodes: [ -> \"nodes\": [\n # 注意避免匹配到像 http:// 这种内容，所以限定在 { 或 , 之后\n potential_json = re.sub(r'([\\{\\,]\\s*)([a-zA-Z_]\\w*)\\s*:', r'\\1\"\\2\":', potential_json)\n \n # 3. 使用 raw_decode 尝试解析\n decoder = json.JSONDecoder()\n \n # 首先尝试直接解析（不做任何预处理）\n try:\n obj = json.loads(potential_json)\n return obj\n except json.JSONDecodeError:\n pass\n \n # 简单预处理：移除对象/列表末位多余逗号\n processed_json = re.sub(r',\\s*([\\]}])', r'\\1', potential_json)\n \n try:\n obj, end_pos = decoder.raw_decode(processed_json)\n return obj\n except json.JSONDecodeError:\n pass\n \n # e. 修复未终止的字符串字面量问题：移除值中的实际换行符\n # LLM 可能在字符串值中生成包含真实 newline 的内容，导致 JSON 非法\n def fix_multiline_strings(s):\n # 简单策略：将字符串值内的换行替换为空格\n lines = s.split('\\n')\n result = []\n in_string = False\n for line in lines:\n # 计算未转义的引号数\n quote_count = line.count('\"') - line.count('\\\\\"')\n if in_string:\n result[-1] += ' ' + line.strip()\n else:\n result.append(line)\n \n if quote_count % 2 == 1:\n in_string = not in_string\n return '\\n'.join(result)\n \n fixed_json = fix_multiline_strings(processed_json)\n \n try:\n obj, end_pos = decoder.raw_decode(fixed_json)\n return obj\n except json.JSONDecodeError:\n try:\n # 4. 尝试处理单引号问题 (JSON 规范要求双引号，但 LLM 常输出单引号)\n # 这是一个简单的替换技巧，仅针对像 {'key': 'value'} 这样的结构\n # 注意：这可能会破坏包含单引号的字符串值，所以作为较后的回退\n fix_quotes = re.sub(r\"'(.*?)':\", r'\"\\1\":', processed_json) # 修复键\n fix_quotes = re.sub(r\":\\s*'(.*?)'\", r': \"\\1\"', fix_quotes) # 修复简单值\n obj, end_pos = decoder.raw_decode(fix_quotes)\n return obj\n except (json.JSONDecodeError, TypeError):\n try:\n # 5. 使用 ast.literal_eval 作为终极回退 (处理 Python 字典格式)\n # 提取第一个匹配的括号对内容\n # 寻找匹配的 { }\n stack = []\n for i, char in enumerate(potential_json):\n if char == '{': stack.append('{')\n elif char == '}':\n if stack: stack.pop()\n if not stack:\n content = potential_json[:i+1]\n return ast.literal_eval(content)\n except (ValueError, SyntaxError, MemoryError) as e:\n logger.warning(f\"All JSON extraction attempts failed: {e}\")\n except Exception as e:\n logger.error(f\"Unexpected error during JSON extraction: {e}\")\n \n return None\n","content_type":"text/x-python; charset=utf-8","language":"python","size":6300,"content_sha256":"2ec52de0a706f64b3466f490dc3f12ad8241f20fa316e858f7e3393afb3062d9"},{"filename":"scripts/utils/llm/capability.py","content":"import os\nfrom typing import Optional, List, Dict, Any\nfrom agno.agent import Agent\nfrom agno.models.base import Model\nfrom loguru import logger\nfrom ..llm.factory import get_model\n\ndef test_tool_call_support(model: Model) -> bool:\n \"\"\"\n 测试模型是否支持原生的 Tool Call (Function Calling)。\n 通过尝试执行一个简单的加法工具来验证。\n \"\"\"\n def get_current_weather(location: str):\n \"\"\"获取指定地点的天气\"\"\"\n return f\"{location} 的天气是晴天，25度。\"\n\n test_agent = Agent(\n model=model,\n tools=[get_current_weather],\n instructions=\"请调用工具查询北京的天气，并直接返回工具的输出结果。\"\n )\n\n try:\n # 运行一个简单的任务，观察是否触发了 tool_call\n response = test_agent.run(\"北京天气怎么样？\")\n \n # 检查 response 中是否包含 tool_calls\n # Agno 的 RunResponse 对象通常包含 messages，我们可以检查最后几条消息\n has_tool_call = False\n for msg in response.messages:\n if hasattr(msg, 'tool_calls') and msg.tool_calls:\n has_tool_call = True\n break\n \n if has_tool_call:\n logger.info(f\"✅ Model {model.id} supports native tool calling.\")\n return True\n else:\n # 如果没有 tool_calls 但返回了正确答案，可能是模型通过纯文本模拟了工具调用（ReAct）\n # 或者根本没用工具。对于原生支持的判断，我们坚持要求有 tool_calls 结构。\n logger.warning(f\"⚠️ Model {model.id} did NOT use native tool calling structure.\")\n return False\n \n except Exception as e:\n logger.error(f\"❌ Error testing tool call for {model.id}: {e}\")\n return False\n\nclass ModelCapabilityRegistry:\n \"\"\"\n 模型能力注册表，用于缓存和管理不同模型的能力测试结果。\n \"\"\"\n _cache = {}\n\n @classmethod\n def get_capabilities(cls, provider: str, model_id: str, **kwargs) -> Dict[str, bool]:\n key = f\"{provider}:{model_id}\"\n if key not in cls._cache:\n logger.info(f\"🔍 Testing capabilities for {key}...\")\n model = get_model(provider, model_id, **kwargs)\n supports_tool_call = test_tool_call_support(model)\n cls._cache[key] = {\n \"supports_tool_call\": supports_tool_call\n }\n return cls._cache[key]\n\nif __name__ == \"__main__\":\n # 简单测试脚本\n from dotenv import load_dotenv\n load_dotenv()\n \n # 测试当前配置的模型\n p = os.getenv(\"LLM_PROVIDER\", \"ust\")\n m = os.getenv(\"LLM_MODEL\", \"Qwen\")\n \n print(f\"Testing {p}/{m}...\")\n res = ModelCapabilityRegistry.get_capabilities(p, m)\n print(f\"Result: {res}\")\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2863,"content_sha256":"40b215695e2b7519224773ec4b280d115770670d1ddfc37e3eea2d23e085e171"},{"filename":"scripts/utils/llm/factory.py","content":"import os\nfrom agno.models.openai import OpenAIChat\nfrom agno.models.ollama import Ollama\nfrom agno.models.dashscope import DashScope\nfrom agno.models.deepseek import DeepSeek\nfrom agno.models.openrouter import OpenRouter\n\ndef get_model(model_provider: str, model_id: str, **kwargs):\n \"\"\"\n Factory to get the appropriate LLM model.\n \n Args:\n model_provider: \"openai\", \"ollama\", \"deepseek\"\n model_id: The specific model ID (e.g., \"gpt-4o\", \"llama3\", \"deepseek-chat\")\n **kwargs: Additional arguments for the model constructor\n \"\"\"\n if model_provider == \"openai\":\n return OpenAIChat(id=model_id, **kwargs)\n \n elif model_provider == \"ollama\":\n return Ollama(id=model_id, **kwargs)\n \n elif model_provider == \"deepseek\":\n # DeepSeek is OpenAI compatible\n api_key = os.getenv(\"DEEPSEEK_API_KEY\")\n if not api_key:\n print(\"Warning: DEEPSEEK_API_KEY not set.\")\n \n return DeepSeek(\n id=model_id,\n api_key=api_key,\n **kwargs\n )\n elif model_provider == \"dashscope\":\n api_key = os.getenv(\"DASHSCOPE_API_KEY\")\n if not api_key:\n print(\"Warning: DASHSCOPE_API_KEY not set.\")\n \n return DashScope(\n id=model_id,\n base_url=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n api_key=api_key,\n **kwargs\n )\n elif model_provider == 'openrouter':\n api_key = os.getenv(\"OPENROUTER_API_KEY\")\n if not api_key:\n print('Warning: OPENROUTER_API_KEY not set.')\n \n return OpenRouter(\n id=model_id,\n api_key=api_key,\n **kwargs\n )\n\n elif model_provider == 'zai':\n api_key = os.getenv(\"ZAI_KEY_API\")\n if not api_key:\n print('Warning: ZAI_KEY_API not set.')\n\n # role_map to ensure compatibility.\n default_role_map = {\n \"system\": \"system\",\n \"user\": \"user\",\n \"assistant\": \"assistant\",\n \"tool\": \"tool\",\n \"model\": \"assistant\",\n }\n\n # Allow callers to override role_map via kwargs, otherwise use default\n role_map = kwargs.pop(\"role_map\", default_role_map)\n \n return OpenAIChat(\n id=model_id,\n base_url=\"https://api.z.ai/api/paas/v4\",\n api_key=api_key,\n timeout=60,\n role_map=role_map,\n extra_body={\"enable_thinking\": False}, # TODO: one more setting for thinking\n **kwargs\n )\n \n elif model_provider == 'ust':\n api_key = os.getenv(\"UST_KEY_API\")\n if not api_key:\n print('Warning: UST_KEY_API not set.')\n \n # Some UST-compatible endpoints expect the standard OpenAI role names\n # (e.g. \"system\", \"user\", \"assistant\") rather than Agno's default\n # mapping which maps \"system\" -> \"developer\". Provide an explicit\n # role_map to ensure compatibility.\n default_role_map = {\n \"system\": \"system\",\n \"user\": \"user\",\n \"assistant\": \"assistant\",\n \"tool\": \"tool\",\n \"model\": \"assistant\",\n }\n\n # Allow callers to override role_map via kwargs, otherwise use default\n role_map = kwargs.pop(\"role_map\", default_role_map)\n\n return OpenAIChat(\n id=model_id,\n api_key=api_key,\n base_url=os.getenv(\"UST_URL\"),\n role_map=role_map,\n extra_body={\"enable_thinking\": False}, # TODO: one more setting for thinking\n **kwargs\n )\n \n else:\n raise ValueError(f\"Unknown model provider: {model_provider}\")\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":3717,"content_sha256":"85065ea8a843075f86a2532d455e64dca2207b657b758196d89cff9fad9e67b0"},{"filename":"scripts/utils/llm/router.py","content":"import os\nfrom typing import Optional, List, Dict, Any, Union\nfrom agno.models.base import Model\nfrom loguru import logger\nfrom dotenv import load_dotenv\nfrom ..llm.factory import get_model\nfrom ..llm.capability import ModelCapabilityRegistry\n\n# 确保在初始化前加载环境变量\nload_dotenv()\n\nclass ModelRouter:\n \"\"\"\n 模型路由管理器\n \n 功能：\n 1. 管理“推理/写作模型” (Reasoning Model) 和“工具调用模型” (Tool Model)。\n 2. 根据任务需求自动选择合适的模型。\n \"\"\"\n \n def __init__(self):\n # 默认从环境变量读取\n self.reasoning_provider = os.getenv(\"REASONING_MODEL_PROVIDER\", os.getenv(\"LLM_PROVIDER\", \"openai\"))\n self.reasoning_id = os.getenv(\"REASONING_MODEL_ID\", os.getenv(\"LLM_MODEL\", \"gpt-4o\"))\n self.reasoning_host = os.getenv(\"REASONING_MODEL_HOST\", os.getenv(\"LLM_HOST\"))\n \n self.tool_provider = os.getenv(\"TOOL_MODEL_PROVIDER\", self.reasoning_provider)\n self.tool_id = os.getenv(\"TOOL_MODEL_ID\", self.reasoning_id)\n self.tool_host = os.getenv(\"TOOL_MODEL_HOST\", self.reasoning_host)\n \n self._reasoning_model = None\n self._tool_model = None\n \n logger.info(f\"🤖 ModelRouter initialized: Reasoning={self.reasoning_id} ({self.reasoning_host or 'default'}), Tool={self.tool_id} ({self.tool_host or 'default'})\")\n\n def get_reasoning_model(self, **kwargs) -> Model:\n if not self._reasoning_model:\n # 优先使用路由配置的 host\n if self.reasoning_host and \"host\" not in kwargs:\n kwargs[\"host\"] = self.reasoning_host\n self._reasoning_model = get_model(self.reasoning_provider, self.reasoning_id, **kwargs)\n return self._reasoning_model\n\n def get_tool_model(self, **kwargs) -> Model:\n if not self._tool_model:\n # 优先使用路由配置的 host\n if self.tool_host and \"host\" not in kwargs:\n kwargs[\"host\"] = self.tool_host\n \n # 检查 tool_model 是否真的支持 tool call\n caps = ModelCapabilityRegistry.get_capabilities(self.tool_provider, self.tool_id, **kwargs)\n if not caps[\"supports_tool_call\"]:\n logger.warning(f\"⚠️ Configured tool model {self.tool_id} might not support native tool calls! Consider using ReAct mode or a different model.\")\n \n self._tool_model = get_model(self.tool_provider, self.tool_id, **kwargs)\n return self._tool_model\n\n def get_model_for_agent(self, has_tools: bool = False, **kwargs) -> Model:\n \"\"\"\n 根据 Agent 是否包含工具来返回合适的模型。\n \"\"\"\n if has_tools:\n return self.get_tool_model(**kwargs)\n return self.get_reasoning_model(**kwargs)\n\n# 全局单例\nrouter = ModelRouter()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":2875,"content_sha256":"0a8dbfa1a9d8d2a5b62ca77bc5a78bbbb1e53706be9cf18a7290e7608649927a"},{"filename":"scripts/utils/logging_setup.py","content":"import os\nimport sys\nfrom datetime import datetime\nfrom typing import Optional\n\nfrom loguru import logger\n\n\ndef setup_file_logging(\n run_id: str,\n log_dir: str = \"logs\",\n level: str = \"INFO\",\n retention: str = \"10 days\",\n rotation: str = \"20 MB\",\n) -> str:\n \"\"\"Configure Loguru to log to stderr + a per-run file.\n\n Returns the log file path.\n \"\"\"\n os.makedirs(log_dir, exist_ok=True)\n\n # Remove default handler to avoid duplicate logs.\n logger.remove()\n\n # Console\n logger.add(sys.stderr, level=level, backtrace=False, diagnose=False)\n\n # File (safe for multi-thread via enqueue)\n log_path = os.path.join(log_dir, f\"signalflux_{run_id}.log\")\n logger.add(\n log_path,\n level=level,\n rotation=rotation,\n retention=retention,\n enqueue=True,\n backtrace=True,\n diagnose=False,\n encoding=\"utf-8\",\n )\n return log_path\n\n\ndef make_run_id(prefix: Optional[str] = None) -> str:\n ts = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n return f\"{prefix}_{ts}\" if prefix else ts\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1072,"content_sha256":"2b0e1a87da39920ed0958e5f9915cf481dccc8016874c91acaf2a0403916eadc"},{"filename":"scripts/utils/predictor/evaluation.py","content":"import os\nimport sys\nimport torch\nimport pandas as pd\nimport numpy as np\nimport glob\nfrom loguru import logger\nfrom datetime import datetime, timedelta\n\n# Setup paths\nKRONOS_DIR = os.path.dirname(os.path.abspath(__file__))\nSRC_DIR = os.path.dirname(os.path.dirname(KRONOS_DIR))\nif SRC_DIR not in sys.path:\n sys.path.insert(0, SRC_DIR)\n\nfrom ..kronos.auto_synthesis_training import AutoSynthesisTrainer\nfrom ..kronos.model import KronosPredictor\nfrom ..visualizer import VisualizerTools\nfrom ..schema.models import ForecastResult, KLinePoint\n\nclass NewsModelEvaluator:\n def __init__(self, model_path=None):\n self.trainer = AutoSynthesisTrainer()\n self.device = self.trainer.device\n \n if model_path is None:\n # Try to find the latest model in exports/models\n model_files = glob.glob(os.path.join(SRC_DIR, \"exports/models/*.pt\"))\n if not model_files:\n logger.warning(\"⚠️ No trained models found in exports/models/. Using base model (zero-init proj).\")\n else:\n model_path = max(model_files, key=os.path.getctime)\n \n if model_path:\n self.load_weights(model_path)\n\n def load_weights(self, path):\n logger.info(f\"🔄 Loading model weights from {path}...\")\n checkpoint = torch.load(path, map_location=self.device)\n self.trainer.model.news_proj.load_state_dict(checkpoint['news_proj_state_dict'])\n logger.success(\"✅ News projection layer loaded.\")\n\n def evaluate_range(self, start_idx=100, end_idx=200, pred_len=5):\n # 1. Fetch Tickers\n res = self.trainer.db.execute_query(\"SELECT code FROM stock_list\")\n all_tickers = [row['code'] for row in res]\n test_tickers = all_tickers[start_idx:end_idx]\n \n if not test_tickers:\n logger.error(f\"No tickers found in range {start_idx}-{end_idx}\")\n return\n\n logger.info(f\"🚀 Evaluating News Model on stocks {start_idx} to {end_idx}...\")\n \n # 2. Discover Shocks\n shocks = self.trainer.discover_shocks(test_tickers, pred_len=pred_len)\n \n # 3. Associate News & Predict\n self.trainer.model.eval()\n predictor = KronosPredictor(self.trainer.model, self.trainer.tokenizer, device=self.device)\n \n save_dir = os.path.join(SRC_DIR, \"exports/evaluation_results\")\n os.makedirs(save_dir, exist_ok=True)\n\n count = 0\n for shock in shocks:\n summary = self.trainer.find_reason_and_verify(shock)\n if not summary:\n continue\n \n logger.info(f\"📈 Testing shock: {shock['ticker']} on {shock['date']}\")\n \n # Embedding news\n news_emb = self.trainer.embedder.encode(summary)\n \n # Prediction\n h = shock['history']\n t = shock['target']\n actuals = t['close'].values[:pred_len]\n \n x_ts = pd.to_datetime(h['date'])\n future_dates = pd.date_range(start=x_ts.iloc[-1] + timedelta(days=1), periods=pred_len, freq='B')\n y_ts = pd.Series(future_dates)\n \n # A. Base Prediction (No news)\n p_base = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=None, verbose=False)\n \n # B. News-Aware Prediction\n p_news = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=news_emb, verbose=False)\n \n # Calculate Improvement\n b_preds = p_base['close'].values[:len(actuals)]\n n_preds = p_news['close'].values[:len(actuals)]\n b_mae = np.mean(np.abs(b_preds - actuals))\n n_mae = np.mean(np.abs(n_preds - actuals))\n improvement = (b_mae - n_mae) / (b_mae + 1e-6) * 100\n\n # C. Visualize\n try:\n def to_kp_list(preds_df):\n points = []\n for idx, row in preds_df.iterrows():\n points.append(KLinePoint(\n date=str(idx)[:10], open=row['open'], high=row['high'],\n low=row['low'], close=row['close'], volume=row.get('volume', 0)\n ))\n return points\n\n forecast_obj = ForecastResult(\n ticker=shock['ticker'],\n base_forecast=to_kp_list(p_base),\n adjusted_forecast=to_kp_list(p_news),\n rationale=summary\n )\n\n chart = VisualizerTools.generate_stock_chart(\n df=h, ticker=shock['ticker'],\n title=f\"Test Eval: {shock['ticker']} ({shock['date']}) Imp: {improvement:.1f}%\",\n forecast=forecast_obj,\n ground_truth=t[['date', 'open', 'high', 'low', 'close', 'volume']]\n )\n \n safe_date = shock['date'].replace(\"-\", \"\")\n filename = f\"test_{shock['ticker']}_{safe_date}.html\"\n VisualizerTools.render_chart_to_file(chart, os.path.join(save_dir, filename))\n \n logger.success(f\"📊 Result for {shock['ticker']} saved. Base MAE: {b_mae:.4f}, News MAE: {n_mae:.4f}\")\n count += 1\n except Exception as e:\n logger.error(f\"Visualization failed: {e}\")\n\n logger.info(f\"🏁 Finished evaluation. {count} cases visualized in {save_dir}\")\n\nif __name__ == \"__main__\":\n # If you have a specific model, pass the path here. Otherwise it picks the latest.\n evaluator = NewsModelEvaluator()\n evaluator.evaluate_range(start_idx=100, end_idx=200, pred_len=1)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":5723,"content_sha256":"c00bbac8132e737bf81917ad79dbc1dfe902f66c9d8b1813649d4bad6cbec78d"},{"filename":"scripts/utils/predictor/kline_generate.py","content":"# Ref: https://github.com/shiyu-coder/Kronos\n\nfrom model import Kronos, KronosTokenizer, KronosPredictor\nimport pandas as pd\nimport sqlite3\nimport torch\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nfrom pandas.tseries.offsets import BusinessDay\nimport numpy as np\n\ndef get_device():\n device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n print(f\"Using device: {device}\")\n return device\n\ndef load_predictor():\n tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\")\n model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\")\n device = get_device()\n tokenizer = tokenizer.to(device)\n model = model.to(device)\n return KronosPredictor(model, tokenizer, device=device, max_context=512)\n\ndef load_data(ticker=\"002111\", db_path=\"AlphaEar/data/signal_flux.db\"):\n with sqlite3.connect(db_path) as conn:\n df = pd.read_sql_query(f\"SELECT * FROM stock_prices WHERE ticker = '{ticker}'\", conn)\n df['date'] = pd.to_datetime(df['date'])\n df = df.sort_values('date').reset_index(drop=True)\n return df\n\ndef plot_kline_matplotlib(ax, ax_vol, dates, df, label_suffix=\"\", color_up='#ef4444', color_down='#22c55e', alpha=1.0, is_prediction=False):\n \"\"\"\n 绘制 K 线图和成交量\n \"\"\"\n # X axis mapping to integers for consistent spacing\n x = np.arange(len(dates))\n \n # K-line data\n opens = df['open'].values\n closes = df['close'].values\n highs = df['high'].values\n lows = df['low'].values\n volumes = df['volume'].values\n \n # Width of the candlestick\n width = 0.6\n \n for i in range(len(x)):\n color = color_up if closes[i] >= opens[i] else color_down\n linestyle = '--' if is_prediction else '-'\n \n # Wick\n ax.vlines(x[i], lows[i], highs[i], color=color, linewidth=1, alpha=alpha, linestyle=linestyle)\n \n # Body\n rect_bottom = min(opens[i], closes[i])\n rect_height = abs(opens[i] - closes[i])\n if rect_height == 0: rect_height = 0.001 # Visual hair\n \n ax.add_patch(plt.Rectangle((x[i] - width/2, rect_bottom), width, rect_height, \n edgecolor=color, facecolor=color if not is_prediction else 'none', \n alpha=alpha, linewidth=1, linestyle=linestyle))\n \n # Volume\n ax_vol.bar(x[i], volumes[i], color=color, alpha=alpha * 0.5, width=width)\n\ndef render_comparison_chart(history_df, actual_df, pred_df, title):\n \"\"\"\n 渲染组合图：历史 K 线 + 真值 K 线 + 预测 K 线\n \"\"\"\n # Combine all dates for X axis\n all_dates = pd.concat([history_df['date'], actual_df['date'] if actual_df is not None else pred_df.index.to_series()]).unique()\n all_dates = sorted(all_dates)\n date_to_idx = {date: i for i, date in enumerate(all_dates)}\n \n fig = plt.figure(figsize=(14, 8), facecolor='white')\n gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1], hspace=0.1)\n ax_main = fig.add_subplot(gs[0])\n ax_vol = fig.add_subplot(gs[1], sharex=ax_main)\n \n # 1. Plot History\n hist_indices = [date_to_idx[d] for d in history_df['date']]\n # We use a custom x for plotting to ensure continuity\n plot_kline_matplotlib(ax_main, ax_vol, history_df['date'], history_df, alpha=0.8)\n \n offset = len(history_df)\n \n # 2. Plot Actual if exists\n if actual_df is not None:\n # Shift indices\n actual_x = np.arange(len(actual_df)) + offset\n # Plotting manually to handle offset\n for i in range(len(actual_df)):\n idx = actual_x[i]\n row = actual_df.iloc[i]\n color = '#ef4444' if row['close'] >= row['open'] else '#22c55e'\n ax_main.vlines(idx, row['low'], row['high'], color=color, linewidth=1, alpha=0.9)\n ax_main.add_patch(plt.Rectangle((idx - 0.3, min(row['open'], row['close'])), 0.6, abs(row['open']-row['close']), \n edgecolor=color, facecolor=color, alpha=0.9))\n ax_vol.bar(idx, row['volume'], color=color, alpha=0.4)\n \n # 3. Plot Prediction\n pred_x = np.arange(len(pred_df)) + offset\n for i in range(len(pred_df)):\n idx = pred_x[i]\n row = pred_df.iloc[i]\n color = '#ff8c00' # Orange for prediction to distinguish\n ax_main.vlines(idx, row['low'], row['high'], color=color, linewidth=1.5, linestyle='--')\n ax_main.add_patch(plt.Rectangle((idx - 0.3, min(row['open'], row['close'])), 0.6, abs(row['open']-row['close']), \n edgecolor=color, facecolor='none', linewidth=1.5, linestyle='--'))\n # Plot secondary prediction line for close\n if i == 0:\n # Connect to history\n ax_main.plot([offset-1, idx], [history_df['close'].iloc[-1], row['close']], color=color, linestyle='--', alpha=0.6)\n elif i > 0:\n ax_main.plot([idx-1, idx], [pred_df['close'].iloc[i-1], row['close']], color=color, linestyle='--', alpha=0.6)\n\n # Styling\n ax_main.set_title(title, fontsize=14, fontweight='bold')\n ax_main.grid(True, linestyle=':', alpha=0.6)\n ax_vol.grid(True, linestyle=':', alpha=0.6)\n ax_vol.set_ylabel('Volume')\n ax_main.set_ylabel('Price')\n \n # Set X ticks\n step = max(1, len(all_dates) // 10)\n ax_vol.set_xticks(np.arange(0, len(all_dates), step))\n ax_vol.set_xticklabels([all_dates[i].strftime('%Y-%m-%d') for i in range(0, len(all_dates), step)], rotation=45)\n \n plt.tight_layout()\n plt.show()\n plt.close()\n\ndef run_backtest(df, predictor, lookback, pred_len, start_index=0):\n total_len = len(df)\n history_start = start_index\n history_end = start_index + lookback \n pred_start = history_end\n \n available_pred_len = total_len - pred_start\n if available_pred_len \u003c= 0: return\n actual_pred_len = min(pred_len, available_pred_len)\n pred_end = pred_start + actual_pred_len\n \n x_df = df.iloc[history_start : history_end].copy()\n y_true_df = df.iloc[pred_start : pred_end].copy()\n y_timestamp = y_true_df['date']\n \n print(f\"Backtesting: {x_df['date'].iloc[0].date()} to {y_timestamp.iloc[-1].date()}\")\n \n pred_df = predictor.predict(\n df=x_df[['open', 'high', 'low', 'close', 'volume']],\n x_timestamp=x_df['date'],\n y_timestamp=y_timestamp,\n pred_len=actual_pred_len,\n T=1.0, top_p=0.9, sample_count=1\n )\n \n render_comparison_chart(x_df, y_true_df, pred_df, f\"Backtest: {TICKER} K-Line Comparison\")\n\ndef run_forecast(df, predictor, lookback, pred_len):\n if len(df) \u003c lookback: return\n x_df = df.iloc[-lookback:].copy()\n last_date = x_df['date'].iloc[-1]\n future_dates = pd.date_range(start=last_date + BusinessDay(1), periods=pred_len, freq='B')\n future_dates = pd.Series(future_dates)\n \n print(f\"Forecasting: Starting from {future_dates.iloc[0].date()}\")\n \n pred_df = predictor.predict(\n df=x_df[['open', 'high', 'low', 'close', 'volume']],\n x_timestamp=x_df['date'],\n y_timestamp=future_dates,\n pred_len=pred_len,\n T=1.0, top_p=0.9, sample_count=1\n )\n \n render_comparison_chart(x_df, None, pred_df, f\"Forecast: {TICKER} Future K-Line\")\n\nif __name__ == \"__main__\":\n LOOKBACK = 20\n PRED_LEN = 10\n TICKER = '002111'\n \n pred_model = load_predictor()\n stock_data = load_data(TICKER)\n \n total_rows = len(stock_data)\n backtest_start = max(0, total_rows - LOOKBACK - PRED_LEN - 10) # Leave some space to see trend\n \n print(\"\\n--- Running Backtest ---\")\n run_backtest(stock_data, pred_model, LOOKBACK, PRED_LEN, start_index=backtest_start)\n \n print(\"\\n--- Running Forecast ---\")\n run_forecast(stock_data, pred_model, LOOKBACK, PRED_LEN)","content_type":"text/x-python; charset=utf-8","language":"python","size":7821,"content_sha256":"83139d1249d2d797969be1207d0dab9f59f9d2563f28799da83b3d0918b39952"},{"filename":"scripts/utils/predictor/model/__init__.py","content":"from .kronos import KronosTokenizer, Kronos, KronosPredictor\n\nmodel_dict = {\n 'kronos_tokenizer': KronosTokenizer,\n 'kronos': Kronos,\n 'kronos_predictor': KronosPredictor\n}\n\n\ndef get_model_class(model_name):\n if model_name in model_dict:\n return model_dict[model_name]\n else:\n print(f\"Model {model_name} not found in model_dict\")\n raise NotImplementedError\n\n","content_type":"text/x-python; charset=utf-8","language":"python","size":394,"content_sha256":"14193881034a621332145ee900e18cda79f72f65b2a2dc98059f8da5305ed336"},{"filename":"scripts/utils/predictor/model/kronos.py","content":"import numpy as np\nimport pandas as pd\nimport torch\nfrom huggingface_hub import PyTorchModelHubMixin\nimport sys\n\nfrom tqdm import trange\n\nsys.path.append(\"../\")\nfrom model.module import *\n\n\nclass KronosTokenizer(nn.Module, PyTorchModelHubMixin):\n \"\"\"\n KronosTokenizer module for tokenizing input data using a hybrid quantization approach.\n\n This tokenizer utilizes a combination of encoder and decoder Transformer blocks\n along with the Binary Spherical Quantization (BSQuantizer) to compress and decompress input data.\n\n Args:\n d_in (int): Input dimension.\n d_model (int): Model dimension.\n n_heads (int): Number of attention heads.\n ff_dim (int): Feed-forward dimension.\n n_enc_layers (int): Number of encoder layers.\n n_dec_layers (int): Number of decoder layers.\n ffn_dropout_p (float): Dropout probability for feed-forward networks.\n attn_dropout_p (float): Dropout probability for attention mechanisms.\n resid_dropout_p (float): Dropout probability for residual connections.\n s1_bits (int): Number of bits for the pre token in BSQuantizer.\n s2_bits (int): Number of bits for the post token in BSQuantizer.\n beta (float): Beta parameter for BSQuantizer.\n gamma0 (float): Gamma0 parameter for BSQuantizer.\n gamma (float): Gamma parameter for BSQuantizer.\n zeta (float): Zeta parameter for BSQuantizer.\n group_size (int): Group size parameter for BSQuantizer.\n\n \"\"\"\n\n def __init__(self, d_in, d_model, n_heads, ff_dim, n_enc_layers, n_dec_layers, ffn_dropout_p, attn_dropout_p, resid_dropout_p, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):\n\n super().__init__()\n self.d_in = d_in\n self.d_model = d_model\n self.n_heads = n_heads\n self.ff_dim = ff_dim\n self.enc_layers = n_enc_layers\n self.dec_layers = n_dec_layers\n self.ffn_dropout_p = ffn_dropout_p\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout_p = resid_dropout_p\n\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.codebook_dim = s1_bits + s2_bits # Total dimension of the codebook after quantization\n self.embed = nn.Linear(self.d_in, self.d_model)\n self.head = nn.Linear(self.d_model, self.d_in)\n\n # Encoder Transformer Blocks\n self.encoder = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.enc_layers - 1)\n ])\n # Decoder Transformer Blocks\n self.decoder = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.dec_layers - 1)\n ])\n self.quant_embed = nn.Linear(in_features=self.d_model, out_features=self.codebook_dim) # Linear layer before quantization\n self.post_quant_embed_pre = nn.Linear(in_features=self.s1_bits, out_features=self.d_model) # Linear layer after quantization (pre part - s1 bits)\n self.post_quant_embed = nn.Linear(in_features=self.codebook_dim, out_features=self.d_model) # Linear layer after quantization (full codebook)\n self.tokenizer = BSQuantizer(self.s1_bits, self.s2_bits, beta, gamma0, gamma, zeta, group_size) # BSQuantizer module\n\n def forward(self, x):\n \"\"\"\n Forward pass of the KronosTokenizer.\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).\n\n Returns:\n tuple: A tuple containing:\n - tuple: (z_pre, z) - Reconstructed outputs from decoder with s1_bits and full codebook respectively,\n both of shape (batch_size, seq_len, d_in).\n - torch.Tensor: bsq_loss - Loss from the BSQuantizer.\n - torch.Tensor: quantized - Quantized representation from BSQuantizer.\n - torch.Tensor: z_indices - Indices from the BSQuantizer.\n \"\"\"\n z = self.embed(x)\n\n for layer in self.encoder:\n z = layer(z)\n\n z = self.quant_embed(z) # (B, T, codebook)\n\n bsq_loss, quantized, z_indices = self.tokenizer(z)\n\n quantized_pre = quantized[:, :, :self.s1_bits] # Extract the first part of quantized representation (s1_bits)\n z_pre = self.post_quant_embed_pre(quantized_pre)\n\n z = self.post_quant_embed(quantized)\n\n # Decoder layers (for pre part - s1 bits)\n for layer in self.decoder:\n z_pre = layer(z_pre)\n z_pre = self.head(z_pre)\n\n # Decoder layers (for full codebook)\n for layer in self.decoder:\n z = layer(z)\n z = self.head(z)\n\n return (z_pre, z), bsq_loss, quantized, z_indices\n\n def indices_to_bits(self, x, half=False):\n \"\"\"\n Converts indices to bit representations and scales them.\n\n Args:\n x (torch.Tensor): Indices tensor.\n half (bool, optional): Whether to process only half of the codebook dimension. Defaults to False.\n\n Returns:\n torch.Tensor: Bit representation tensor.\n \"\"\"\n if half:\n x1 = x[0] # Assuming x is a tuple of indices if half is True\n x2 = x[1]\n mask = 2 ** torch.arange(self.codebook_dim//2, device=x1.device, dtype=torch.long) # Create a mask for bit extraction\n x1 = (x1.unsqueeze(-1) & mask) != 0 # Extract bits for the first half\n x2 = (x2.unsqueeze(-1) & mask) != 0 # Extract bits for the second half\n x = torch.cat([x1, x2], dim=-1) # Concatenate the bit representations\n else:\n mask = 2 ** torch.arange(self.codebook_dim, device=x.device, dtype=torch.long) # Create a mask for bit extraction\n x = (x.unsqueeze(-1) & mask) != 0 # Extract bits\n\n x = x.float() * 2 - 1 # Convert boolean to bipolar (-1, 1)\n q_scale = 1. / (self.codebook_dim ** 0.5) # Scaling factor\n x = x * q_scale\n return x\n\n def encode(self, x, half=False):\n \"\"\"\n Encodes the input data into quantized indices.\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).\n half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.\n\n Returns:\n torch.Tensor: Quantized indices from BSQuantizer.\n \"\"\"\n z = self.embed(x)\n for layer in self.encoder:\n z = layer(z)\n z = self.quant_embed(z)\n\n bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)\n return z_indices\n\n def decode(self, x, half=False):\n \"\"\"\n Decodes quantized indices back to the input data space.\n\n Args:\n x (torch.Tensor): Quantized indices tensor.\n half (bool, optional): Whether the indices were generated with half quantization. Defaults to False.\n\n Returns:\n torch.Tensor: Reconstructed output tensor of shape (batch_size, seq_len, d_in).\n \"\"\"\n quantized = self.indices_to_bits(x, half)\n z = self.post_quant_embed(quantized)\n for layer in self.decoder:\n z = layer(z)\n z = self.head(z)\n return z\n\n\nclass Kronos(nn.Module, PyTorchModelHubMixin):\n \"\"\"\n Kronos Model.\n\n Args:\n s1_bits (int): Number of bits for pre tokens.\n s2_bits (int): Number of bits for post tokens.\n n_layers (int): Number of Transformer blocks.\n d_model (int): Dimension of the model's embeddings and hidden states.\n n_heads (int): Number of attention heads in the MultiheadAttention layers.\n ff_dim (int): Dimension of the feedforward network in the Transformer blocks.\n ffn_dropout_p (float): Dropout probability for the feedforward network.\n attn_dropout_p (float): Dropout probability for the attention layers.\n resid_dropout_p (float): Dropout probability for residual connections.\n token_dropout_p (float): Dropout probability for token embeddings.\n learn_te (bool): Whether to use learnable temporal embeddings.\n \"\"\"\n\n def __init__(self, s1_bits, s2_bits, n_layers, d_model, n_heads, ff_dim, ffn_dropout_p, attn_dropout_p, resid_dropout_p, token_dropout_p, learn_te, news_dim=None):\n super().__init__()\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.n_layers = n_layers\n self.d_model = d_model\n self.n_heads = n_heads\n self.learn_te = learn_te\n self.ff_dim = ff_dim\n self.ffn_dropout_p = ffn_dropout_p\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout_p = resid_dropout_p\n self.token_dropout_p = token_dropout_p\n self.news_dim = news_dim\n\n self.s1_vocab_size = 2 ** self.s1_bits\n self.token_drop = nn.Dropout(self.token_dropout_p)\n self.embedding = HierarchicalEmbedding(self.s1_bits, self.s2_bits, self.d_model)\n self.time_emb = TemporalEmbedding(self.d_model, self.learn_te)\n self.transformer = nn.ModuleList([\n TransformerBlock(self.d_model, self.n_heads, self.ff_dim, self.ffn_dropout_p, self.attn_dropout_p, self.resid_dropout_p)\n for _ in range(self.n_layers)\n ])\n self.norm = RMSNorm(self.d_model)\n self.dep_layer = DependencyAwareLayer(self.d_model)\n self.head = DualHead(self.s1_bits, self.s2_bits, self.d_model)\n\n if self.news_dim is not None:\n self.news_proj = nn.Linear(self.news_dim, self.d_model)\n else:\n self.news_proj = None\n\n self.apply(self._init_weights)\n\n def _init_weights(self, module):\n\n if isinstance(module, nn.Linear):\n nn.init.xavier_normal_(module.weight)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n elif isinstance(module, nn.Embedding):\n nn.init.normal_(module.weight, mean=0, std=self.embedding.d_model ** -0.5)\n elif isinstance(module, nn.LayerNorm):\n nn.init.ones_(module.weight)\n nn.init.zeros_(module.bias)\n elif isinstance(module, RMSNorm):\n nn.init.ones_(module.weight)\n\n def forward(self, s1_ids, s2_ids, stamp=None, padding_mask=None, use_teacher_forcing=False, s1_targets=None, news_emb=None):\n \"\"\"\n Args:\n s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]\n stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n use_teacher_forcing (bool, optional): Whether to use teacher forcing for s1 decoding. Defaults to False.\n s1_targets (torch.Tensor, optional): Target s1 token IDs for teacher forcing. Shape: [batch_size, seq_len]. Defaults to None.\n news_emb (torch.Tensor, optional): News embedding tensor. Shape: [batch_size, news_dim]. Defaults to None.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]:\n - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]\n - s2_logits: Logits for s2 token predictions, conditioned on s1. Shape: [batch_size, seq_len, s2_vocab_size]\n \"\"\"\n x = self.embedding([s1_ids, s2_ids])\n if stamp is not None:\n time_embedding = self.time_emb(stamp)\n x = x + time_embedding\n x = self.token_drop(x)\n\n for layer in self.transformer:\n x = layer(x, key_padding_mask=padding_mask)\n\n x = self.norm(x)\n\n if news_emb is not None and self.news_proj is not None:\n news_bias = self.news_proj(news_emb).unsqueeze(1) # [B, 1, d_model]\n x = x + news_bias\n\n s1_logits = self.head(x)\n\n if use_teacher_forcing:\n sibling_embed = self.embedding.emb_s1(s1_targets)\n else:\n s1_probs = F.softmax(s1_logits.detach(), dim=-1)\n sample_s1_ids = torch.multinomial(s1_probs.view(-1, self.s1_vocab_size), 1).view(s1_ids.shape)\n sibling_embed = self.embedding.emb_s1(sample_s1_ids)\n\n x2 = self.dep_layer(x, sibling_embed, key_padding_mask=padding_mask) # Dependency Aware Layer: Condition on s1 embeddings\n s2_logits = self.head.cond_forward(x2)\n return s1_logits, s2_logits\n\n def decode_s1(self, s1_ids, s2_ids, stamp=None, padding_mask=None, news_emb=None):\n \"\"\"\n Decodes only the s1 tokens.\n\n This method performs a forward pass to predict only s1 tokens. It returns the s1 logits\n and the context representation from the Transformer, which can be used for subsequent s2 decoding.\n\n Args:\n s1_ids (torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n s2_ids (torch.Tensor): Input tensor of s2 token IDs. Shape: [batch_size, seq_len]\n stamp (torch.Tensor, optional): Temporal stamp tensor. Shape: [batch_size, seq_len]. Defaults to None.\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n news_emb (torch.Tensor, optional): News embedding tensor. Shape: [batch_size, news_dim]. Defaults to None.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]:\n - s1 logits: Logits for s1 token predictions. Shape: [batch_size, seq_len, s1_vocab_size]\n - context: Context representation from the Transformer. Shape: [batch_size, seq_len, d_model]\n \"\"\"\n x = self.embedding([s1_ids, s2_ids])\n if stamp is not None:\n time_embedding = self.time_emb(stamp)\n x = x + time_embedding\n x = self.token_drop(x)\n\n for layer in self.transformer:\n x = layer(x, key_padding_mask=padding_mask)\n\n x = self.norm(x)\n\n if news_emb is not None and self.news_proj is not None:\n news_bias = self.news_proj(news_emb).unsqueeze(1) # [B, 1, d_model]\n x = x + news_bias\n\n s1_logits = self.head(x)\n return s1_logits, x\n\n def decode_s2(self, context, s1_ids, padding_mask=None):\n \"\"\"\n Decodes the s2 tokens, conditioned on the context and s1 tokens.\n\n This method decodes s2 tokens based on a pre-computed context representation (typically from `decode_s1`)\n and the s1 token IDs. It uses the dependency-aware layer and the conditional s2 head to predict s2 tokens.\n\n Args:\n context (torch.Tensor): Context representation from the transformer (output of decode_s1).\n Shape: [batch_size, seq_len, d_model]\n s1_ids (torch.torch.Tensor): Input tensor of s1 token IDs. Shape: [batch_size, seq_len]\n padding_mask (torch.Tensor, optional): Mask for padding tokens. Shape: [batch_size, seq_len]. Defaults to None.\n\n Returns:\n torch.Tensor: s2 logits. Shape: [batch_size, seq_len, s2_vocab_size]\n \"\"\"\n sibling_embed = self.embedding.emb_s1(s1_ids)\n x2 = self.dep_layer(context, sibling_embed, key_padding_mask=padding_mask)\n return self.head.cond_forward(x2)\n\n\ndef top_k_top_p_filtering(\n logits,\n top_k: int = 0,\n top_p: float = 1.0,\n filter_value: float = -float(\"Inf\"),\n min_tokens_to_keep: int = 1,\n):\n \"\"\"Filter a distribution of logits using top-k and/or nucleus (top-p) filtering\n Args:\n logits: logits distribution shape (batch size, vocabulary size)\n if top_k > 0: keep only top k tokens with highest probability (top-k filtering).\n if top_p \u003c 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).\n Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)\n Make sure we keep at least min_tokens_to_keep per batch example in the output\n From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317\n \"\"\"\n if top_k > 0:\n top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1)) # Safety check\n # Remove all tokens with a probability less than the last token of the top-k\n indices_to_remove = logits \u003c torch.topk(logits, top_k)[0][..., -1, None]\n logits[indices_to_remove] = filter_value\n return logits\n\n if top_p \u003c 1.0:\n sorted_logits, sorted_indices = torch.sort(logits, descending=True)\n cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n\n # Remove tokens with cumulative probability above the threshold (token with 0 are kept)\n sorted_indices_to_remove = cumulative_probs > top_p\n if min_tokens_to_keep > 1:\n # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)\n sorted_indices_to_remove[..., :min_tokens_to_keep] = 0\n # Shift the indices to the right to keep also the first token above the threshold\n sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n sorted_indices_to_remove[..., 0] = 0\n\n # scatter sorted tensors to original indexing\n indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n logits[indices_to_remove] = filter_value\n return logits\n\n\ndef sample_from_logits(logits, temperature=1.0, top_k=None, top_p=None, sample_logits=True):\n logits = logits / temperature\n if top_k is not None or top_p is not None:\n if top_k > 0 or top_p \u003c 1.0:\n logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)\n\n probs = F.softmax(logits, dim=-1)\n\n if not sample_logits:\n _, x = top_k(probs, k=1, dim=-1)\n else:\n x = torch.multinomial(probs, num_samples=1)\n\n return x\n\n\ndef auto_regressive_inference(tokenizer, model, x, x_stamp, y_stamp, max_context, pred_len, clip=5, T=1.0, top_k=0, top_p=0.99, sample_count=5, verbose=False, news_emb=None):\n with torch.no_grad():\n x = torch.clip(x, -clip, clip)\n\n device = x.device\n x = x.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x.size(1), x.size(2)).to(device)\n x_stamp = x_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, x_stamp.size(1), x_stamp.size(2)).to(device)\n y_stamp = y_stamp.unsqueeze(1).repeat(1, sample_count, 1, 1).reshape(-1, y_stamp.size(1), y_stamp.size(2)).to(device)\n\n x_token = tokenizer.encode(x, half=True)\n \n initial_seq_len = x.size(1)\n batch_size = x_token[0].size(0)\n total_seq_len = initial_seq_len + pred_len\n full_stamp = torch.cat([x_stamp, y_stamp], dim=1)\n\n generated_pre = x_token[0].new_empty(batch_size, pred_len)\n generated_post = x_token[1].new_empty(batch_size, pred_len)\n\n pre_buffer = x_token[0].new_zeros(batch_size, max_context)\n post_buffer = x_token[1].new_zeros(batch_size, max_context)\n buffer_len = min(initial_seq_len, max_context)\n if buffer_len > 0:\n start_idx = max(0, initial_seq_len - max_context)\n pre_buffer[:, :buffer_len] = x_token[0][:, start_idx:start_idx + buffer_len]\n post_buffer[:, :buffer_len] = x_token[1][:, start_idx:start_idx + buffer_len]\n\n if verbose:\n ran = trange\n else:\n ran = range\n for i in ran(pred_len):\n current_seq_len = initial_seq_len + i\n window_len = min(current_seq_len, max_context)\n\n if current_seq_len \u003c= max_context:\n input_tokens = [\n pre_buffer[:, :window_len],\n post_buffer[:, :window_len]\n ]\n else:\n input_tokens = [pre_buffer, post_buffer]\n\n context_end = current_seq_len\n context_start = max(0, context_end - max_context)\n current_stamp = full_stamp[:, context_start:context_end, :].contiguous()\n\n s1_logits, context = model.decode_s1(input_tokens[0], input_tokens[1], current_stamp, news_emb=news_emb)\n s1_logits = s1_logits[:, -1, :]\n sample_pre = sample_from_logits(s1_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)\n\n s2_logits = model.decode_s2(context, sample_pre)\n s2_logits = s2_logits[:, -1, :]\n sample_post = sample_from_logits(s2_logits, temperature=T, top_k=top_k, top_p=top_p, sample_logits=True)\n\n generated_pre[:, i] = sample_pre.squeeze(-1)\n generated_post[:, i] = sample_post.squeeze(-1)\n\n if current_seq_len \u003c max_context:\n pre_buffer[:, current_seq_len] = sample_pre.squeeze(-1)\n post_buffer[:, current_seq_len] = sample_post.squeeze(-1)\n else:\n pre_buffer.copy_(torch.roll(pre_buffer, shifts=-1, dims=1))\n post_buffer.copy_(torch.roll(post_buffer, shifts=-1, dims=1))\n pre_buffer[:, -1] = sample_pre.squeeze(-1)\n post_buffer[:, -1] = sample_post.squeeze(-1)\n\n full_pre = torch.cat([x_token[0], generated_pre], dim=1)\n full_post = torch.cat([x_token[1], generated_post], dim=1)\n\n context_start = max(0, total_seq_len - max_context)\n input_tokens = [\n full_pre[:, context_start:total_seq_len].contiguous(),\n full_post[:, context_start:total_seq_len].contiguous()\n ]\n z = tokenizer.decode(input_tokens, half=True)\n z = z.reshape(-1, sample_count, z.size(1), z.size(2))\n preds = z.cpu().numpy()\n preds = np.mean(preds, axis=1)\n\n return preds\n\n\ndef calc_time_stamps(x_timestamp):\n time_df = pd.DataFrame()\n time_df['minute'] = x_timestamp.dt.minute\n time_df['hour'] = x_timestamp.dt.hour\n time_df['weekday'] = x_timestamp.dt.weekday\n time_df['day'] = x_timestamp.dt.day\n time_df['month'] = x_timestamp.dt.month\n return time_df\n\n\nclass KronosPredictor:\n\n def __init__(self, model, tokenizer, device=\"cuda:0\", max_context=512, clip=5):\n self.tokenizer = tokenizer\n self.model = model\n self.max_context = max_context\n self.clip = clip\n self.price_cols = ['open', 'high', 'low', 'close']\n self.vol_col = 'volume'\n self.amt_vol = 'amount'\n self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month']\n self.device = device\n\n self.tokenizer = self.tokenizer.to(self.device)\n self.model = self.model.to(self.device)\n\n def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, news_emb=None):\n\n x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device)\n x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device)\n y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device)\n\n preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len,\n self.clip, T, top_k, top_p, sample_count, verbose, news_emb=news_emb)\n preds = preds[:, -pred_len:, :]\n return preds\n\n def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True, news_emb=None):\n\n if not isinstance(df, pd.DataFrame):\n raise ValueError(\"Input must be a pandas DataFrame.\")\n\n if not all(col in df.columns for col in self.price_cols):\n raise ValueError(f\"Price columns {self.price_cols} not found in DataFrame.\")\n\n df = df.copy()\n if self.vol_col not in df.columns:\n df[self.vol_col] = 0.0 # Fill missing volume with zeros\n df[self.amt_vol] = 0.0 # Fill missing amount with zeros\n if self.amt_vol not in df.columns and self.vol_col in df.columns:\n df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)\n\n if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():\n raise ValueError(\"Input DataFrame contains NaN values in price or volume columns.\")\n\n x_time_df = calc_time_stamps(x_timestamp)\n y_time_df = calc_time_stamps(y_timestamp)\n\n x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)\n x_stamp = x_time_df.values.astype(np.float32)\n y_stamp = y_time_df.values.astype(np.float32)\n\n x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)\n\n x = (x - x_mean) / (x_std + 1e-5)\n x = np.clip(x, -self.clip, self.clip)\n\n x = x[np.newaxis, :]\n x_stamp = x_stamp[np.newaxis, :]\n y_stamp = y_stamp[np.newaxis, :]\n\n if news_emb is not None:\n news_emb_tensor = torch.from_numpy(np.array(news_emb).astype(np.float32)).to(self.device)\n # Ensure batch dimension for news_emb if only one sample\n if news_emb_tensor.ndim == 1:\n news_emb_tensor = news_emb_tensor.unsqueeze(0)\n else:\n news_emb_tensor = None\n\n preds = self.generate(x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose, news_emb=news_emb_tensor)\n\n preds = preds.squeeze(0)\n preds = preds * (x_std + 1e-5) + x_mean\n\n pred_df = pd.DataFrame(preds, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp)\n return pred_df\n\n\n def predict_batch(self, df_list, x_timestamp_list, y_timestamp_list, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True):\n \"\"\"\n Perform parallel (batch) prediction on multiple time series. All series must have the same historical length and prediction length (pred_len).\n\n Args:\n df_list (List[pd.DataFrame]): List of input DataFrames, each containing price columns and optional volume/amount columns.\n x_timestamp_list (List[pd.DatetimeIndex or Series]): List of timestamps corresponding to historical data, length should match the number of rows in each DataFrame.\n y_timestamp_list (List[pd.DatetimeIndex or Series]): List of future prediction timestamps, length should equal pred_len.\n pred_len (int): Number of prediction steps.\n T (float): Sampling temperature.\n top_k (int): Top-k filtering threshold.\n top_p (float): Top-p (nucleus sampling) threshold.\n sample_count (int): Number of parallel samples per series, automatically averaged internally.\n verbose (bool): Whether to display autoregressive progress.\n\n Returns:\n List[pd.DataFrame]: List of prediction results in the same order as input, each DataFrame contains\n `open, high, low, close, volume, amount` columns, indexed by corresponding `y_timestamp`.\n \"\"\"\n # Basic validation\n if not isinstance(df_list, (list, tuple)) or not isinstance(x_timestamp_list, (list, tuple)) or not isinstance(y_timestamp_list, (list, tuple)):\n raise ValueError(\"df_list, x_timestamp_list, y_timestamp_list must be list or tuple types.\")\n if not (len(df_list) == len(x_timestamp_list) == len(y_timestamp_list)):\n raise ValueError(\"df_list, x_timestamp_list, y_timestamp_list must have consistent lengths.\")\n\n num_series = len(df_list)\n\n x_list = []\n x_stamp_list = []\n y_stamp_list = []\n means = []\n stds = []\n seq_lens = []\n y_lens = []\n\n for i in range(num_series):\n df = df_list[i]\n if not isinstance(df, pd.DataFrame):\n raise ValueError(f\"Input at index {i} is not a pandas DataFrame.\")\n if not all(col in df.columns for col in self.price_cols):\n raise ValueError(f\"DataFrame at index {i} is missing price columns {self.price_cols}.\")\n\n df = df.copy()\n if self.vol_col not in df.columns:\n df[self.vol_col] = 0.0\n df[self.amt_vol] = 0.0\n if self.amt_vol not in df.columns and self.vol_col in df.columns:\n df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1)\n\n if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any():\n raise ValueError(f\"DataFrame at index {i} contains NaN values in price or volume columns.\")\n\n x_timestamp = x_timestamp_list[i]\n y_timestamp = y_timestamp_list[i]\n\n x_time_df = calc_time_stamps(x_timestamp)\n y_time_df = calc_time_stamps(y_timestamp)\n\n x = df[self.price_cols + [self.vol_col, self.amt_vol]].values.astype(np.float32)\n x_stamp = x_time_df.values.astype(np.float32)\n y_stamp = y_time_df.values.astype(np.float32)\n\n if x.shape[0] != x_stamp.shape[0]:\n raise ValueError(f\"Inconsistent lengths at index {i}: x has {x.shape[0]} vs x_stamp has {x_stamp.shape[0]}.\")\n if y_stamp.shape[0] != pred_len:\n raise ValueError(f\"y_timestamp length at index {i} should equal pred_len={pred_len}, got {y_stamp.shape[0]}.\")\n\n x_mean, x_std = np.mean(x, axis=0), np.std(x, axis=0)\n x_norm = (x - x_mean) / (x_std + 1e-5)\n x_norm = np.clip(x_norm, -self.clip, self.clip)\n\n x_list.append(x_norm)\n x_stamp_list.append(x_stamp)\n y_stamp_list.append(y_stamp)\n means.append(x_mean)\n stds.append(x_std)\n\n seq_lens.append(x_norm.shape[0])\n y_lens.append(y_stamp.shape[0])\n\n # Require all series to have consistent historical and prediction lengths for batch processing\n if len(set(seq_lens)) != 1:\n raise ValueError(f\"Parallel prediction requires all series to have consistent historical lengths, got: {seq_lens}\")\n if len(set(y_lens)) != 1:\n raise ValueError(f\"Parallel prediction requires all series to have consistent prediction lengths, got: {y_lens}\")\n\n x_batch = np.stack(x_list, axis=0).astype(np.float32) # (B, seq_len, feat)\n x_stamp_batch = np.stack(x_stamp_list, axis=0).astype(np.float32) # (B, seq_len, time_feat)\n y_stamp_batch = np.stack(y_stamp_list, axis=0).astype(np.float32) # (B, pred_len, time_feat)\n\n preds = self.generate(x_batch, x_stamp_batch, y_stamp_batch, pred_len, T, top_k, top_p, sample_count, verbose)\n # preds: (B, pred_len, feat)\n\n pred_dfs = []\n for i in range(num_series):\n preds_i = preds[i] * (stds[i] + 1e-5) + means[i]\n pred_df = pd.DataFrame(preds_i, columns=self.price_cols + [self.vol_col, self.amt_vol], index=y_timestamp_list[i])\n pred_dfs.append(pred_df)\n\n return pred_dfs\n","content_type":"text/x-python; charset=utf-8","language":"python","size":31093,"content_sha256":"b339c8685f40322663eb6c3f0d56dc91e64b65fae1f9cc22e414d6c76b9c7b78"},{"filename":"scripts/utils/predictor/model/module.py","content":"import math\n\nfrom einops import rearrange, reduce\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import Function\nimport torch.nn.functional as F\n\n\nclass DifferentiableEntropyFunction(Function):\n @staticmethod\n def forward(ctx, zq, basis, K, eps):\n zb = (zq + 1) / 2\n zi = ((zb * basis).sum(-1)).to(torch.int64)\n cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),\n 0,\n zi.flatten(),\n torch.ones_like(zi.flatten()).to(zq.dtype),\n 'sum')\n prob = (cnt + eps) / (cnt + eps).sum()\n H = -(prob * torch.log(prob)).sum()\n ctx.save_for_backward(zq, zi, prob)\n ctx.K = K\n return H\n\n @staticmethod\n def backward(ctx, grad_output):\n zq, zi, prob = ctx.saved_tensors\n grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K\n reord_grad = grad_array[zi.flatten()].reshape(zi.shape)\n grad_input = reord_grad.unsqueeze(-1) * zq\n return grad_input, None, None, None, None\n\n\ndef codebook_entropy(zq, basis, K, eps=1e-4):\n return DifferentiableEntropyFunction.apply(zq, basis, K, eps)\n\n\nclass BinarySphericalQuantizer(nn.Module):\n def __init__(self, embed_dim, beta, gamma0, gamma, zeta,\n input_format='bchw',\n soft_entropy=True, group_size=9,\n persample_entropy_compute='analytical',\n cb_entropy_compute='group',\n l2_norm=True,\n inv_temperature=1):\n \"\"\"\n Paper link: https://arxiv.org/pdf/2406.07548.pdf\n Here we use the official implementation of the BinarySphericalQuantizer.\n \"\"\"\n super().__init__()\n self.embed_dim = embed_dim\n self.beta = beta # loss weight for commit loss\n self.gamma0 = gamma0 # loss weight for entropy penalty\n self.gamma = gamma # loss weight for entropy penalty\n self.zeta = zeta # loss weight for entire entropy penalty\n self.input_format = input_format\n assert self.embed_dim % group_size == 0, \"embed_dim must be divisible by group_size\"\n self.num_groups = self.embed_dim // group_size\n self.group_size = group_size\n assert persample_entropy_compute in ['group', 'analytical'], \"persample_entropy_compute must be either 'group' or 'analytical'\"\n assert cb_entropy_compute in ['group', 'nce'], \"cb_entropy_compute must be either 'group' or 'nce'\"\n self.persample_entropy_compute = persample_entropy_compute\n self.cb_entropy_compute = cb_entropy_compute\n self.l2_norm = l2_norm\n self.inv_temperature = inv_temperature\n\n self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))\n self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))\n\n self.num_dimensions = 2 ** embed_dim\n self.bits_per_index = embed_dim\n\n # we only need to keep the codebook portion up to the group size\n # because we approximate the H loss with this subcode\n group_codes = torch.arange(2 ** self.group_size)\n group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]\n self.register_buffer('group_codebook', group_codebook, persistent=False)\n\n self.soft_entropy = soft_entropy # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf\n\n def quantize(self, z):\n assert z.shape[-1] == self.embed_dim, f\"Expected {self.embed_dim} dimensions, got {z.shape[-1]}\"\n\n zhat = torch.where(z > 0,\n torch.tensor(1, dtype=z.dtype, device=z.device),\n torch.tensor(-1, dtype=z.dtype, device=z.device))\n return z + (zhat - z).detach()\n\n def forward(self, z, collect_metrics=True):\n # if self.input_format == 'bchw':\n # z = rearrange(z, 'b c h w -> b h w c')\n zq = self.quantize(z)\n\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n\n zq = zq * q_scale\n\n if not collect_metrics:\n return zq, zq.new_zeros(()), {}\n\n indices = self.codes_to_indexes(zq.detach())\n group_indices = self.codes_to_group_indexes(zq.detach())\n if not self.training:\n used_codes = torch.unique(indices, return_counts=False)\n else:\n used_codes = None\n\n if self.soft_entropy:\n persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)\n entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy\n else:\n zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)\n persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)\n cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)\n entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy\n\n # commit loss\n commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))\n\n # if self.input_format == 'bchw':\n # zq = rearrange(zq, 'b h w c -> b c h w')\n\n return (\n zq,\n commit_loss + self.zeta * entropy_penalty / self.inv_temperature,\n {\"H\": cb_entropy, \"used_codes\": used_codes, \"indices\": indices, \"group_indices\": group_indices,\n \"avg_prob\": avg_prob}\n )\n\n def soft_entropy_loss(self, z):\n # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size\n # the sub-code is the last group_size bits of the full code\n group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)\n divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)\n\n # we calculate the distance between the divided_z and the codebook for each subgroup\n distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)\n prob = (-distance * self.inv_temperature).softmax(dim=-1)\n if self.persample_entropy_compute == 'analytical':\n if self.l2_norm:\n p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)\n else:\n p = torch.sigmoid(-4 * z * self.inv_temperature)\n prob = torch.stack([p, 1 - p], dim=-1)\n per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()\n else:\n per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()\n\n # macro average of the probability of each subgroup\n avg_prob = reduce(prob, '... g d ->g d', 'mean')\n codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)\n\n # the approximation of the entropy is the sum of the entropy of each subgroup\n return per_sample_entropy, codebook_entropy.sum(), avg_prob\n\n def get_hard_per_sample_entropy(self, zb_by_sample):\n probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]\n persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)\n persample_entropy = persample_entropy.sum(-1)\n return persample_entropy.mean()\n\n def codes_to_indexes(self, zhat):\n \"\"\"Converts a `code` to an index in the codebook.\n Args:\n zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}\n \"\"\"\n assert zhat.shape[-1] == self.embed_dim, f\"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}\"\n return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)\n\n def codes_to_group_indexes(self, zhat):\n \"\"\"Converts a `code` to a list of indexes (in groups) in the codebook.\n Args:\n zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}\n \"\"\"\n zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)\n return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)\n\n def indexes_to_codes(self, indices):\n \"\"\"Inverse of `indexes_to_codes`.\"\"\"\n indices = indices.unsqueeze(-1)\n codes_non_centered = torch.remainder(\n torch.floor_divide(indices, self.basis), 2\n )\n return codes_non_centered * 2 - 1\n\n def group_indexes_to_codes(self, group_indices):\n \"\"\"Inverse of `group_indexes_to_codes`.\"\"\"\n group_indices = group_indices.unsqueeze(-1)\n codes_non_centered = torch.remainder(\n torch.floor_divide(group_indices, self.group_basis), 2\n )\n codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')\n return codes_non_centered * 2 - 1\n\n def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):\n if normalize:\n probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)\n else:\n probs = count\n H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)\n return H\n\n def get_group_codebook_entry(self, group_indices):\n z_q = self.group_indexes_to_codes(group_indices)\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n z_q = z_q * q_scale\n if self.input_format == 'bchw':\n h, w = int(z_q.shape[1] ** 0.5)\n assert h * w == z_q.shape[1], 'Invalid sequence length'\n z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)\n return z_q\n\n def get_codebook_entry(self, indices):\n z_q = self.indexes_to_codes(indices)\n q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.\n z_q = z_q * q_scale\n if self.input_format == 'bchw':\n h, w = int(z_q.shape[1] ** 0.5)\n assert h * w == z_q.shape[1], 'Invalid sequence length'\n z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)\n return z_q\n\n\nclass BSQuantizer(nn.Module):\n\n def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):\n super().__init__()\n self.codebook_dim = s1_bits + s2_bits\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)\n\n def bits_to_indices(self, bits):\n bits = (bits >= 0).to(torch.long)\n indices = 2 ** torch.arange(\n 0,\n bits.shape[-1],\n 1,\n dtype=torch.long,\n device=bits.device,\n )\n return (bits * indices).sum(-1)\n\n def forward(self, z, half=False, collect_metrics=True):\n z = F.normalize(z, dim=-1)\n quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)\n if half:\n q_pre = quantized[:, :, :self.s1_bits]\n q_post = quantized[:, :, self.s1_bits:]\n z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]\n else:\n z_indices = self.bits_to_indices(quantized)\n return bsq_loss, quantized, z_indices\n\n\nclass RMSNorm(torch.nn.Module):\n def __init__(self, dim: int, eps: float = 1e-5):\n super().__init__()\n self.eps = eps\n self.weight = nn.Parameter(torch.ones(dim))\n\n def _norm(self, x):\n return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)\n\n def forward(self, x):\n output = self._norm(x.float()).type_as(x)\n return output * self.weight\n\n\nclass FeedForward(nn.Module):\n def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):\n super().__init__()\n\n self.w1 = nn.Linear(d_model, ff_dim, bias=False)\n self.w3 = nn.Linear(d_model, ff_dim, bias=False)\n self.w2 = nn.Linear(ff_dim, d_model, bias=False)\n self.ffn_dropout = nn.Dropout(ffn_dropout_p)\n\n def forward(self, x):\n return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))\n\n\nclass RotaryPositionalEmbedding(nn.Module):\n def __init__(self, dim):\n super().__init__()\n inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))\n self.register_buffer(\"inv_freq\", inv_freq)\n self.seq_len_cached = None\n self.cos_cached = None\n self.sin_cached = None\n\n def _update_cos_sin_cache(self, x, seq_len):\n if seq_len != self.seq_len_cached:\n self.seq_len_cached = seq_len\n t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)\n freqs = torch.einsum('i,j->ij', t, self.inv_freq)\n emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n self.cos_cached = emb.cos()[None, None, :, :]\n self.sin_cached = emb.sin()[None, None, :, :]\n return self.cos_cached, self.sin_cached\n\n def forward(self, q, k):\n cos, sin = self._update_cos_sin_cache(q, q.shape[-2])\n return (\n (q * cos) + (self._rotate_half(q) * sin),\n (k * cos) + (self._rotate_half(k) * sin),\n )\n\n def _rotate_half(self, x):\n x1, x2 = x.chunk(2, dim=-1)\n return torch.cat((-x2, x1), dim=-1)\n\n\nclass MultiHeadAttentionWithRoPE(nn.Module):\n def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):\n super().__init__()\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n\n self.q_proj = nn.Linear(d_model, d_model)\n self.k_proj = nn.Linear(d_model, d_model)\n self.v_proj = nn.Linear(d_model, d_model)\n self.out_proj = nn.Linear(d_model, d_model)\n self.rotary = RotaryPositionalEmbedding(self.head_dim)\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout = nn.Dropout(resid_dropout_p)\n\n def forward(self, x, key_padding_mask=None):\n batch_size, seq_len, _ = x.shape\n\n q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n\n q, k = self.rotary(q, k)\n\n if key_padding_mask is not None:\n attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]\n attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1) # [batch, n_heads, q_len, k_len]\n else:\n attn_mask = None\n\n attn_output = F.scaled_dot_product_attention(\n q, k, v,\n attn_mask=attn_mask,\n dropout_p=self.attn_dropout_p if self.training else 0.0,\n is_causal=True\n )\n\n attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)\n return self.resid_dropout(self.out_proj(attn_output))\n\n\nclass MultiHeadCrossAttentionWithRoPE(nn.Module):\n def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):\n super().__init__()\n self.d_model = d_model\n self.n_heads = n_heads\n self.head_dim = d_model // n_heads\n\n self.q_proj = nn.Linear(d_model, d_model)\n self.k_proj = nn.Linear(d_model, d_model)\n self.v_proj = nn.Linear(d_model, d_model)\n self.out_proj = nn.Linear(d_model, d_model)\n self.rotary = RotaryPositionalEmbedding(self.head_dim)\n self.attn_dropout_p = attn_dropout_p\n self.resid_dropout = nn.Dropout(resid_dropout)\n\n def forward(self, query, key, value, key_padding_mask=None):\n batch_size, q_len, _ = query.shape\n _, seq_len, _ = key.shape\n\n q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)\n k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)\n\n q, k = self.rotary(q, k)\n\n if key_padding_mask is not None:\n attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)\n attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)\n else:\n attn_mask = None\n\n is_causal_flag = self.training\n\n attn_output = F.scaled_dot_product_attention(\n q, k, v,\n attn_mask=attn_mask,\n dropout_p=self.attn_dropout_p if self.training else 0.0,\n is_causal=is_causal_flag\n )\n\n attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)\n return self.resid_dropout(self.out_proj(attn_output))\n\n\nclass HierarchicalEmbedding(nn.Module):\n def __init__(self, s1_bits, s2_bits, d_model=256):\n super().__init__()\n self.s1_bits = s1_bits\n self.s2_bits = s2_bits\n\n vocab_s1 = 2 ** s1_bits\n vocab_s2 = 2 ** s2_bits\n\n self.emb_s1 = nn.Embedding(vocab_s1, d_model)\n self.emb_s2 = nn.Embedding(vocab_s2, d_model)\n self.d_model = d_model\n self.fusion_proj = nn.Linear(d_model * 2, d_model)\n\n nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)\n nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)\n\n def split_token(self, token_ids: torch.Tensor, s2_bits: int):\n \"\"\"Inputs:\n token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].\n s2_bits (int): Number of low bits used for the fine token (s2).\n \"\"\"\n assert isinstance(s2_bits, int) and s2_bits > 0, \"s2_bits must be a positive integer\"\n\n t = token_ids.long()\n mask = (1 \u003c\u003c s2_bits) - 1\n s2_ids = t & mask # extract low bits\n s1_ids = t >> s2_bits # extract high bits\n return s1_ids, s2_ids\n\n def forward(self, token_ids):\n \"\"\"Inputs:\n token_ids:\n - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or\n - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.\n Output: [batch_size, seq_len, d_model]\n \"\"\"\n if isinstance(token_ids, tuple) or isinstance(token_ids, list):\n s1_ids, s2_ids = token_ids\n else:\n s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)\n s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)\n s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)\n return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))\n\n\nclass DependencyAwareLayer(nn.Module):\n def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):\n super().__init__()\n self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)\n self.norm = RMSNorm(d_model)\n\n def forward(self, hidden_states, sibling_embed, key_padding_mask=None):\n \"\"\"hidden_states: [batch, seq_len, d_model]\n sibling_embed: Embedding from another subtoken\n \"\"\"\n attn_out = self.cross_attn(\n query=sibling_embed,\n key=hidden_states,\n value=hidden_states,\n key_padding_mask=key_padding_mask\n )\n return self.norm(hidden_states + attn_out)\n\n\nclass TransformerBlock(nn.Module):\n def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):\n super().__init__()\n self.norm1 = RMSNorm(d_model)\n self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)\n self.norm2 = RMSNorm(d_model)\n self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)\n\n def forward(self, x, key_padding_mask=None):\n residual = x\n x = self.norm1(x)\n attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)\n x = residual + attn_out\n\n residual = x\n x = self.norm2(x)\n ffn_out = self.ffn(x)\n x = residual + ffn_out\n return x\n\n\nclass DualHead(nn.Module):\n def __init__(self, s1_bits, s2_bits, d_model):\n super().__init__()\n self.vocab_s1 = 2 ** s1_bits\n self.vocab_s2 = 2 ** s2_bits\n self.proj_s1 = nn.Linear(d_model, self.vocab_s1)\n self.proj_s2 = nn.Linear(d_model, self.vocab_s2)\n\n def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):\n if padding_mask is not None:\n valid_mask = (padding_mask == 0)\n s1_logits = s1_logits[valid_mask]\n s2_logits = s2_logits[valid_mask]\n s1_targets = s1_targets[valid_mask]\n s2_targets = s2_targets[valid_mask]\n ce_s1 = F.cross_entropy(s1_logits, s1_targets)\n ce_s2 = F.cross_entropy(s2_logits, s2_targets)\n else:\n ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))\n ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))\n ce_loss = (ce_s1 + ce_s2) / 2\n return ce_loss, ce_s1, ce_s2\n\n def forward(self, x):\n return self.proj_s1(x)\n\n def cond_forward(self, x2):\n return self.proj_s2(x2)\n\n\nclass FixedEmbedding(nn.Module):\n def __init__(self, c_in, d_model):\n super(FixedEmbedding, self).__init__()\n\n w = torch.zeros(c_in, d_model).float()\n w.require_grad = False\n\n position = torch.arange(0, c_in).float().unsqueeze(1)\n div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()\n\n w[:, 0::2] = torch.sin(position * div_term)\n w[:, 1::2] = torch.cos(position * div_term)\n\n self.emb = nn.Embedding(c_in, d_model)\n self.emb.weight = nn.Parameter(w, requires_grad=False)\n\n def forward(self, x):\n return self.emb(x).detach()\n\n\nclass TemporalEmbedding(nn.Module):\n def __init__(self, d_model, learn_pe):\n super(TemporalEmbedding, self).__init__()\n\n minute_size = 60\n hour_size = 24\n weekday_size = 7\n day_size = 32\n month_size = 13\n\n Embed = FixedEmbedding if not learn_pe else nn.Embedding\n self.minute_embed = Embed(minute_size, d_model)\n self.hour_embed = Embed(hour_size, d_model)\n self.weekday_embed = Embed(weekday_size, d_model)\n self.day_embed = Embed(day_size, d_model)\n self.month_embed = Embed(month_size, d_model)\n\n def forward(self, x):\n x = x.long()\n\n minute_x = self.minute_embed(x[:, :, 0])\n hour_x = self.hour_embed(x[:, :, 1])\n weekday_x = self.weekday_embed(x[:, :, 2])\n day_x = self.day_embed(x[:, :, 3])\n month_x = self.month_embed(x[:, :, 4])\n\n return hour_x + weekday_x + day_x + month_x + minute_x","content_type":"text/x-python; charset=utf-8","language":"python","size":22847,"content_sha256":"91e8ff4ea1d77a95e39be79a217c360165e585d90475c7a8369c2833953ba181"},{"filename":"scripts/utils/predictor/training.py","content":"import os\nimport sys\nimport time\nimport torch\nimport torch.nn as nn\nimport pandas as pd\nimport numpy as np\nimport json\nimport random\nfrom loguru import logger\nfrom datetime import datetime, timedelta\nfrom sentence_transformers import SentenceTransformer\nfrom dotenv import load_dotenv\n\n# Load environment variables\nload_dotenv()\n\n# Setup paths\nKRONOS_DIR = os.path.dirname(os.path.abspath(__file__))\nSRC_DIR = os.path.dirname(os.path.dirname(KRONOS_DIR))\nif SRC_DIR not in sys.path:\n sys.path.insert(0, SRC_DIR)\n\nfrom ..kronos.model import Kronos, KronosTokenizer, KronosPredictor\nfrom ..database_manager import DatabaseManager\nfrom ..stock_tools import StockTools\nfrom ..search_tools import SearchTools\nfrom ..llm.factory import get_model\nfrom ..visualizer import VisualizerTools\nfrom ..schema.models import ForecastResult, KLinePoint\nfrom agno.agent import Agent\n\nclass AutoSynthesisTrainer:\n def __init__(self, news_dim=384):\n self.device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n self.db = DatabaseManager()\n self.tools = StockTools(self.db)\n self.searcher = SearchTools(self.db)\n # Try loading from local cache first to avoid network timeouts\n model_name = os.getenv('EMBEDDING_MODEL', 'sentence-transformers/all-MiniLM-L6-v2')\n try:\n logger.info(f\"🔄 Attempting to load {model_name} from local cache...\")\n self.embedder = SentenceTransformer(model_name, device=self.device, local_files_only=True)\n logger.success(\"✅ Model loaded from local cache.\")\n except Exception:\n logger.warning(\"⚠️ Local cache not found or incomplete. Attempting to download...\")\n self.embedder = SentenceTransformer(model_name, device=self.device)\n self.news_dim = news_dim\n \n # Try loading from local cache first to avoid network timeouts\n try:\n logger.info(\"🔄 Attempting to load Kronos and Tokenizer from local cache...\")\n self.tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\", local_files_only=True).to(self.device)\n base_model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\", local_files_only=True)\n logger.success(\"✅ Kronos and Tokenizer loaded from local cache.\")\n except Exception:\n logger.warning(\"⚠️ Local Kronos/Tokenizer not found or incomplete. Attempting to download...\")\n self.tokenizer = KronosTokenizer.from_pretrained(\"NeoQuasar/Kronos-Tokenizer-base\").to(self.device)\n base_model = Kronos.from_pretrained(\"NeoQuasar/Kronos-base\")\n \n self.model = Kronos(\n base_model.s1_bits, base_model.s2_bits, base_model.n_layers, \n base_model.d_model, base_model.n_heads, base_model.ff_dim,\n base_model.ffn_dropout_p, base_model.attn_dropout_p,\n base_model.resid_dropout_p, base_model.token_dropout_p,\n base_model.learn_te, news_dim=self.news_dim\n ).to(self.device)\n self.model.load_state_dict(base_model.state_dict(), strict=False)\n \n # LLM for causality verification\n provider = os.getenv(\"LLM_PROVIDER\", \"ust\")\n model_id = os.getenv(\"LLM_MODEL\", \"Qwen\")\n self.llm_agent = Agent(model=get_model(provider, model_id))\n\n def discover_shocks(self, ticker_list, threshold=2.0, limit_per_stock=5, days=365, pred_len=5):\n \"\"\"1. Find days with significant price movements (Look back 1 year)\"\"\"\n shocks = []\n end_date = datetime.now().strftime(\"%Y-%m-%d\")\n start_date = (datetime.now() - timedelta(days=days)).strftime(\"%Y-%m-%d\")\n \n for ticker in ticker_list:\n df = self.tools.get_stock_price(ticker, start_date=start_date, end_date=end_date)\n if df.empty or len(df) \u003c 60:\n continue\n \n # Look for big moves\n moves = df[df['change_pct'].abs() > threshold].copy()\n if moves.empty: continue\n \n count = 0\n for idx, row in moves.iterrows():\n # Ensure we have history before this day AND enough future days for eval\n date_idx = df.index.get_loc(idx)\n if date_idx \u003c 50 or date_idx + pred_len > len(df): continue\n \n shocks.append({\n 'ticker': ticker,\n 'date': row['date'],\n 'change': row['change_pct'],\n 'history': df.iloc[date_idx-50:date_idx],\n 'target': df.iloc[date_idx:date_idx + pred_len] # Now capturing pred_len days\n })\n count += 1\n if count >= limit_per_stock: break\n \n logger.info(f\"✨ Discovered {len(shocks)} potential price shocks over the last {days} days.\")\n return shocks\n\n def find_reason_and_verify(self, shock):\n \"\"\"2. Search for reasons and verify causality using LLM\"\"\"\n ticker_info = self.db.get_stock_by_code(shock['ticker'])\n name = ticker_info['name'] if ticker_info else shock['ticker']\n date_str = shock['date']\n \n # Try multiple query variations and engines\n queries = [\n f\"{name} ({shock['ticker']}) {date_str} 为什么涨跌原因\",\n f\"{name} {date_str} 异动原因\",\n f\"{shock['ticker']} {date_str} 新闻\"\n ]\n \n search_results = []\n for query in queries:\n logger.info(f\"🔍 Searching for reason: {query}\")\n # Try alternate engines\n for engine in [\"baidu\"]:\n try:\n results = self.searcher.search_list(query, engine=engine, max_results=3, enrich=False)\n if results:\n search_results = results\n break\n except Exception as e:\n logger.warning(f\"Search failed for {query} on {engine}: {e}\")\n \n if search_results:\n break\n time.sleep(random.uniform(1.0, 2.0))\n \n if not search_results:\n logger.warning(f\"⚠️ No search results found for {name} on {date_str} after multiple attempts.\")\n return None\n \n context = \"\\n\".join([f\"- {r['title']}: {r.get('content', '')[:300]}\" for r in search_results])\n \n prompt = f\"\"\"\n 任务：判断以下新闻是否解释了该股票在 {date_str} 的 {shock['change']:.2f}% 价格变动。\n \n 股票：{name}\n 日期：{date_str}\n 变动：{shock['change']:.2f}%\n \n 搜索结果：\n {context}\n \n 要求：\n 1. 该新闻是否在该日期左右发生？\n 2. 该新闻是否能逻辑上解释这种大幅波动（如财报、利好政策、重组、大环境暴跌等）？\n 3. 如果是，请总结一段 100 字以内的“核心推动原因”。\n 4. 返回 JSON: {{\"is_causal\": true/false, \"summary\": \"原因摘要\"}}\n \"\"\"\n \n try:\n res = self.llm_agent.run(prompt)\n data = json.loads(res.content.replace('```json', '').replace('```', '').strip())\n if data.get('is_causal'):\n logger.success(f\"✅ Verified cause for {name} on {date_str}: {data['summary']}\")\n return data['summary']\n else:\n logger.warning(f\"❌ Verified cause for {name} on {date_str}: {data['summary']}\")\n return None\n except Exception as e:\n logger.warning(f\"Verification failed: {e}\")\n return None\n\n def save_model(self, path=None):\n \"\"\"Save the news_proj weights\"\"\"\n if path is None:\n save_dir = os.path.join(SRC_DIR, \"exports/models\")\n os.makedirs(save_dir, exist_ok=True)\n path = os.path.join(save_dir, f\"kronos_news_v1_{datetime.now().strftime('%Y%m%d_%H%M')}.pt\")\n \n # We only really need to save the news_proj part as it's the only one we train\n torch.save({\n 'news_proj_state_dict': self.model.news_proj.state_dict(),\n 'news_dim': self.news_dim,\n 'd_model': self.model.d_model\n }, path)\n logger.success(f\"💾 Model weights saved to {path}\")\n return path\n\n def run_synthesis_and_train(self, tickers, pred_len=5):\n # 1. Discovery\n shocks = self.discover_shocks(tickers, pred_len=pred_len)\n print(f'find {len(shocks)} shocks')\n \n # 2. News Association & Verification\n dataset = []\n max_news_items = 200 # Limit to 200 news items per session to avoid search bans\n \n logger.info(f\"🧬 Starting News Association for {len(shocks)} shocks (Max limit: {max_news_items})\")\n \n for i, shock in enumerate(shocks):\n if len(dataset) >= max_news_items:\n logger.info(\"Reached maximum news items limit for this session.\")\n break\n \n summary = self.find_reason_and_verify(shock)\n if summary:\n # 3. Embedding news\n emb = self.embedder.encode(summary)\n dataset.append({\n 'history': shock['history'],\n 'target': shock['target'],\n 'news_emb': emb,\n 'summary': summary\n })\n \n # Add delay after search with randomness to avoid being blocked\n if i \u003c len(shocks) - 1:\n delay = random.uniform(2.0, 4.0)\n time.sleep(delay)\n \n if not dataset:\n logger.error(\"❌ No verified news-price pairs found. Adjust threshold or check if news is available in that period.\")\n return\n\n # 4. Train/Val Split\n random.seed(42)\n random.shuffle(dataset)\n \n if len(dataset) \u003c 2:\n train_set = dataset\n val_set = []\n logger.warning(f\"⚠️ Only {len(dataset)} sample(s) found. Training on all, skipping validation.\")\n else:\n split_idx = max(1, int(len(dataset) * 0.8))\n if split_idx >= len(dataset):\n split_idx = len(dataset) - 1\n \n train_set = dataset[:split_idx]\n val_set = dataset[split_idx:]\n logger.info(f\"🏗️ Dataset Split: {len(train_set)} samples for training, {len(val_set)} for validation.\")\n\n if not train_set:\n logger.error(\"❌ No samples for training.\")\n return\n\n # 5. Training (Few-shot)\n optimizer = torch.optim.Adam(self.model.news_proj.parameters(), lr=1e-3)\n criterion = nn.CrossEntropyLoss()\n self.model.train()\n \n loss_history = []\n logger.info(f\"🚀 Training for 30 epochs...\")\n for epoch in range(30):\n total_loss = 0\n for item in train_set:\n optimizer.zero_grad()\n \n # Prep Data\n hist_df = item['history']\n # For training, we still focus on the immediate next point (teacher forcing)\n target_df = item['target'].iloc[:1]\n \n hist_raw = hist_df[['open', 'high', 'low', 'close', 'volume']].values.astype(np.float32)\n hist_raw = np.column_stack([hist_raw, hist_raw[:, 3] * hist_raw[:, 4]]) \n \n mean, std = hist_raw.mean(axis=0), hist_raw.std(axis=0) + 1e-5\n hist_norm = torch.from_numpy((hist_raw - mean) / std).unsqueeze(0).to(self.device)\n \n target_raw = target_df[['open', 'high', 'low', 'close', 'volume']].values.astype(np.float32)\n target_raw = np.column_stack([target_raw, target_raw[:, 3] * target_raw[:, 4]])\n target_norm = torch.from_numpy((target_raw - mean) / std).unsqueeze(0).to(self.device)\n \n with torch.no_grad():\n z_indices = self.tokenizer.encode(hist_norm, half=True)\n t_indices = self.tokenizer.encode(target_norm, half=True)\n s1_ids, s2_ids = z_indices[0], z_indices[1]\n t_s1, t_s2 = t_indices[0], t_indices[1]\n \n news_t = torch.from_numpy(item['news_emb']).unsqueeze(0).to(self.device)\n s1_logits, s2_logits = self.model(s1_ids, s2_ids, news_emb=news_t, use_teacher_forcing=True, s1_targets=t_s1)\n \n loss = (criterion(s1_logits[:, -1, :], t_s1[:, 0]) + criterion(s2_logits[:, -1, :], t_s2[:, 0])) / 2\n loss.backward()\n optimizer.step()\n total_loss += loss.item()\n \n avg_epoch_loss = total_loss / max(1, len(train_set))\n loss_history.append(avg_epoch_loss)\n \n if (epoch + 1) % 10 == 0:\n logger.info(f\"Epoch {epoch+1} Loss: {avg_epoch_loss:.4f}\")\n\n # 5.1 Visualize Loss Curve\n loss_chart = VisualizerTools.generate_loss_chart(loss_history)\n VisualizerTools.render_chart_to_file(loss_chart, os.path.join(SRC_DIR, \"exports/training_results/loss_curve.html\"))\n\n # 5.2 Save final model\n self.save_model()\n\n # 6. Final Evaluation on Validation Set\n if not val_set:\n logger.warning(\"⚠️ Validation set is empty. Skipping statistical analysis.\")\n return\n\n logger.info(f\"🧪 Final Evaluation: Base vs News-Integrated ({pred_len}-day Window)\")\n self.model.eval()\n predictor = KronosPredictor(self.model, self.tokenizer, device=self.device)\n \n base_maes = []\n news_maes = []\n \n print(\"\\n\" + \"=\"*90)\n print(f\"{'Date':\u003c12} | {'Ticker':\u003c8} | {'Base MAE':\u003c15} | {'News MAE':\u003c15} | {'Improvement'}\")\n print(\"-\" * 90)\n\n for item in val_set:\n h = item['history']\n t = item['target']\n actuals = t['close'].values[:pred_len]\n \n x_ts = pd.to_datetime(h['date'])\n # Future timestamps: handle business days if possible, or just simple offset\n future_dates = pd.date_range(start=x_ts.iloc[-1] + timedelta(days=1), periods=pred_len, freq='B')\n y_ts = pd.Series(future_dates)\n \n # A. Base Prediction\n p_base = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=None, verbose=False)\n b_preds = p_base['close'].values[:len(actuals)]\n \n # B. News-Aware Prediction\n p_news = predictor.predict(h, x_ts, y_ts, pred_len=pred_len, news_emb=item['news_emb'], verbose=False)\n n_preds = p_news['close'].values[:len(actuals)]\n \n # Calculate MAE over the window\n b_mae = np.mean(np.abs(b_preds - actuals))\n n_mae = np.mean(np.abs(n_preds - actuals))\n \n base_maes.append(b_mae)\n news_maes.append(n_mae)\n \n improvement = (b_mae - n_mae) / (b_mae + 1e-6) * 100\n \n date_str = str(t['date'].values[0])[:10]\n ticker = h.iloc[-1]['ticker'] if 'ticker' in h.columns else \"Stock\"\n print(f\"{date_str:\u003c12} | {ticker:\u003c8} | {b_mae:\u003c15.4f} | {n_mae:\u003c15.4f} | {improvement:>+7.1f}%\")\n\n # C. Generate Visualization for this case\n try:\n # Helper to convert DF to KLinePoints\n def to_kp_list(preds_df):\n points = []\n for idx, row in preds_df.iterrows():\n points.append(KLinePoint(\n date=str(idx)[:10],\n open=row['open'],\n high=row['high'],\n low=row['low'],\n close=row['close'],\n volume=row['volume'] if 'volume' in row else 0\n ))\n return points\n\n forecast_obj = ForecastResult(\n ticker=ticker,\n base_forecast=to_kp_list(p_base),\n adjusted_forecast=to_kp_list(p_news),\n rationale=item['summary']\n )\n\n # Ground truth for visualizer expects a DataFrame with 'date' and 'close'\n gt_df = t[['date', 'open', 'high', 'low', 'close', 'volume']]\n \n chart = VisualizerTools.generate_stock_chart(\n df=h, \n ticker=ticker, \n title=f\"Training Eval: {ticker} ({date_str}) Improvement: {improvement:.1f}%\",\n forecast=forecast_obj,\n ground_truth=gt_df\n )\n \n safe_date = date_str.replace(\"-\", \"\")\n filename = f\"eval_{ticker}_{safe_date}.html\"\n VisualizerTools.render_chart_to_file(chart, os.path.join(SRC_DIR, f\"exports/training_results/{filename}\"))\n except Exception as e:\n logger.error(f\"Failed to generate eval chart for {ticker}: {e}\")\n\n # Summary Statistics\n avg_base_err = sum(base_maes) / max(1, len(base_maes))\n avg_news_err = sum(news_maes) / max(1, len(news_maes))\n overall_imp = (avg_base_err - avg_news_err) / (avg_base_err + 1e-6) * 100\n \n print(\"-\" * 90)\n print(f\"{'AVERAGE':\u003c12} | {'-':\u003c8} | {avg_base_err:\u003c15.4f} | {avg_news_err:\u003c15.4f} | {overall_imp:>+7.1f}%\")\n print(\"=\"*90 + \"\\n\")\n \n logger.success(f\"🏁 Statistical Analysis Complete. Avg Error Reduction ({pred_len}-day): {overall_imp:.2f}%\")\n logger.info(f\"📊 Visualization results saved to: {os.path.join(SRC_DIR, 'exports/training_results/')}\")\n\nif __name__ == \"__main__\":\n trainer = AutoSynthesisTrainer()\n \n logger.info(\"📂 Fetching all stock codes from database...\")\n res = trainer.db.execute_query(\"SELECT code FROM stock_list\")\n all_tickers = [row['code'] for row in res]\n \n if not all_tickers:\n logger.warning(\"⚠️ No tickers found in stock_list table. Trying to sync...\")\n trainer.tools._check_and_update_stock_list(force=True)\n res = trainer.db.execute_query(\"SELECT code FROM stock_list\")\n all_tickers = [row['code'] for row in res]\n\n logger.info(f\"🚀 Starting training on potential stocks (1-year scan)...\")\n # 为了演示，我们扫描前 100 个股票，寻找最近一年的冲击点\n trainer.run_synthesis_and_train(all_tickers[:100], pred_len=1)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":18696,"content_sha256":"8cbc820f3142abce635dc9b8baf680434a3947219476cc5d430ed799549905a9"},{"filename":"scripts/utils/search_tools.py","content":"import os\nimport hashlib\nimport json\nimport re\nimport requests\nimport time\nimport threading\nfrom typing import List, Dict, Optional, Any\nfrom agno.tools.duckduckgo import DuckDuckGoTools\nfrom agno.tools.baidusearch import BaiduSearchTools\nfrom agno.agent import Agent\nfrom loguru import logger\nfrom datetime import datetime\nfrom .database_manager import DatabaseManager\nfrom .content_extractor import ContentExtractor\nfrom .llm.factory import get_model\nfrom .hybrid_search import LocalNewsSearch\n\n# 默认搜索缓存 TTL（秒），可通过环境变量覆盖\nDEFAULT_SEARCH_TTL = int(os.getenv(\"SEARCH_CACHE_TTL\", \"3600\")) # 默认 1 小时\n\n\nclass JinaSearchEngine:\n \"\"\"Jina Search API 封装 - 使用 s.jina.ai 进行网络搜索\"\"\"\n \n JINA_SEARCH_URL = \"https://s.jina.ai/\"\n \n # 速率限制配置\n _rate_limit_no_key = 10 # 无 key 时每分钟最大请求数\n _rate_window = 60.0\n _min_interval = 2.0\n _request_times = []\n _last_request_time = 0.0\n _lock = threading.Lock()\n \n def __init__(self):\n self.api_key = os.getenv(\"JINA_API_KEY\", \"\").strip()\n self.has_api_key = bool(self.api_key)\n if self.has_api_key:\n logger.info(\"✅ Jina Search API key configured\")\n \n @classmethod\n def _wait_for_rate_limit(cls, has_api_key: bool) -> None:\n \"\"\"等待以满足速率限制\"\"\"\n if has_api_key:\n time.sleep(0.3)\n return\n \n with cls._lock:\n current_time = time.time()\n cls._request_times = [t for t in cls._request_times if current_time - t \u003c cls._rate_window]\n \n if len(cls._request_times) >= cls._rate_limit_no_key:\n oldest = cls._request_times[0]\n wait_time = cls._rate_window - (current_time - oldest) + 1.0\n if wait_time > 0:\n logger.warning(f\"⏳ Jina Search rate limit, waiting {wait_time:.1f}s...\")\n time.sleep(wait_time)\n current_time = time.time()\n cls._request_times = [t for t in cls._request_times if current_time - t \u003c cls._rate_window]\n \n time_since_last = current_time - cls._last_request_time\n if time_since_last \u003c cls._min_interval:\n time.sleep(cls._min_interval - time_since_last)\n \n cls._request_times.append(time.time())\n cls._last_request_time = time.time()\n \n def search(self, query: str, max_results: int = 5) -> List[Dict]:\n \"\"\"\n 使用 Jina Search API 执行搜索\n \n Args:\n query: 搜索关键词\n max_results: 返回结果数量\n \n Returns:\n 搜索结果列表，每个结果包含 title, url, content\n \"\"\"\n if not query:\n return []\n \n logger.info(f\"🔍 Jina Search: {query}\")\n \n # 等待速率限制\n self._wait_for_rate_limit(self.has_api_key)\n \n headers = {\n \"Accept\": \"application/json\",\n \"X-Retain-Images\": \"none\",\n }\n \n if self.has_api_key:\n headers[\"Authorization\"] = f\"Bearer {self.api_key}\"\n \n try:\n # Jina Search API: https://s.jina.ai/{query}\n import urllib.parse\n encoded_query = urllib.parse.quote(query)\n url = f\"{self.JINA_SEARCH_URL}{encoded_query}\"\n \n response = requests.get(url, headers=headers, timeout=30)\n \n if response.status_code == 429:\n logger.warning(\"⚠️ Jina Search rate limited (429), waiting 30s...\")\n time.sleep(30)\n return self.search(query, max_results)\n \n if response.status_code != 200:\n logger.warning(f\"Jina Search failed (Status {response.status_code})\")\n return []\n \n # 解析响应\n try:\n data = response.json()\n except json.JSONDecodeError:\n # 如果返回纯文本，尝试解析\n data = {\"data\": [{\"title\": \"Search Result\", \"url\": \"\", \"content\": response.text}]}\n \n results = []\n \n # Jina 返回格式可能是 {\"data\": [...]} 或直接是列表\n items = data.get(\"data\", []) if isinstance(data, dict) else data\n if not isinstance(items, list):\n items = [items] if items else []\n \n for i, item in enumerate(items[:max_results]):\n if isinstance(item, dict):\n results.append({\n \"title\": item.get(\"title\", f\"Result {i+1}\"),\n \"url\": item.get(\"url\", \"\"),\n \"href\": item.get(\"url\", \"\"), # 兼容性\n \"content\": item.get(\"content\", item.get(\"description\", \"\")),\n \"body\": item.get(\"content\", item.get(\"description\", \"\")), # 兼容性\n })\n elif isinstance(item, str):\n results.append({\n \"title\": f\"Result {i+1}\",\n \"url\": \"\",\n \"content\": item\n })\n \n logger.info(f\"✅ Jina Search returned {len(results)} results\")\n return results\n \n except requests.exceptions.Timeout:\n logger.error(\"Jina Search timeout\")\n return []\n except requests.exceptions.RequestException as e:\n logger.error(f\"Jina Search request error: {e}\")\n return []\n except Exception as e:\n logger.error(f\"Jina Search unexpected error: {e}\")\n return []\n\nclass SearchTools:\n \"\"\"扩展性搜索工具库 - 支持多引擎聚合与内容缓存\"\"\"\n \n def __init__(self, db: DatabaseManager):\n self.db = db\n \n # 检查 Jina API Key 是否配置\n jina_api_key = os.getenv(\"JINA_API_KEY\", \"\").strip()\n self._jina_enabled = bool(jina_api_key)\n \n self._engines = {\n \"ddg\": DuckDuckGoTools(),\n \"baidu\": BaiduSearchTools(),\n \"local\": LocalNewsSearch(db)\n }\n \n # 如果配置了 Jina API Key，添加 Jina 引擎\n if self._jina_enabled:\n self._engines[\"jina\"] = JinaSearchEngine()\n logger.info(\"🚀 Jina Search engine enabled (JINA_API_KEY configured)\")\n \n # 确定默认搜索引擎\n self._default_engine = \"jina\" if self._jina_enabled else \"ddg\"\n\n def _generate_hash(self, query: str, engine: str, max_results: int) -> str:\n return hashlib.md5(f\"{engine}:{query}:{max_results}\".encode()).hexdigest()\n\n def search(self, query: str, engine: str = None, max_results: int = 5, ttl: Optional[int] = None) -> str:\n \"\"\"\n 使用指定搜索引擎执行网络搜索，结果会被缓存以提高效率。\n \n Args:\n query: 搜索关键词，如 \"英伟达财报\" 或 \"光伏行业政策\"。\n engine: 搜索引擎选择。可选值: \n \"jina\" (Jina Search，需配置 JINA_API_KEY，LLM友好输出),\n \"ddg\" (DuckDuckGo，推荐英文/国际搜索), \n \"baidu\" (百度，推荐中文/国内搜索),\n \"local\" (本地历史新闻搜索，基于向量+BM25)。\n 默认: 若配置了 JINA_API_KEY 则使用 \"jina\"，否则 \"ddg\"。\n max_results: 期望返回的结果数量，默认 5 条。\n ttl: 缓存有效期（秒）。如果缓存超过此时间会重新搜索。\n 默认使用环境变量 SEARCH_CACHE_TTL 或 3600 秒。\n 设为 0 可强制刷新。\n \n Returns:\n 搜索结果的文本描述，包含标题、摘要和链接。\n \"\"\"\n # 使用默认引擎（如果配置了 Jina 则优先使用 Jina）\n if engine is None:\n engine = self._default_engine\n \n if engine not in self._engines:\n return f\"Error: Unsupported engine '{engine}'. Available: {list(self._engines.keys())}\"\n\n query_hash = self._generate_hash(query, engine, max_results)\n effective_ttl = ttl if ttl is not None else DEFAULT_SEARCH_TTL\n \n # 1. 尝试从缓存读取 (local 引擎不缓存，因为它本身就是查库)\n if engine != \"local\":\n cache = self.db.get_search_cache(query_hash, ttl_seconds=effective_ttl if effective_ttl > 0 else None)\n if cache and effective_ttl != 0:\n logger.info(f\"ℹ️ Found search results in cache for: {query} ({engine})\")\n return cache['results']\n\n # 2. 执行真实搜索\n logger.info(f\"📡 Searching {engine} for: {query}\")\n try:\n tool = self._engines[engine]\n if engine == \"jina\":\n # Jina Search 返回 List[Dict]\n jina_results = tool.search(query, max_results=max_results)\n results = []\n for r in jina_results:\n results.append({\n \"title\": r.get(\"title\", \"\"),\n \"href\": r.get(\"url\", \"\"),\n \"body\": r.get(\"content\", \"\")\n })\n elif engine == \"ddg\":\n results = tool.duckduckgo_search(query, max_results=max_results)\n elif engine == \"baidu\":\n results = tool.baidu_search(query, max_results=max_results)\n elif engine == \"local\":\n # LocalNewsSearch 返回的是 List[Dict]\n local_results = tool.search(query, top_n=max_results)\n results = []\n for r in local_results:\n results.append({\n \"title\": r.get(\"title\"),\n \"href\": r.get(\"url\", \"local\"),\n \"body\": r.get(\"content\", \"\")\n })\n else:\n results = \"Search not implemented for this engine.\"\n \n results_str = str(results)\n if engine != \"local\":\n self.db.save_search_cache(query_hash, query, engine, results_str)\n return results_str\n \n except Exception as e:\n # 搜索失败时的降级策略\n if engine == \"jina\":\n logger.warning(f\"⚠️ Jina search failed, falling back to ddg: {query} ({e})\")\n try:\n return self.search(query, engine=\"ddg\", max_results=max_results, ttl=ttl)\n except Exception as e2:\n logger.error(f\"❌ DDG fallback also failed for {query}: {e2}\")\n elif engine == \"ddg\":\n logger.warning(f\"⚠️ DDG search failed, falling back to baidu: {query} ({e})\")\n try:\n return self.search(query, engine=\"baidu\", max_results=max_results, ttl=ttl)\n except Exception as e2:\n logger.error(f\"❌ Baidu fallback also failed for {query}: {e2}\")\n\n logger.error(f\"❌ Search failed for {query}: {e}\")\n return f\"Error occurred during search: {str(e)}\"\n\n def search_list(self, query: str, engine: str = None, max_results: int = 5, ttl: Optional[int] = None, enrich: bool = True) -> List[Dict]:\n \"\"\"\n 执行搜索并返回结构化列表 (List[Dict])。\n Dict 包含: title, href (or url), body (or snippet)\n \n Args:\n engine: 搜索引擎，默认使用配置的默认引擎（Jina 优先）\n enrich: 是否抓取正文内容 (默认 True)\n \"\"\"\n # 使用默认引擎\n if engine is None:\n engine = self._default_engine\n \n if engine not in self._engines:\n logger.error(f\"Unsupported engine {engine}\")\n return []\n \n # 不同的 hash 以区分是否 enrichment\n enrich_suffix = \":enriched\" if enrich else \"\"\n query_hash = self._generate_hash(query, engine + enrich_suffix, max_results)\n effective_ttl = ttl if ttl is not None else DEFAULT_SEARCH_TTL\n \n # 1. 尝试从缓存读取\n cache = self.db.get_search_cache(query_hash, ttl_seconds=effective_ttl if effective_ttl > 0 else None)\n if cache and effective_ttl != 0:\n try:\n cached_data = json.loads(cache['results'])\n if isinstance(cached_data, list):\n logger.info(f\"ℹ️ Found structured search cache for: {query}\")\n return cached_data\n except:\n pass\n \n # 1.5 Smart Cache (Fuzzy + LLM)\n if effective_ttl != 0:\n try:\n # 1. Similar cached queries\n similar_queries = self.db.find_similar_queries(query, limit=3)\n # Filter by TTL\n valid_candidates = []\n for q in similar_queries:\n if q['query'] == query: continue \n q_time = datetime.fromisoformat(q['timestamp'])\n if effective_ttl and (datetime.now() - q_time).total_seconds() > effective_ttl:\n continue\n q['type'] = 'cached_search'\n valid_candidates.append(q)\n\n # 2. Relevant local news (as search results)\n local_news = self.db.search_local_news(query, limit=3)\n if local_news:\n # Group local news as a single \"candidate\" source? Or individual?\n # Better to treat \"Local News Database\" as one candidate source that contains X items.\n # Or just add them to candidates list?\n # Let's package strictly relevant news as a \"local_news_bundle\"\n valid_candidates.append({\n 'type': 'local_news',\n 'query': 'Local Database News',\n 'items': local_news,\n 'timestamp': datetime.now().isoformat()\n })\n \n if valid_candidates:\n logger.info(f\"🤔 Found {len(valid_candidates)} smart cache candidates (Queries/News). Asking LLM...\")\n evaluation = self._evaluate_cache_relevance(query, valid_candidates)\n \n if evaluation and evaluation.get('reuse', False):\n idx = evaluation.get('index', -1)\n if 0 \u003c= idx \u003c len(valid_candidates):\n chosen = valid_candidates[idx]\n logger.info(f\"🤖 LLM suggested reusing: '{chosen.get('query')}' ({chosen['type']})\")\n \n if chosen['type'] == 'cached_search':\n # Load the chosen cache\n cache = self.db.get_search_cache(chosen['query_hash']) \n if cache:\n try:\n cached_data = json.loads(cache['results'])\n if isinstance(cached_data, list):\n return cached_data\n except:\n pass\n elif chosen['type'] == 'local_news':\n # Convert local news items to search result format\n news_results = []\n for i, news in enumerate(chosen['items'], 1):\n news_results.append({\n \"id\": news.get('id'),\n \"rank\": i,\n \"title\": news.get('title'),\n \"url\": news.get('url'),\n \"content\": news.get('content'),\n \"original_snippet\": news.get('content')[:200] if news.get('content') else '',\n \"source\": f\"Local News ({news.get('source')})\",\n \"publish_time\": news.get('publish_time'),\n \"crawl_time\": news.get('crawl_time'),\n \"sentiment_score\": news.get('sentiment_score', 0),\n \"meta_data\": {\"origin\": \"local_db\"}\n })\n return news_results\n\n except Exception as e:\n logger.warning(f\"Smart cache check failed: {e}\")\n \n # 2. 执行搜索\n logger.info(f\"📡 Searching {engine} (structured) for: {query}\")\n try:\n tool = self._engines[engine]\n results = []\n if engine == \"jina\":\n # Jina Search 直接返回结构化数据\n jina_results = tool.search(query, max_results=max_results)\n for r in jina_results:\n results.append({\n \"title\": r.get(\"title\", \"\"),\n \"url\": r.get(\"url\", \"\"),\n \"href\": r.get(\"url\", \"\"),\n \"body\": r.get(\"content\", \"\"),\n \"content\": r.get(\"content\", \"\"),\n \"source\": \"Jina Search\"\n })\n elif engine == \"ddg\":\n results = tool.duckduckgo_search(query, max_results=max_results)\n elif engine == \"baidu\":\n results = tool.baidu_search(query, max_results=max_results)\n elif engine == \"local\":\n # LocalNewsSearch 返回的是 List[Dict]\n local_results = tool.search(query, top_n=max_results)\n results = []\n for r in local_results:\n results.append({\n \"title\": r.get(\"title\"),\n \"url\": r.get(\"url\", \"local\"),\n \"body\": r.get(\"content\", \"\")[:500],\n \"source\": f\"Local ({r.get('source', 'db')})\",\n \"publish_time\": r.get(\"publish_time\")\n })\n \n # 处理字符串类型的 JSON 返回 (Baidu 常返 JSON 字符串)\n if isinstance(results, str) and engine not in [\"local\", \"jina\"]:\n try:\n results = json.loads(results)\n except:\n pass\n \n # 转为统一格式\n normalized_results = []\n if isinstance(results, list):\n \n for i, r in enumerate(results, 1):\n title = r.get('title', '')\n url = r.get('href') or r.get('url') or r.get('link', '')\n content = r.get('body') or r.get('snippet') or r.get('abstract', '')\n \n if title and url:\n normalized_results.append({\n \"id\": self._generate_hash(url + query, \"search_item\", i),\n \"rank\": i,\n \"title\": title,\n \"url\": url,\n \"content\": content,\n \"original_snippet\": content, # 保留摘要\n \"source\": f\"Search ({engine})\",\n \"publish_time\": datetime.now().isoformat(), # 暂用当前时间\n \"crawl_time\": datetime.now().isoformat(),\n \"meta_data\": {\"query\": query, \"engine\": engine}\n })\n \n # Fallback if still string and failed to parse\n elif isinstance(results, str) and results:\n normalized_results.append({\"title\": query, \"url\": \"\", \"content\": results, \"source\": engine})\n\n # 3. 抓取正文 & 计算情绪 (Enrichment)\n # 注意：如果使用 Jina Search，内容已经是 LLM 友好格式，可选择跳过 enrichment\n skip_content_enrichment = (engine == \"jina\")\n \n if enrich and normalized_results:\n logger.info(f\"🕸️ Enriching {len(normalized_results)} search results with Jina & Sentiment...\")\n extractor = ContentExtractor()\n \n # Lazy load sentiment tool\n if not hasattr(self, 'sentiment_tool') or self.sentiment_tool is None:\n from ..sentiment_tools import SentimentTools\n self.sentiment_tool = SentimentTools(self.db)\n \n for item in normalized_results:\n if item.get(\"url\"):\n try:\n # 如果是 Jina Search，内容已经足够好，跳过额外抓取\n if skip_content_enrichment and item.get(\"content\") and len(item.get(\"content\", \"\")) > 100:\n full_content = item[\"content\"]\n else:\n # Use Jina Reader to get full content\n full_content = extractor.extract_with_jina(item[\"url\"], timeout=60)\n \n if full_content and len(full_content) > 100:\n item[\"content\"] = full_content\n \n # Calculate sentiment\n # Use title + snippet of content for efficiency\n text_to_analyze = f\"{item['title']} {full_content[:500]}\"\n sent_result = self.sentiment_tool.analyze_sentiment(text_to_analyze) # Using self.sentiment_tool\n score = sent_result.get('score', 0.0)\n item[\"sentiment_score\"] = float(score)\n \n logger.info(f\" ✅ Enriched: {item['title'][:20]}... (Sentiment: {score:.2f})\")\n else:\n # Fallback: Use snippet for sentiment\n logger.info(f\" ⚠️ Content short/failed for {item['url']}, using snippet for sentiment.\")\n text_to_analyze = f\"{item['title']} {item['content']}\" # content is snippet here\n sent_result = self.sentiment_tool.analyze_sentiment(text_to_analyze)\n score = sent_result.get('score', 0.0)\n item[\"sentiment_score\"] = float(score)\n\n except Exception as e:\n # Fallback: Use snippet for sentiment on error\n logger.warning(f\"Failed to enrich {item['url']}: {e}. Using snippet.\")\n text_to_analyze = f\"{item['title']} {item['content']}\"\n sent_result = self.sentiment_tool.analyze_sentiment(text_to_analyze)\n score = sent_result.get('score', 0.0)\n item[\"sentiment_score\"] = float(score)\n \n # 缓存结果 list\n if normalized_results:\n # Pass list directly, DB manager will handle JSON dump for main cache and populate search_details\n # Only cache if NOT from local news reuse (though this logic path is for fresh search)\n self.db.save_search_cache(query_hash, query, engine, normalized_results)\n \n return normalized_results\n \n except Exception as e:\n # 搜索失败时的降级策略\n if engine == \"jina\":\n logger.warning(f\"⚠️ Jina search_list failed, falling back to ddg: {query} ({e})\")\n try:\n return self.search_list(query, engine=\"ddg\", max_results=max_results, ttl=ttl, enrich=enrich)\n except Exception as e2:\n logger.error(f\"❌ DDG fallback (search_list) also failed for {query}: {e2}\")\n elif engine == \"ddg\":\n logger.warning(f\"⚠️ DDG search_list failed, falling back to baidu: {query} ({e})\")\n try:\n return self.search_list(query, engine=\"baidu\", max_results=max_results, ttl=ttl, enrich=enrich)\n except Exception as e2:\n logger.error(f\"❌ Baidu fallback (search_list) also failed for {query}: {e2}\")\n\n logger.error(f\"❌ Structured search failed for {query}: {e}\")\n return []\n\n def _evaluate_cache_relevance(self, current_query: str, candidates: List[Dict]) -> Dict:\n \"\"\"\n 使用 LLM 评估缓存候选是否足以回答当前问题。\n \"\"\"\n try:\n # Prepare candidates text\n candidates_desc = []\n for i, c in enumerate(candidates):\n if c['type'] == 'cached_search':\n # Preview cached results if available? \n # Maybe just use the query string as a proxy for what's in there.\n # Or peek at 'results' snippet.\n preview = \"\"\n try:\n # Attempt to peek first result title from JSON string\n # Note: c.get('results') might be a stringified JSON list\n res_list = json.loads(c.get('results', '[]'))\n if res_list and isinstance(res_list, list) and len(res_list) > 0:\n first_item = res_list[0]\n if isinstance(first_item, dict) and 'title' in first_item:\n preview = f\" (Contains: {first_item.get('title', '')[:50]}...)\"\n except:\n pass\n candidates_desc.append(f\"[{i}] Old Search Query: '{c['query']}' {preview} (Time: {c['timestamp']})\")\n elif c['type'] == 'local_news':\n # List titles of local news\n titles = [item['title'] for item in c['items'][:3]]\n candidates_desc.append(f\"[{i}] Local Database News: {', '.join(titles)}... (Time: {c['timestamp']})\")\n\n prompt = f\"\"\"\n Task: Decide if existing information is sufficient for the new search query.\n \n New Query: \"{current_query}\"\n \n Available Information Candidates:\n {chr(10).join(candidates_desc)}\n \n Instructions:\n 1. Analyze if any candidate provides ENOUGH up-to-date info for the \"New Query\".\n 2. If yes, choose the best one.\n 3. If the query implies needing LATEST real-time info and candidates are old, choose none.\n 4. Return strictly JSON: {{\"reuse\": true/false, \"index\": \u003ccandidate_index_int>, \"reason\": \"short explanation\"}}\n \"\"\"\n # 初始化模型\n provider = os.getenv(\"LLM_PROVIDER\", \"ust\")\n model_id = os.getenv(\"LLM_MODEL\", \"Qwen\")\n host = os.getenv(\"LLM_HOST\")\n if host:\n model = get_model(provider, model_id, host=host)\n else:\n model = get_model(provider, model_id)\n \n agent = Agent(model=model, markdown=True)\n \n response = agent.run(prompt)\n content = response.content\n \n # Parse JSON\n json_match = re.search(r'```json\\s*(.*?)\\s*```', content, re.DOTALL)\n if json_match:\n return json.loads(json_match.group(1))\n elif '{' in content:\n # Fallback for cases where LLM doesn't wrap in ```json\n return json.loads(content[content.find('{'):content.rfind('}')+1])\n return {\"reuse\": False}\n \n except Exception as e:\n logger.warning(f\"LLM evaluation failed: {e}\")\n return {\"reuse\": False}\n\n def aggregate_search(self, query: str, engines: Optional[List[str]] = None, max_results: int = 5) -> str:\n \"\"\"\n 使用多个搜索引擎同时搜索并聚合结果，获得更全面的信息覆盖。\n \n Args:\n query: 搜索关键词。\n engines: 要使用的搜索引擎列表。可选值: [\"ddg\", \"baidu\"]。\n 默认同时使用 ddg 和 baidu。\n max_results: 每个引擎期望返回的结果数量。\n \n Returns:\n 聚合后的搜索结果，按引擎分组显示。\n \"\"\"\n engines = engines or [\"ddg\", \"baidu\"]\n aggregated_results = []\n for engine in engines:\n res = self.search(query, engine=engine, max_results=max_results)\n aggregated_results.append(f\"--- Results from {engine.upper()} ---\\n{res}\")\n \n return \"\\n\\n\".join(aggregated_results)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":29258,"content_sha256":"839c578de49ee139f9c7a4d36cdcafddc833ef407f73fdb0e38e0c02abac6ad9"},{"filename":"scripts/utils/stock_tools.py","content":"from datetime import datetime, timedelta\nfrom typing import List, Dict, Optional\nimport akshare as ak\nimport pandas as pd\nimport re\nimport sqlite3\nfrom requests.exceptions import RequestException\nfrom loguru import logger\nfrom .database_manager import DatabaseManager\nimport os\nfrom contextlib import contextmanager\n\n@contextmanager\ndef temporary_no_proxy():\n \"\"\"Context manager to temporarily unset proxy environment variables.\"\"\"\n proxies = {k: os.environ.get(k) for k in ['http_proxy', 'https_proxy', 'HTTP_PROXY', 'HTTPS_PROXY']}\n for k in proxies:\n if k in os.environ:\n del os.environ[k]\n try:\n yield\n finally:\n for k, v in proxies.items():\n if v is not None:\n os.environ[k] = v\n\nclass StockTools:\n \"\"\"金融分析股票工具 - 结合高性能数据库缓存与增量更新\"\"\"\n \n def __init__(self, db: DatabaseManager, auto_update: bool = True):\n \"\"\"\n 初始化股票工具\n \n Args:\n db: 数据库管理器\n auto_update: 是否在列表为空时自动更新，默认 True\n \"\"\"\n self.db = db\n if auto_update:\n self._check_and_update_stock_list()\n\n def _check_and_update_stock_list(self, force: bool = False):\n \"\"\"检查并更新股票列表。仅在列表为空或 force=True 时从网络拉取。\"\"\"\n # 直接查询表中记录数\n cursor = self.db.conn.cursor()\n cursor.execute(\"SELECT COUNT(*) FROM stock_list\")\n count = cursor.fetchone()[0]\n \n if count > 0 and not force:\n logger.info(f\"ℹ️ Stock list already cached ({count} stocks)\")\n return\n \n logger.info(\"📡 Updating A-share and HK-share stock list from akshare...\")\n \n def fetch_data():\n # A-share\n df_a = ak.stock_zh_a_spot_em()\n df_a = df_a[['代码', '名称']].copy()\n df_a.columns = ['code', 'name']\n \n # HK-share\n df_hk = ak.stock_hk_spot_em()\n df_hk = df_hk[['代码', '名称']].copy()\n df_hk.columns = ['code', 'name']\n \n # Combine\n return pd.concat([df_a, df_hk], ignore_index=True)\n\n try:\n try:\n df_combined = fetch_data()\n except (RequestException, Exception) as e:\n if \"Proxy\" in str(e) or \"proxy\" in str(e):\n logger.warning(f\"⚠️ Proxy error detected: {e}. Retrying with proxy disabled...\")\n with temporary_no_proxy():\n df_combined = fetch_data()\n else:\n raise e\n \n self.db.save_stock_list(df_combined)\n logger.info(f\"✅ Cached {len(df_combined)} stocks (A-share + HK) to database.\")\n \n except Exception as e:\n logger.error(f\"❌ Failed to sync stock list: {e}\")\n\n\n def search_ticker(self, query: str, limit: int = 5) -> List[Dict]:\n \"\"\"\n 模糊搜索 A 股股票代码或名称，支持常见缩写。\n \"\"\"\n # 清洗后缀 (如 CATL.SZ -> CATL, 000001.SZ -> 000001)\n clean_query = re.sub(r'\\.(SZ|SH|HK|US)

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.

, '', query, flags=re.IGNORECASE)\n \n # 常见缩写映射\n aliases = {\n \"CATL\": \"宁德时代\",\n \"BYD\": \"比亚迪\",\n \"TSLA\": \"特斯拉\",\n \"Moutai\": \"贵州茅台\",\n \"Tencent\": \"腾讯\",\n \"Alibaba\": \"阿里巴巴\",\n \"Meituan\": \"美团\",\n }\n \n search_query = aliases.get(clean_query.upper(), clean_query)\n \n # Robustness: if regex-like ticker code is embedded in query (e.g. \"300364 中文在线\"), try to extract it\n if not search_query.isdigit():\n # Extract explicit 5-6 digit codes\n match = re.search(r'\\b(\\d{5,6})\\b', clean_query)\n if match:\n search_query = match.group(1)\n\n return self.db.search_stock(search_query, limit)\n\n def get_stock_price(\n self,\n ticker: str,\n start_date: Optional[str] = None,\n end_date: Optional[str] = None,\n force_sync: bool = False,\n ) -> pd.DataFrame:\n \"\"\"\n 获取指定股票的历史价格数据。优先从本地缓存读取，缺失时自动从网络补齐。\n \n Args:\n ticker: 股票代码，如 \"600519\"（贵州茅台）或 \"000001\"（平安银行）。\n start_date: 开始日期，格式 \"YYYY-MM-DD\"。默认为 90 天前。\n end_date: 结束日期，格式 \"YYYY-MM-DD\"。默认为今天。\n \n Returns:\n 包含 date, open, close, high, low, volume, change_pct 列的 DataFrame。\n \"\"\"\n now = datetime.now()\n if not end_date:\n end_date = now.strftime('%Y-%m-%d')\n if not start_date:\n start_date = (now - timedelta(days=90)).strftime('%Y-%m-%d')\n\n df_db = self.db.get_stock_prices(ticker, start_date, end_date)\n \n need_update = False\n if df_db.empty:\n need_update = True\n else:\n db_latest = pd.to_datetime(df_db['date'].max())\n req_latest = pd.to_datetime(end_date)\n if (req_latest - db_latest).days > 2:\n need_update = True\n\n if force_sync:\n need_update = True\n\n if need_update:\n logger.info(f\"📡 Data stale or missing for {ticker}, syncing from network...\")\n \n # 清洗 ticker，确保只包含数字（Akshare A 股接口通常只需要数字代码）\n clean_ticker = \"\".join(filter(str.isdigit, ticker))\n if not clean_ticker:\n # Non A/H numeric tickers are not supported by the current data source.\n logger.warning(f\"⚠️ Unsupported ticker format (A/H only): {ticker}\")\n return df_db\n\n try:\n s_fmt = start_date.replace(\"-\", \"\")\n e_fmt = end_date.replace(\"-\", \"\")\n \n df_remote = None\n \n def fetch_data():\n if len(clean_ticker) == 5:\n # HK Stock\n return ak.stock_hk_hist(\n symbol=clean_ticker, period=\"daily\",\n start_date=s_fmt, end_date=e_fmt,\n adjust=\"qfq\"\n )\n else:\n # A-share Stock\n return ak.stock_zh_a_hist(\n symbol=clean_ticker, period=\"daily\",\n start_date=s_fmt, end_date=e_fmt,\n adjust=\"qfq\"\n )\n\n try:\n df_remote = fetch_data()\n except (RequestException, Exception) as e:\n if \"Proxy\" in str(e) or \"proxy\" in str(e):\n logger.warning(f\"⚠️ Proxy error detected: {e}. Retrying with proxy disabled...\")\n with temporary_no_proxy():\n df_remote = fetch_data()\n else:\n raise e\n \n if df_remote is not None and not df_remote.empty:\n df_remote = df_remote.rename(columns={\n '日期': 'date', '开盘': 'open', '收盘': 'close',\n '最高': 'high', '最低': 'low', '成交量': 'volume',\n '涨跌幅': 'change_pct'\n })\n # 确保日期格式正确\n df_remote['date'] = pd.to_datetime(df_remote['date']).dt.strftime('%Y-%m-%d')\n \n # 只有在获取到有意义的数据时才保存\n self.db.save_stock_prices(clean_ticker, df_remote) # 保存时使用清洗后的 clean_ticker\n \n # 重新查询数据库返回结果，保证一致性\n return self.db.get_stock_prices(clean_ticker, start_date, end_date)\n else:\n logger.warning(f\"⚠️ Akshare returned empty data for {clean_ticker}\")\n \n except KeyError as e:\n # Akshare 有时在某些股票无数据时会抛出 KeyError\n logger.warning(f\"⚠️ Akshare data missing for {clean_ticker}: {e}\")\n except (RequestException, ConnectionError) as e:\n logger.error(f\"❌ Network error during Akshare sync for {clean_ticker}: {e}\")\n except sqlite3.Error as e:\n logger.error(f\"❌ Database error during Akshare sync for {clean_ticker}: {e}\")\n except Exception as e:\n logger.error(f\"❌ Unexpected error during Akshare sync for {clean_ticker}: {e}\")\n \n return df_db\n\n\ndef get_stock_analysis(ticker: str, db: DatabaseManager) -> str:\n \"\"\"\n 生成指定股票的分析摘要报告。\n \n Args:\n ticker: 股票代码\n db: 数据库管理器实例\n \n Returns:\n Markdown 格式的分析报告，包含价格走势和关键指标。\n \"\"\"\n tools = StockTools(db)\n df = tools.get_stock_price(ticker)\n \n if df.empty:\n return f\"❌ 未能获取 {ticker} 的股价数据。\"\n \n latest = df.iloc[-1]\n change = ((latest['close'] - df.iloc[0]['close']) / df.iloc[0]['close']) * 100\n \n report = [\n f\"## 📊 {ticker} 分析报告\",\n f\"- **查询时段**: {df.iloc[0]['date']} -> {latest['date']}\",\n f\"- **当前价**: ¥{latest['close']:.2f}\",\n f\"- **时段涨跌**: {change:+.2f}%\",\n f\"- **最高/最低**: ¥{df['high'].max():.2f} / ¥{df['low'].min():.2f}\",\n \"\\n### 最近交易概览\",\n \"```\",\n df.tail(5)[['date', 'close', 'change_pct', 'volume']].to_string(index=False),\n \"```\"\n ]\n return \"\\n\".join(report)\n","content_type":"text/x-python; charset=utf-8","language":"python","size":10072,"content_sha256":"8db165681a6faf28a55767ed2409869cd20467a7b01f9fdb462a9bfa84b85c95"},{"filename":"tests/test_predictor.py","content":"import sys\nimport os\nimport unittest\n\n# Add skill root to path\nsys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))\n\ntry:\n from scripts.kronos_predictor import KronosPredictorUtility\n from scripts.utils.database_manager import DatabaseManager\nexcept ImportError as e:\n print(f\"Import Error: {e}\")\n sys.exit(1)\n\nclass TestPredictor(unittest.TestCase):\n def test_init(self):\n print(\"Testing KronosPredictorUtility Iteration...\")\n db = DatabaseManager(\":memory:\")\n # Kronos might need model files, but init should pass if we don't call predict?\n # Note: Kronos loads model in init. This might fail if model path is invalid.\n # We wrap in try-except to catch model loading errors which are expected in this env\n try:\n tools = KronosPredictorUtility()\n self.assertIsNotNone(tools)\n except Exception as e:\n print(f\"Kronos Init failed (expected if no model): {e}\")\n\nif __name__ == '__main__':\n unittest.main()\n","content_type":"text/x-python; charset=utf-8","language":"python","size":1030,"content_sha256":"b7293345a8fb9a0b04f693bd709f838155a96634202ee3da8c049a1b7b835242"}],"content_json":{"type":"doc","content":[{"type":"heading","attrs":{"level":1},"content":[{"text":"AlphaEar Predictor Skill","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Overview","type":"text"}]},{"type":"paragraph","content":[{"text":"This skill utilizes the Kronos model (via ","type":"text"},{"text":"KronosPredictorUtility","type":"text","marks":[{"type":"code_inline"}]},{"text":") to perform time-series forecasting and adjust predictions based on news sentiment.","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Capabilities","type":"text"}]},{"type":"heading","attrs":{"level":3},"content":[{"text":"1. Forecast Market Trends","type":"text"}]},{"type":"heading","attrs":{"level":3},"content":[{"text":"1. Forecast Market Trends","type":"text"}]},{"type":"paragraph","content":[{"text":"Workflow:","type":"text","marks":[{"type":"strong"}]}]},{"type":"ordered_list","attrs":{"order":1,"listStyle":"number"},"content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Generate Base Forecast","type":"text","marks":[{"type":"strong"}]},{"text":": Use ","type":"text"},{"text":"scripts/kronos_predictor.py","type":"text","marks":[{"type":"code_inline"}]},{"text":" (via ","type":"text"},{"text":"KronosPredictorUtility","type":"text","marks":[{"type":"code_inline"}]},{"text":") to generate the technical/quantitative forecast.","type":"text"}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Adjust Forecast (Agentic)","type":"text","marks":[{"type":"strong"}]},{"text":": Use the ","type":"text"},{"text":"Forecast Adjustment Prompt","type":"text","marks":[{"type":"strong"}]},{"text":" in ","type":"text"},{"text":"references/PROMPTS.md","type":"text","marks":[{"type":"code_inline"}]},{"text":" to subjectively adjust the numbers based on latest news/logic.","type":"text"}]}]}]},{"type":"paragraph","content":[{"text":"Key Tools:","type":"text","marks":[{"type":"strong"}]}]},{"type":"bullet_list","content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"KronosPredictorUtility.get_base_forecast(df, lookback, pred_len, news_text)","type":"text","marks":[{"type":"code_inline"}]},{"text":": Returns ","type":"text"},{"text":"List[KLinePoint]","type":"text","marks":[{"type":"code_inline"}]},{"text":".","type":"text"}]}]}]},{"type":"paragraph","content":[{"text":"Example Usage (Python):","type":"text","marks":[{"type":"strong"}]}]},{"type":"code_block","attrs":{"wrap":false,"language":"python"},"content":[{"text":"from scripts.utils.kronos_predictor import KronosPredictorUtility\nfrom scripts.utils.database_manager import DatabaseManager\n\ndb = DatabaseManager()\npredictor = KronosPredictorUtility()\n\n# Forecast\nforecast = predictor.predict(\"600519\", horizon=\"7d\")\nprint(forecast)","type":"text"}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Configuration","type":"text"}]},{"type":"paragraph","content":[{"text":"This skill requires the ","type":"text"},{"text":"Kronos","type":"text","marks":[{"type":"strong"}]},{"text":" model and an embedding model.","type":"text"}]},{"type":"ordered_list","attrs":{"order":1,"listStyle":"number"},"content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Kronos Model","type":"text","marks":[{"type":"strong"}]},{"text":":","type":"text"}]},{"type":"bullet_list","content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Ensure ","type":"text"},{"text":"exports/models","type":"text","marks":[{"type":"code_inline"}]},{"text":" directory exists in the project root.","type":"text"}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Place trained news projector weights (e.g., ","type":"text"},{"text":"kronos_news_v1.pt","type":"text","marks":[{"type":"code_inline"}]},{"text":") in ","type":"text"},{"text":"exports/models/","type":"text","marks":[{"type":"code_inline"}]},{"text":".","type":"text"}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Or depend on the base model (automatically downloaded).","type":"text"}]}]}]}]}]},{"type":"blockquote","content":[{"type":"paragraph","content":[{"text":"[!CAUTION] ","type":"text"},{"text":"Model Security","type":"text","marks":[{"type":"strong"}]},{"text":": This skill loads model weights from ","type":"text"},{"text":"exports/models","type":"text","marks":[{"type":"code_inline"}]},{"text":". We use ","type":"text"},{"text":"weights_only=True","type":"text","marks":[{"type":"code_inline"}]},{"text":" and only scan for the ","type":"text"},{"text":"kronos_news_*.pt","type":"text","marks":[{"type":"code_inline"}]},{"text":" pattern. Ensure you only place trusted checkpoints in this directory.","type":"text"}]}]},{"type":"ordered_list","attrs":{"order":2,"listStyle":"number"},"content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"Environment Variables","type":"text","marks":[{"type":"strong"}]},{"text":":","type":"text"}]},{"type":"bullet_list","content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"EMBEDDING_MODEL","type":"text","marks":[{"type":"code_inline"}]},{"text":": Path or name of the embedding model (default: ","type":"text"},{"text":"sentence-transformers/all-MiniLM-L6-v2","type":"text","marks":[{"type":"code_inline"}]},{"text":").","type":"text"}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"KRONOS_MODEL_PATH","type":"text","marks":[{"type":"code_inline"}]},{"text":": Optional path to override model loading.","type":"text"}]}]}]}]}]},{"type":"heading","attrs":{"level":2},"content":[{"text":"Dependencies","type":"text"}]},{"type":"bullet_list","content":[{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"torch","type":"text","marks":[{"type":"code_inline"}]}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"transformers","type":"text","marks":[{"type":"code_inline"}]}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"sentence-transformers","type":"text","marks":[{"type":"code_inline"}]}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"pandas","type":"text","marks":[{"type":"code_inline"}]}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"numpy","type":"text","marks":[{"type":"code_inline"}]}]}]},{"type":"list_item","content":[{"type":"paragraph","content":[{"text":"scikit-learn","type":"text","marks":[{"type":"code_inline"}]}]}]}]}]},"metadata":{"date":"2026-11-02","name":"alphaear-predictor","author":"@skillopedia","source":{"stars":3096,"repo_name":"Awesome-finance-skills","origin_url":"https://github.com/RKiding/Awesome-finance-skills/blob/HEAD/skills/alphaear-predictor/SKILL.md","repo_owner":"RKiding","body_sha256":"cbba615e281bee983651a8cc5c2e1f30c961156b767ce661aa29adbd74908953","cluster_key":"a396a731a730e09245c2741e023b9e22075b3fdbfc0738d0ca4e3c08bca346be","clean_bundle":{"format":"clean-skill-bundle-v1","source":"RKiding/Awesome-finance-skills/skills/alphaear-predictor/SKILL.md","text_bytes":299116,"source_commit":"853f09b4d0baae747759ed31e21ed5c5b2316a5f","skills_sh_rank":17881,"attachment_count":33,"text_attachments":33,"attachment_storage":"skill_attachments","binary_attachments":0},"cluster_size":1,"skill_md_path":"skills/alphaear-predictor/SKILL.md","import_metadata":{"date":"2026-11-02","author":"@skillopedia","version":"v1","category":"finance-legal-compliance","category_label":"Finance"},"exact_dupes_collapsed_into_this":0},"version":"v1","category":"finance-legal-compliance","import_tag":"skills-sh-top20k-2026-11","description":"Market prediction skill using Kronos. Use when user needs finance market time-series forecasting or news-aware finance market adjustments."}},"renderedAt":1793671414806}

Important: agents should read /llm.txt, /llms.txt, or /.well-known/skills.json to discover the public Skillopedia API.