40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
# app/agents/vectorization_agent.py
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=100,
|
|
)
|
|
|
|
def agent_vectorize_and_store(
|
|
doc_id: str,
|
|
text: str,
|
|
category: str,
|
|
language: str,
|
|
embedding_model,
|
|
vector_store
|
|
):
|
|
print(f"--- [Background Task] Starting vectorization (ID: {doc_id})...")
|
|
|
|
try:
|
|
chunks = text_splitter.split_text(text)
|
|
if not chunks:
|
|
print(f"--- [Background task] document is empty, skip vectorization. (ID: {doc_id})")
|
|
return
|
|
|
|
chunk_ids = [f"{doc_id}_{i}" for i in range(len(chunks))]
|
|
metadatas = [{"doc_id": doc_id, "category": category, "language": language, "chunk_number": i} for i in
|
|
range(len(chunks))]
|
|
|
|
embeddings = embedding_model.embed_documents(chunks)
|
|
|
|
vector_store.add(
|
|
ids=chunk_ids,
|
|
embeddings=embeddings,
|
|
documents=chunks,
|
|
metadatas=metadatas
|
|
)
|
|
print(f"--- [Background Task] Document {doc_id} vectorized。")
|
|
except Exception as e:
|
|
print(f"--- [Background Task] Document vectorization failed (ID: {doc_id}): {e}")
|