Initial commit

novel_generator/vectorstore_utils.py

@@ -0,0 +1,244 @@
+#novel_generator/vectorstore_utils.py
+# -*- coding: utf-8 -*-
+"""
+向量库相关操作（初始化、更新、检索、清空、文本切分等）
+"""
+import os
+import logging
+import traceback
+import nltk
+import numpy as np
+import re
+import ssl
+import requests
+import warnings
+from langchain_chroma import Chroma
+
+# 禁用特定的Torch警告
+warnings.filterwarnings('ignore', message='.*Torch was not compiled with flash attention.*')
+os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 禁用tokenizer并行警告
+
+from chromadb.config import Settings
+from langchain.docstore.document import Document
+from sklearn.metrics.pairwise import cosine_similarity
+from .common import call_with_retry
+
+def get_vectorstore_dir(filepath: str) -> str:
+    """获取 vectorstore 路径"""
+    return os.path.join(filepath, "vectorstore")
+
+def clear_vector_store(filepath: str) -> bool:
+    """清空 清空向量库"""
+    import shutil
+    store_dir = get_vectorstore_dir(filepath)
+    if not os.path.exists(store_dir):
+        logging.info("No vector store found to clear.")
+        return False
+    try:
+        shutil.rmtree(store_dir)
+        logging.info(f"Vector store directory '{store_dir}' removed.")
+        return True
+    except Exception as e:
+        logging.error(f"无法删除向量库文件夹，请关闭程序后手动删除 {store_dir}。\n {str(e)}")
+        traceback.print_exc()
+        return False
+
+def init_vector_store(embedding_adapter, texts, filepath: str):
+    """
+    在 filepath 下创建/加载一个 Chroma 向量库并插入 texts。
+    如果Embedding失败，则返回 None，不中断任务。
+    """
+    from langchain.embeddings.base import Embeddings as LCEmbeddings
+
+    store_dir = get_vectorstore_dir(filepath)
+    os.makedirs(store_dir, exist_ok=True)
+    documents = [Document(page_content=str(t)) for t in texts]
+
+    try:
+        class LCEmbeddingWrapper(LCEmbeddings):
+            def embed_documents(self, texts):
+                return call_with_retry(
+                    func=embedding_adapter.embed_documents,
+                    max_retries=3,
+                    fallback_return=[],
+                    texts=texts
+                )
+            def embed_query(self, query: str):
+                res = call_with_retry(
+                    func=embedding_adapter.embed_query,
+                    max_retries=3,
+                    fallback_return=[],
+                    query=query
+                )
+                return res
+
+        chroma_embedding = LCEmbeddingWrapper()
+        vectorstore = Chroma.from_documents(
+            documents,
+            embedding=chroma_embedding,
+            persist_directory=store_dir,
+            client_settings=Settings(anonymized_telemetry=False),
+            collection_name="novel_collection"
+        )
+        return vectorstore
+    except Exception as e:
+        logging.warning(f"Init vector store failed: {e}")
+        traceback.print_exc()
+        return None
+
+def load_vector_store(embedding_adapter, filepath: str):
+    """
+    读取已存在的 Chroma 向量库。若不存在则返回 None。
+    如果加载失败（embedding 或IO问题），则返回 None。
+    """
+    from langchain.embeddings.base import Embeddings as LCEmbeddings
+    store_dir = get_vectorstore_dir(filepath)
+    if not os.path.exists(store_dir):
+        logging.info("Vector store not found. Will return None.")
+        return None
+
+    try:
+        class LCEmbeddingWrapper(LCEmbeddings):
+            def embed_documents(self, texts):
+                return call_with_retry(
+                    func=embedding_adapter.embed_documents,
+                    max_retries=3,
+                    fallback_return=[],
+                    texts=texts
+                )
+            def embed_query(self, query: str):
+                res = call_with_retry(
+                    func=embedding_adapter.embed_query,
+                    max_retries=3,
+                    fallback_return=[],
+                    query=query
+                )
+                return res
+
+        chroma_embedding = LCEmbeddingWrapper()
+        return Chroma(
+            persist_directory=store_dir,
+            embedding_function=chroma_embedding,
+            client_settings=Settings(anonymized_telemetry=False),
+            collection_name="novel_collection"
+        )
+    except Exception as e:
+        logging.warning(f"Failed to load vector store: {e}")
+        traceback.print_exc()
+        return None
+
+def split_by_length(text: str, max_length: int = 500):
+    """按照 max_length 切分文本"""
+    segments = []
+    start_idx = 0
+    while start_idx < len(text):
+        end_idx = min(start_idx + max_length, len(text))
+        segment = text[start_idx:end_idx]
+        segments.append(segment.strip())
+        start_idx = end_idx
+    return segments
+
+def split_text_for_vectorstore(chapter_text: str, max_length: int = 500, similarity_threshold: float = 0.7):
+    """
+    对新的章节文本进行分段后,再用于存入向量库。
+    使用 embedding 进行文本相似度计算。
+    """
+    if not chapter_text.strip():
+        return []
+    
+    nltk.download('punkt', quiet=True)
+    nltk.download('punkt_tab', quiet=True)
+    sentences = nltk.sent_tokenize(chapter_text)
+    if not sentences:
+        return []
+    
+    # 直接按长度分段,不做相似度合并
+    final_segments = []
+    current_segment = []
+    current_length = 0
+    
+    for sentence in sentences:
+        sentence_length = len(sentence)
+        if current_length + sentence_length > max_length:
+            if current_segment:
+                final_segments.append(" ".join(current_segment))
+            current_segment = [sentence]
+            current_length = sentence_length
+        else:
+            current_segment.append(sentence)
+            current_length += sentence_length
+    
+    if current_segment:
+        final_segments.append(" ".join(current_segment))
+    
+    return final_segments
+
+def update_vector_store(embedding_adapter, new_chapter: str, filepath: str):
+    """
+    将最新章节文本插入到向量库中。
+    若库不存在则初始化；若初始化/更新失败，则跳过。
+    """
+    from utils import read_file, clear_file_content, save_string_to_txt
+    splitted_texts = split_text_for_vectorstore(new_chapter)
+    if not splitted_texts:
+        logging.warning("No valid text to insert into vector store. Skipping.")
+        return
+
+    store = load_vector_store(embedding_adapter, filepath)
+    if not store:
+        logging.info("Vector store does not exist or failed to load. Initializing a new one for new chapter...")
+        store = init_vector_store(embedding_adapter, splitted_texts, filepath)
+        if not store:
+            logging.warning("Init vector store failed, skip embedding.")
+        else:
+            logging.info("New vector store created successfully.")
+        return
+
+    try:
+        docs = [Document(page_content=str(t)) for t in splitted_texts]
+        store.add_documents(docs)
+        logging.info("Vector store updated with the new chapter splitted segments.")
+    except Exception as e:
+        logging.warning(f"Failed to update vector store: {e}")
+        traceback.print_exc()
+
+def get_relevant_context_from_vector_store(embedding_adapter, query: str, filepath: str, k: int = 2) -> str:
+    """
+    从向量库中检索与 query 最相关的 k 条文本，拼接后返回。
+    如果向量库加载/检索失败，则返回空字符串。
+    最终只返回最多2000字符的检索片段。
+    """
+    store = load_vector_store(embedding_adapter, filepath)
+    if not store:
+        logging.info("No vector store found or load failed. Returning empty context.")
+        return ""
+
+    try:
+        docs = store.similarity_search(query, k=k)
+        if not docs:
+            logging.info(f"No relevant documents found for query '{query}'. Returning empty context.")
+            return ""
+        combined = "\n".join([d.page_content for d in docs])
+        if len(combined) > 2000:
+            combined = combined[:2000]
+        return combined
+    except Exception as e:
+        logging.warning(f"Similarity search failed: {e}")
+        traceback.print_exc()
+        return ""
+
+def _get_sentence_transformer(model_name: str = 'paraphrase-MiniLM-L6-v2'):
+    """获取sentence transformer模型，处理SSL问题"""
+    try:
+        # 设置torch环境变量
+        os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "0"
+        os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "0"
+        
+        # 禁用SSL验证
+        ssl._create_default_https_context = ssl._create_unverified_context
+        
+        # ...existing code...
+    except Exception as e:
+        logging.error(f"Failed to load sentence transformer model: {e}")
+        traceback.print_exc()
+        return None