app/document_processor.py - main - cppv/hackathon13

app/document_processor.py
66 строк3 KB
cppv
hackathon13
25 май 2025, 11:54
25 май 2025, 11:548af8c2d
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import FakeEmbeddings
import os

class DocumentProcessor:
    # тут всякие штуки для работы с доками лангчейна
    # делаем из них поиск по векторам
    
    def __init__(self, docs_dir, index_path):
        # путь где лежат доки и куда сохранять индекс
        self.docs_dir = docs_dir
        self.index_path = index_path
        
    def load_documents(self):
        # загружаем все markdown файлы из папки
        loader = DirectoryLoader(self.docs_dir, glob="**/*.md", loader_cls=TextLoader)
        documents = loader.load()
        return documents
    
    def split_documents(self, documents, chunk_size=1000, chunk_overlap=200):
        # режем доки на кусочки чтоб в память влезали
        # overlap это насколько куски пересекаются
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
        )
        chunks = text_splitter.split_documents(documents)
        return chunks
    
    def create_index(self, chunks, embeddings):
        # делаем из кусочков поисковый индекс
        # embeddings это штука которая текст в векторы превращает
        vectorstore = FAISS.from_documents(chunks, embeddings)
        return vectorstore
    
    def save_index(self, vectorstore):
        # сохраняем индекс на диск
        vectorstore.save_local(self.index_path)
        
    def load_index(self, embeddings):
        # пытаемся загрузить индекс если он уже есть
        index_faiss = os.path.join(self.index_path, "index.faiss")
        index_pkl = os.path.join(self.index_path, "index.pkl")
        if os.path.exists(index_faiss) and os.path.exists(index_pkl):
            vectorstore = FAISS.load_local(self.index_path, embeddings, allow_dangerous_deserialization=True)
            return vectorstore
        return None
    
    def process_documents(self, embeddings):
        # тут все вместе - загружаем или создаем индекс
        
        # сначала смотрим может он уже есть
        vectorstore = self.load_index(embeddings)
        if vectorstore:
            return vectorstore
        
        # если нет - делаем с нуля
        documents = self.load_documents()
        chunks = self.split_documents(documents)
        vectorstore = self.create_index(chunks, embeddings)
        self.save_index(vectorstore)
        
        return vectorstore