Update indexing.py

indexing.py

@@ -1,63 +1,32 @@
-"""
-Indexing with vector database
-"""
-
-from pathlib import Path
-import re
-
+# indexing.py
 import chromadb
-
-from unidecode import unidecode
-
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_chroma import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
-
-
-
-def load_doc(list_file_path, chunk_size, chunk_overlap):
-    """Carrega documentos PDF e divide em chunks"""
-    loaders = [PyPDFLoader(x) for x in list_file_path]
-    pages = []
-    for loader in loaders:
-        pages.extend(loader.load())
-
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size,
-        chunk_overlap=chunk_overlap
-    )
-    doc_splits = text_splitter.split_documents(pages)
-    return doc_splits
-
-
-def create_collection_name(filepath):
-    """Cria nome da coleção com base no nome do arquivo PDF"""
-    collection_name = Path(filepath).stem
-    collection_name = collection_name.replace(" ", "-")
-    collection_name = unidecode(collection_name)
-    collection_name = re.sub("[^A-Za-z0-9]+", "-", collection_name)
-    collection_name = collection_name[:50]
-    if len(collection_name) < 3:
-        collection_name += "xyz"
-    if not collection_name[0].isalnum():
-        collection_name = "A" + collection_name[1:]
-    if not collection_name[-1].isalnum():
-        collection_name = collection_name[:-1] + "Z"
-    return collection_name
-
-
-def create_db(splits, collection_name):
-    """Cria banco vetorial com os embeddings"""
-    embedding = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-    )
-    chromadb.api.client.SharedSystemClient.clear_system_cache()
-    new_client = chromadb.EphemeralClient()
-    vectordb = Chroma.from_documents(
-        documents=splits,
-        embedding=embedding,
-        client=new_client,
-        collection_name=collection_name,
-    )
-    return vectordb
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+
+# Inicializa o modelo de embedding
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+# Inicializa o cliente ChromaDB com persistência local
+chroma_client = chromadb.Client(Settings(
+    chroma_db_impl="duckdb+parquet",
+    persist_directory="./chroma_db"  # Certifique-se de que essa pasta está no .gitignore
+))
+
+# Cria ou recupera a coleção
+collection = chroma_client.get_or_create_collection(name="document_chunks")
+
+def index_chunks(chunks):
+    """
+    Recebe uma lista de chunks de texto e adiciona ao banco vetorial.
+    """
+    embeddings = embedding_model.encode(chunks).tolist()
+    ids = [f"chunk_{i}" for i in range(len(chunks))]
+    collection.add(documents=chunks, embeddings=embeddings, ids=ids)
+
+def query_relevant_chunk(query_text, n_results=1):
+    """
+    Realiza uma busca vetorial e retorna os chunks mais relevantes.
+    """
+    query_embedding = embedding_model.encode([query_text]).tolist()[0]
+    results = collection.query(query_embeddings=[query_embedding], n_results=n_results)
+    return results["documents"][0] if results["documents"] else []