ДОбился запуска
This commit is contained in:
@@ -1,23 +1,32 @@
|
||||
import os
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import Distance, VectorParams, PointStruct
|
||||
from langchain_community.document_loaders import TextLoader, PyPDFLoader, DirectoryLoader
|
||||
from langchain.text_splitter import Language
|
||||
from langchain_community.document_loaders import (
|
||||
TextLoader,
|
||||
PyPDFLoader,
|
||||
DirectoryLoader,
|
||||
)
|
||||
from langchain_text_splitters import Language, RecursiveCharacterTextSplitter
|
||||
from langchain_community.embeddings import OllamaEmbeddings
|
||||
|
||||
|
||||
class LocalKnowledgeBase:
|
||||
def __init__(self, collection_name="documents", persist_dir="./qdrant_data"):
|
||||
# Инициализируем клиент Qdrant
|
||||
self.client = QdrantClient(
|
||||
path=persist_dir, # Локальное хранение
|
||||
prefer_grpc=True
|
||||
path=persist_dir, prefer_grpc=True # Локальное хранение
|
||||
)
|
||||
|
||||
|
||||
self.collection_name = collection_name
|
||||
self.embeddings = embeddings
|
||||
|
||||
|
||||
# Инициализация эмбеддингов
|
||||
self.embeddings = OllamaEmbeddings(
|
||||
model="nomic-embed-text", # Хорошие локальные эмбеддинги
|
||||
)
|
||||
|
||||
# Создаем коллекцию если её нет
|
||||
self._create_collection()
|
||||
|
||||
|
||||
def _create_collection(self):
|
||||
try:
|
||||
self.client.get_collection(self.collection_name)
|
||||
@@ -28,26 +37,25 @@ class LocalKnowledgeBase:
|
||||
collection_name=self.collection_name,
|
||||
vectors_config=VectorParams(
|
||||
size=768, # Размерность эмбеддингов nomic-embed-text
|
||||
distance=Distance.COSINE
|
||||
)
|
||||
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"),
|
||||
'.hpp': lambda path: TextLoader(path, encoding='utf-8'),
|
||||
'.cpp': lambda path: TextLoader(path, encoding='utf-8'), # добавить
|
||||
'.h': lambda path: TextLoader(path, encoding='utf-8'), # добавить
|
||||
'.cc': lambda path: TextLoader(path, encoding='utf-8'), # добавить
|
||||
|
||||
".pdf": PyPDFLoader,
|
||||
".txt": lambda path: DirectoryLoader(path, glob="**/*.txt"),
|
||||
".docx": lambda path: DirectoryLoader(path, glob="**/*.docx"),
|
||||
".hpp": lambda path: TextLoader(path, encoding="utf-8"),
|
||||
".cpp": lambda path: TextLoader(path, encoding="utf-8"), # добавить
|
||||
".h": lambda path: TextLoader(path, encoding="utf-8"), # добавить
|
||||
".cc": lambda path: TextLoader(path, encoding="utf-8"), # добавить
|
||||
}
|
||||
|
||||
|
||||
all_documents = []
|
||||
|
||||
|
||||
for ext, loader_class in loaders.items():
|
||||
for file_path in os.listdir(directory_path):
|
||||
if file_path.endswith(ext):
|
||||
@@ -59,31 +67,31 @@ class LocalKnowledgeBase:
|
||||
print(f"Загружен {file_path}: {len(documents)} страниц")
|
||||
except Exception as e:
|
||||
print(f"Ошибка загрузки {file_path}: {e}")
|
||||
|
||||
|
||||
# Разбиваем на чанки
|
||||
text_splitter = RecursiveCharacterTextSplitter.from_language(
|
||||
language=Language.CPP,
|
||||
chunk_size=1000,
|
||||
chunk_overlap=200,
|
||||
length_function=len
|
||||
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,
|
||||
@@ -91,55 +99,45 @@ class LocalKnowledgeBase:
|
||||
"text": doc.page_content,
|
||||
"source": doc.metadata.get("source", "unknown"),
|
||||
"page": doc.metadata.get("page", 0),
|
||||
"file_type": "cpp"
|
||||
}
|
||||
"file_type": "cpp",
|
||||
},
|
||||
)
|
||||
points.append(point)
|
||||
|
||||
|
||||
# Пакетная загрузка каждые 100 точек
|
||||
if len(points) >= 100:
|
||||
self.client.upsert(
|
||||
collection_name=self.collection_name,
|
||||
points=points
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
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(
|
||||
search_result = self.client.query_points(
|
||||
collection_name=self.collection_name,
|
||||
query_vector=query_embedding,
|
||||
limit=top_k
|
||||
query=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")
|
||||
})
|
||||
|
||||
for hit in search_result.points:
|
||||
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")
|
||||
Reference in New Issue
Block a user