136 lines
5.1 KiB
Python
136 lines
5.1 KiB
Python
import os
|
||
from qdrant_client import QdrantClient
|
||
from qdrant_client.models import Distance, VectorParams, PointStruct
|
||
|
||
class LocalKnowledgeBase:
|
||
def __init__(self, collection_name="documents", persist_dir="./qdrant_data"):
|
||
# Инициализируем клиент Qdrant
|
||
self.client = QdrantClient(
|
||
path=persist_dir, # Локальное хранение
|
||
prefer_grpc=True
|
||
)
|
||
|
||
self.collection_name = collection_name
|
||
self.embeddings = embeddings
|
||
|
||
# Создаем коллекцию если её нет
|
||
self._create_collection()
|
||
|
||
def _create_collection(self):
|
||
try:
|
||
self.client.get_collection(self.collection_name)
|
||
print(f"Коллекция {self.collection_name} уже существует")
|
||
except:
|
||
# Создаем новую коллекцию
|
||
self.client.create_collection(
|
||
collection_name=self.collection_name,
|
||
vectors_config=VectorParams(
|
||
size=768, # Размерность эмбеддингов nomic-embed-text
|
||
distance=Distance.COSINE
|
||
)
|
||
)
|
||
print(f"Создана коллекция {self.collection_name}")
|
||
|
||
def load_documents(self, directory_path: str):
|
||
"""Загрузка документов из директории"""
|
||
loaders = {
|
||
'.pdf': PyPDFLoader,
|
||
'.txt': lambda path: DirectoryLoader(path, glob="**/*.txt"),
|
||
'.docx': lambda path: DirectoryLoader(path, glob="**/*.docx"),
|
||
}
|
||
|
||
all_documents = []
|
||
|
||
for ext, loader_class in loaders.items():
|
||
for file_path in os.listdir(directory_path):
|
||
if file_path.endswith(ext):
|
||
full_path = os.path.join(directory_path, file_path)
|
||
try:
|
||
loader = loader_class(full_path)
|
||
documents = loader.load()
|
||
all_documents.extend(documents)
|
||
print(f"Загружен {file_path}: {len(documents)} страниц")
|
||
except Exception as e:
|
||
print(f"Ошибка загрузки {file_path}: {e}")
|
||
|
||
# Разбиваем на чанки
|
||
text_splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=1000,
|
||
chunk_overlap=200,
|
||
length_function=len
|
||
)
|
||
|
||
chunks = text_splitter.split_documents(all_documents)
|
||
print(f"Всего чанков: {len(chunks)}")
|
||
|
||
# Создаем эмбеддинги и сохраняем в Qdrant
|
||
self._index_documents(chunks)
|
||
|
||
return chunks
|
||
|
||
def _index_documents(self, documents):
|
||
"""Индексация документов в Qdrant"""
|
||
points = []
|
||
|
||
for i, doc in enumerate(documents):
|
||
# Создаем эмбеддинг для каждого чанка
|
||
embedding = self.embeddings.embed_query(doc.page_content)
|
||
|
||
point = PointStruct(
|
||
id=i,
|
||
vector=embedding,
|
||
payload={
|
||
"text": doc.page_content,
|
||
"source": doc.metadata.get("source", "unknown"),
|
||
"page": doc.metadata.get("page", 0)
|
||
}
|
||
)
|
||
points.append(point)
|
||
|
||
# Пакетная загрузка каждые 100 точек
|
||
if len(points) >= 100:
|
||
self.client.upsert(
|
||
collection_name=self.collection_name,
|
||
points=points
|
||
)
|
||
points = []
|
||
print(f"Индексировано {i+1} документов")
|
||
|
||
# Загружаем оставшиеся
|
||
if points:
|
||
self.client.upsert(
|
||
collection_name=self.collection_name,
|
||
points=points
|
||
)
|
||
|
||
print(f"Индексация завершена. Всего документов: {len(documents)}")
|
||
|
||
def search(self, query: str, top_k: int = 5):
|
||
"""Поиск в базе знаний"""
|
||
# Создаем эмбеддинг запроса
|
||
query_embedding = self.embeddings.embed_query(query)
|
||
|
||
# Ищем в Qdrant
|
||
search_result = self.client.search(
|
||
collection_name=self.collection_name,
|
||
query_vector=query_embedding,
|
||
limit=top_k
|
||
)
|
||
|
||
# Форматируем результаты
|
||
results = []
|
||
for hit in search_result:
|
||
results.append({
|
||
"text": hit.payload["text"],
|
||
"score": hit.score,
|
||
"source": hit.payload.get("source", "unknown")
|
||
})
|
||
|
||
return results
|
||
|
||
# Инициализируем базу знаний
|
||
kb = LocalKnowledgeBase()
|
||
|
||
# Загружаем документы (если есть)
|
||
if os.path.exists("./documents"):
|
||
kb.load_documents("./documents") |