lgcosta
8eeb6b89f4
Usa o modelo multilingual, ignora trechos curtos e o template do Llama 3.2 para a resposta sair em uma frase completa. O chat deixa a API livre durante a geração, o upload indexa o PDF e os agendamentos não são gravados duas vezes.
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
import json
|
|
import os
|
|
import uuid
|
|
|
|
from app.services.pdf_processor import PDFProcessor
|
|
from app.services.rag_service import QdrantService
|
|
from app.utils.embeddings import EmbeddingService
|
|
|
|
|
|
def book_dirs() -> list[str]:
|
|
paths = [
|
|
os.getenv("BOOKS_DIR", "/app/books"),
|
|
os.getenv("UPLOAD_DIR", "/app/uploads"),
|
|
]
|
|
seen = []
|
|
for path in paths:
|
|
if path and path not in seen and os.path.isdir(path):
|
|
seen.append(path)
|
|
return seen
|
|
|
|
|
|
def main():
|
|
embedding_service = EmbeddingService()
|
|
pdf_processor = PDFProcessor(embedding_service)
|
|
qdrant = QdrantService(os.getenv("QDRANT_URL", "http://qdrant:6333"))
|
|
qdrant.client.delete_collection(qdrant.collection_name)
|
|
qdrant._ensure_collection()
|
|
|
|
books = []
|
|
for directory in book_dirs():
|
|
for name in sorted(os.listdir(directory)):
|
|
if name.lower().endswith(".pdf"):
|
|
books.append(os.path.join(directory, name))
|
|
|
|
print(f"Found {len(books)} books")
|
|
os.makedirs(os.getenv("INDEX_DIR", "/app/index"), exist_ok=True)
|
|
|
|
for book_path in books:
|
|
book_name = os.path.basename(book_path)
|
|
book_metadata = {
|
|
"title": book_name.replace(".pdf", "").replace("-", " ").replace("_", " "),
|
|
"author": "Desconhecido",
|
|
"type": "book",
|
|
}
|
|
print(f"Processing: {book_name}")
|
|
chunks = pdf_processor.load_and_split(book_path, book_metadata)
|
|
chunks_with_embeddings = pdf_processor.embed_chunks(chunks)
|
|
for chunk in chunks_with_embeddings:
|
|
chunk["id"] = str(uuid.uuid4())
|
|
|
|
qdrant.add_chunks(chunks_with_embeddings)
|
|
|
|
chunks_file = os.path.join(os.getenv("INDEX_DIR", "/app/index"), f"{book_name}.json")
|
|
with open(chunks_file, "w", encoding="utf-8") as handle:
|
|
json.dump(
|
|
[
|
|
{"content": chunk["content"], "metadata": chunk["metadata"]}
|
|
for chunk in chunks_with_embeddings
|
|
],
|
|
handle,
|
|
ensure_ascii=False,
|
|
)
|
|
print(f"Indexed {len(chunks_with_embeddings)} chunks from {book_name}")
|
|
|
|
print("Indexing complete!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|