39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
# app/agents/vectorization_agent.py
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
from ..core.vector_store import vector_store, embedding_model
|
|
|
|
# 初始化文本分割器，用于将长文档切成小块
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=500, # 每个块的大小（字符数）
|
|
chunk_overlap=50, # 块之间的重叠部分
|
|
)
|
|
|
|
def agent_vectorize_and_store(doc_id: str, text: str, category: str):
|
|
"""Agent 4: 向量化并存储 (真实实现)"""
|
|
print(f"--- [Agent 4] 正在向量化文档 (ID: {doc_id})...")
|
|
|
|
# 1. 将文档文本分割成块
|
|
chunks = text_splitter.split_text(text)
|
|
print(f"--- [Agent 4] 文档被切分为 {len(chunks)} 个块。")
|
|
|
|
if not chunks:
|
|
print(f"--- [Agent 4] 文档内容为空，跳过向量化。")
|
|
return
|
|
|
|
# 2. 为每个块创建唯一的ID和元数据
|
|
chunk_ids = [f"{doc_id}_{i}" for i in range(len(chunks))]
|
|
metadatas = [{"doc_id": doc_id, "category": category, "chunk_number": i} for i in range(len(chunks))]
|
|
|
|
# 3. 使用嵌入模型为所有块生成向量
|
|
embeddings = embedding_model.embed_documents(chunks)
|
|
|
|
# 4. 将ID、向量、元数据和文本块本身添加到ChromaDB
|
|
vector_store.add(
|
|
ids=chunk_ids,
|
|
embeddings=embeddings,
|
|
documents=chunks,
|
|
metadatas=metadatas
|
|
)
|
|
|
|
print(f"--- [Agent 4] 文档 {doc_id} 的向量已存入ChromaDB。")
|