First!

2026-03-31 23:08:06 +05:00
commit 2ae08db4b2
4 changed files with 413 additions and 0 deletions
--- a/local_knowledge_base.py
+++ b/local_knowledge_base.py
@@ -0,0 +1,136 @@
+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")