142 lines
3.6 KiB
Python
142 lines
3.6 KiB
Python
"""文本组装与向量生成。
|
|
|
|
核心逻辑:
|
|
1. compose_text: 从 EnTranslatedArticle 拼接中文文本供嵌入
|
|
2. embed_article: 单篇文章嵌入
|
|
3. embed_articles: 批量嵌入（自动分块）
|
|
"""
|
|
|
|
import logging
|
|
|
|
from openai import OpenAI
|
|
|
|
from embedding.client import EmbeddingConfig, embed_batch
|
|
from embedding.models import EmbeddingResult
|
|
from llm.models import EnTranslatedArticle
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 嵌入文本最大字符数（DashScope text-embedding-v3 支持 8192 token，保守取 4000 字符）
|
|
MAX_EMBED_CHARS = 4000
|
|
|
|
|
|
def compose_text(article: EnTranslatedArticle, max_chars: int = MAX_EMBED_CHARS) -> str:
|
|
"""把 EnTranslatedArticle 组装成单段嵌入文本。
|
|
|
|
拼接策略:
|
|
- 中文标题（最高信号）
|
|
- 事件摘要（语义浓缩）
|
|
- 中文正文（截断）
|
|
|
|
Args:
|
|
article: M4 输出的双语文章
|
|
max_chars: 整段最大字符数
|
|
|
|
Returns:
|
|
拼接后的中文嵌入文本
|
|
"""
|
|
parts: list[str] = []
|
|
|
|
# 1. 标题（中文）
|
|
if article.title_zh:
|
|
parts.append(f"标题: {article.title_zh}")
|
|
|
|
# 2. 事件标签（语义浓缩）
|
|
for ev in article.events:
|
|
event_parts = [
|
|
f"[{ev.sentiment.value}]",
|
|
f"{ev.event_type}",
|
|
f"重要度{ev.importance}",
|
|
]
|
|
if ev.stock_codes:
|
|
event_parts.append("代码:" + ",".join(ev.stock_codes[:5]))
|
|
if ev.summary_zh:
|
|
event_parts.append(ev.summary_zh)
|
|
parts.append("事件: " + " ".join(event_parts))
|
|
|
|
# 3. 中文正文
|
|
if article.content_zh:
|
|
body = article.content_zh
|
|
parts.append(f"正文: {body}")
|
|
|
|
text = "\n".join(parts)
|
|
|
|
# 截断保护
|
|
if len(text) > max_chars:
|
|
logger.debug("嵌入文本超长 %d → %d", len(text), max_chars)
|
|
text = text[:max_chars]
|
|
|
|
return text
|
|
|
|
|
|
def embed_article(
|
|
client: OpenAI,
|
|
config: EmbeddingConfig,
|
|
article: EnTranslatedArticle,
|
|
) -> EmbeddingResult:
|
|
"""单篇文章嵌入。
|
|
|
|
Args:
|
|
client: OpenAI 客户端
|
|
config: Embedding 配置
|
|
article: M4 输出的双语文章
|
|
|
|
Returns:
|
|
EmbeddingResult 含向量 + 元信息
|
|
"""
|
|
text = compose_text(article)
|
|
|
|
vectors = embed_batch(client, config, [text])
|
|
if not vectors:
|
|
raise ValueError(f"嵌入返回空结果: {article.url_hash}")
|
|
|
|
return EmbeddingResult(
|
|
url_hash=article.url_hash,
|
|
source_id=article.source_id,
|
|
vector=vectors[0],
|
|
dimension=len(vectors[0]),
|
|
embedded_text=text,
|
|
provider=config.provider,
|
|
model=config.model,
|
|
)
|
|
|
|
|
|
def embed_articles(
|
|
client: OpenAI,
|
|
config: EmbeddingConfig,
|
|
articles: list[EnTranslatedArticle],
|
|
) -> list[EmbeddingResult]:
|
|
"""批量嵌入多篇文章（自动分块+重试）。
|
|
|
|
Args:
|
|
client: OpenAI 客户端
|
|
config: Embedding 配置
|
|
articles: M4 输出的双语文章列表
|
|
|
|
Returns:
|
|
EmbeddingResult 列表
|
|
"""
|
|
if not articles:
|
|
return []
|
|
|
|
# 组装所有嵌入文本
|
|
texts = [compose_text(a) for a in articles]
|
|
|
|
# 批量嵌入
|
|
vectors = embed_batch(client, config, texts)
|
|
|
|
# 组装结果
|
|
results: list[EmbeddingResult] = []
|
|
for article, vec in zip(articles, vectors):
|
|
results.append(EmbeddingResult(
|
|
url_hash=article.url_hash,
|
|
source_id=article.source_id,
|
|
vector=vec,
|
|
dimension=len(vec),
|
|
embedded_text=compose_text(article),
|
|
provider=config.provider,
|
|
model=config.model,
|
|
))
|
|
|
|
return results
|