327 lines
12 KiB
Python
327 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
向量存储 - 使用 Ollama 生成向量，FAISS 进行索引和检索
|
||
"""
|
||
import os
|
||
import json
|
||
import numpy as np
|
||
import requests
|
||
from config import INDEX_DIR, OLLAMA_URL, OLLAMA_EMBED_MODEL, TOP_K
|
||
|
||
class VectorStore:
|
||
def __init__(self):
|
||
self.index = None
|
||
self.documents = [] # 存储文档内容和元数据
|
||
self.dimension = 768 # nomic-embed-text 的向量维度
|
||
self.index_file = os.path.join(INDEX_DIR, "faiss.index")
|
||
self.docs_file = os.path.join(INDEX_DIR, "documents.json")
|
||
self.faiss = None
|
||
|
||
def _load_faiss(self):
|
||
"""懒加载 FAISS"""
|
||
if self.faiss is None:
|
||
import faiss
|
||
self.faiss = faiss
|
||
|
||
def _embed_with_ollama(self, text, retry_count=3):
|
||
"""使用 Ollama 生成向量，带重试机制"""
|
||
import time
|
||
import urllib.request
|
||
import urllib.error
|
||
|
||
url = f"{OLLAMA_URL}/api/embeddings"
|
||
|
||
# 确保文本不为空且是字符串
|
||
if not text or not isinstance(text, str):
|
||
text = "empty"
|
||
|
||
# 清理文本中的特殊字符
|
||
text = text.replace('\x00', '') # 移除 null 字符
|
||
|
||
# 截断过长的文本（nomic-embed-text 上下文限制约 2048 tokens）
|
||
# 中文约 1.5 字符/token，保守设置为 1000 字符
|
||
max_length = 1000
|
||
if len(text) > max_length:
|
||
text = text[:max_length]
|
||
|
||
payload = {
|
||
"model": OLLAMA_EMBED_MODEL,
|
||
"prompt": text
|
||
}
|
||
|
||
last_error = None
|
||
for attempt in range(retry_count):
|
||
try:
|
||
# 使用 urllib 代替 requests，避免潜在的编码问题
|
||
data = json.dumps(payload, ensure_ascii=False).encode('utf-8')
|
||
req = urllib.request.Request(
|
||
url,
|
||
data=data,
|
||
headers={'Content-Type': 'application/json; charset=utf-8'},
|
||
method='POST'
|
||
)
|
||
|
||
with urllib.request.urlopen(req, timeout=120) as response:
|
||
result = json.loads(response.read().decode('utf-8'))
|
||
return result.get("embedding", [])
|
||
|
||
except urllib.error.HTTPError as e:
|
||
last_error = e
|
||
error_body = e.read().decode('utf-8') if e.fp else 'N/A'
|
||
print(f"Ollama HTTP 错误 (尝试 {attempt+1}/{retry_count}): {e.code} {e.reason}")
|
||
print(f"响应内容: {error_body[:500]}")
|
||
print(f"请求文本长度: {len(text)}")
|
||
if attempt < retry_count - 1:
|
||
wait_time = (attempt + 1) * 2
|
||
print(f"等待 {wait_time} 秒后重试...")
|
||
time.sleep(wait_time)
|
||
except Exception as e:
|
||
last_error = e
|
||
print(f"Ollama 嵌入失败 (尝试 {attempt+1}/{retry_count}): {e}")
|
||
if attempt < retry_count - 1:
|
||
wait_time = (attempt + 1) * 2
|
||
print(f"等待 {wait_time} 秒后重试...")
|
||
time.sleep(wait_time)
|
||
|
||
raise last_error
|
||
|
||
def _embed_batch(self, texts):
|
||
"""批量生成向量"""
|
||
import time
|
||
embeddings = []
|
||
for i, text in enumerate(texts):
|
||
# 打印文本信息用于调试
|
||
print(f" 生成向量 {i+1}/{len(texts)}...")
|
||
print(f" 文本长度: {len(text)}, 前50字符: {repr(text[:50])}")
|
||
embedding = self._embed_with_ollama(text)
|
||
embeddings.append(embedding)
|
||
# 添加小延迟避免请求过快
|
||
if i < len(texts) - 1:
|
||
time.sleep(1.0)
|
||
return embeddings
|
||
for i, text in enumerate(texts):
|
||
print(f" 生成向量 {i+1}/{len(texts)}...")
|
||
embedding = self._embed_with_ollama(text)
|
||
embeddings.append(embedding)
|
||
# 添加小延迟避免请求过快
|
||
if i < len(texts) - 1:
|
||
time.sleep(0.5)
|
||
return embeddings
|
||
|
||
def _init_index(self):
|
||
"""初始化 FAISS 索引"""
|
||
self._load_faiss()
|
||
if self.index is None:
|
||
self.index = self.faiss.IndexFlatIP(self.dimension)
|
||
|
||
def load_index(self):
|
||
"""从磁盘加载索引"""
|
||
self._load_faiss()
|
||
|
||
if os.path.exists(self.index_file) and os.path.exists(self.docs_file):
|
||
try:
|
||
print("正在加载已有索引...")
|
||
|
||
# FAISS 在 Windows 上不支持中文路径，使用临时文件
|
||
import tempfile
|
||
import shutil
|
||
|
||
try:
|
||
# 复制到临时文件再读取
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix='.index') as tmp:
|
||
tmp_path = tmp.name
|
||
shutil.copy2(self.index_file, tmp_path)
|
||
self.index = self.faiss.read_index(tmp_path)
|
||
os.unlink(tmp_path)
|
||
except Exception as e:
|
||
print(f"临时文件方式失败，尝试直接读取: {e}")
|
||
self.index = self.faiss.read_index(self.index_file)
|
||
|
||
with open(self.docs_file, 'r', encoding='utf-8') as f:
|
||
self.documents = json.load(f)
|
||
print(f"索引加载完成，共 {len(self.documents)} 个文档块")
|
||
return True
|
||
except Exception as e:
|
||
print(f"加载索引失败: {e}")
|
||
self._init_index()
|
||
self.documents = []
|
||
return False
|
||
else:
|
||
print("未找到已有索引，创建新索引")
|
||
self._init_index()
|
||
self.documents = []
|
||
return False
|
||
|
||
def save_index(self):
|
||
"""保存索引到磁盘"""
|
||
self._load_faiss()
|
||
if self.index is not None:
|
||
# 确保目录存在
|
||
os.makedirs(os.path.dirname(self.index_file), exist_ok=True)
|
||
|
||
# FAISS 在 Windows 上不支持中文路径，使用临时文件再移动
|
||
import tempfile
|
||
import shutil
|
||
|
||
try:
|
||
# 先写入临时文件
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix='.index') as tmp:
|
||
tmp_path = tmp.name
|
||
|
||
self.faiss.write_index(self.index, tmp_path)
|
||
|
||
# 移动到目标位置
|
||
shutil.move(tmp_path, self.index_file)
|
||
except Exception as e:
|
||
# 如果临时文件方式失败，尝试直接写入
|
||
print(f"临时文件方式失败，尝试直接写入: {e}")
|
||
self.faiss.write_index(self.index, self.index_file)
|
||
|
||
with open(self.docs_file, 'w', encoding='utf-8') as f:
|
||
json.dump(self.documents, f, ensure_ascii=False, indent=2)
|
||
print(f"索引已保存，共 {len(self.documents)} 个文档块")
|
||
|
||
def add_documents(self, chunks, metadata=None):
|
||
"""添加文档块到索引"""
|
||
if not chunks:
|
||
return 0
|
||
|
||
self._load_faiss()
|
||
self._init_index()
|
||
|
||
# 使用 Ollama 生成向量
|
||
print(f"正在为 {len(chunks)} 个文本块生成向量...")
|
||
embeddings = self._embed_batch(chunks)
|
||
|
||
# 检查向量维度
|
||
if embeddings and len(embeddings[0]) != self.dimension:
|
||
self.dimension = len(embeddings[0])
|
||
self.index = self.faiss.IndexFlatIP(self.dimension)
|
||
print(f"更新向量维度为: {self.dimension}")
|
||
|
||
# 归一化向量（用于余弦相似度）
|
||
embeddings_np = np.array(embeddings).astype('float32')
|
||
norms = np.linalg.norm(embeddings_np, axis=1, keepdims=True)
|
||
embeddings_np = embeddings_np / (norms + 1e-10)
|
||
|
||
# 添加到索引
|
||
start_idx = len(self.documents)
|
||
self.index.add(embeddings_np)
|
||
|
||
# 保存文档内容和元数据
|
||
for i, chunk in enumerate(chunks):
|
||
doc = {
|
||
'id': start_idx + i,
|
||
'content': chunk,
|
||
'metadata': metadata or {}
|
||
}
|
||
self.documents.append(doc)
|
||
|
||
# 自动保存
|
||
self.save_index()
|
||
|
||
return len(chunks)
|
||
|
||
def search(self, query, top_k=None):
|
||
"""搜索相关文档"""
|
||
if top_k is None:
|
||
top_k = TOP_K
|
||
|
||
if self.index is None or self.index.ntotal == 0:
|
||
return []
|
||
|
||
self._load_faiss()
|
||
|
||
# 生成查询向量
|
||
query_embedding = self._embed_with_ollama(query)
|
||
query_np = np.array([query_embedding]).astype('float32')
|
||
|
||
# 归一化
|
||
norm = np.linalg.norm(query_np)
|
||
query_np = query_np / (norm + 1e-10)
|
||
|
||
# 搜索
|
||
k = min(top_k, self.index.ntotal)
|
||
scores, indices = self.index.search(query_np, k)
|
||
|
||
# 构建结果
|
||
results = []
|
||
for i, idx in enumerate(indices[0]):
|
||
if idx < len(self.documents) and idx >= 0:
|
||
doc = self.documents[idx]
|
||
results.append({
|
||
'content': doc['content'],
|
||
'score': float(scores[0][i]),
|
||
'metadata': doc.get('metadata', {})
|
||
})
|
||
|
||
return results
|
||
|
||
def delete_by_filename(self, filename):
|
||
"""删除指定文件的所有文档块"""
|
||
if not self.documents:
|
||
return 0
|
||
|
||
self._load_faiss()
|
||
|
||
# 找出要保留的文档
|
||
remaining_docs = []
|
||
deleted_count = 0
|
||
|
||
for doc in self.documents:
|
||
if doc.get('metadata', {}).get('filename') != filename:
|
||
remaining_docs.append(doc)
|
||
else:
|
||
deleted_count += 1
|
||
|
||
if deleted_count > 0:
|
||
# 重建索引
|
||
self.documents = []
|
||
self.index = self.faiss.IndexFlatIP(self.dimension)
|
||
|
||
if remaining_docs:
|
||
chunks = [doc['content'] for doc in remaining_docs]
|
||
metadatas = [doc.get('metadata', {}) for doc in remaining_docs]
|
||
|
||
embeddings = self._embed_batch(chunks)
|
||
embeddings_np = np.array(embeddings).astype('float32')
|
||
norms = np.linalg.norm(embeddings_np, axis=1, keepdims=True)
|
||
embeddings_np = embeddings_np / (norms + 1e-10)
|
||
self.index.add(embeddings_np)
|
||
|
||
for i, (chunk, meta) in enumerate(zip(chunks, metadatas)):
|
||
self.documents.append({
|
||
'id': i,
|
||
'content': chunk,
|
||
'metadata': meta
|
||
})
|
||
|
||
self.save_index()
|
||
|
||
return deleted_count
|
||
|
||
def clear(self):
|
||
"""清空所有索引"""
|
||
self._load_faiss()
|
||
self.index = self.faiss.IndexFlatIP(self.dimension)
|
||
self.documents = []
|
||
self.save_index()
|
||
print("索引已清空")
|
||
|
||
def get_stats(self):
|
||
"""获取索引统计信息"""
|
||
files = set()
|
||
for doc in self.documents:
|
||
filename = doc.get('metadata', {}).get('filename')
|
||
if filename:
|
||
files.add(filename)
|
||
|
||
return {
|
||
'total_chunks': len(self.documents),
|
||
'total_files': len(files),
|
||
'files': list(files)
|
||
}
|
||
|
||
# 全局实例
|
||
vector_store = VectorStore()
|