quivr/backend/parsers/common.py

46 lines
1.3 KiB
Python
Raw Normal View History

2023-05-18 02:22:13 +03:00
import time
from langchain.schema import Document
2023-06-28 20:39:27 +03:00
from models.brains import Brain
from models.files import File
2023-06-19 23:46:25 +03:00
from models.settings import CommonsDep
2023-06-28 20:39:27 +03:00
from utils.vectors import Neurons
2023-05-18 02:22:13 +03:00
2023-05-22 09:39:55 +03:00
async def process_file(
commons: CommonsDep,
2023-06-28 20:39:27 +03:00
file: File,
loader_class,
enable_summarization,
2023-06-28 20:39:27 +03:00
brain_id,
user_openai_api_key,
):
2023-05-18 02:22:13 +03:00
dateshort = time.strftime("%Y%m%d")
2023-05-22 09:39:55 +03:00
2023-06-28 20:39:27 +03:00
file.compute_documents(loader_class)
2023-05-22 09:39:55 +03:00
2023-06-28 20:39:27 +03:00
for doc in file.documents:
2023-05-22 09:39:55 +03:00
metadata = {
2023-06-28 20:39:27 +03:00
"file_sha1": file.file_sha1,
"file_size": file.file_size,
"file_name": file.file_name,
"chunk_size": file.chunk_size,
"chunk_overlap": file.chunk_overlap,
2023-05-22 09:39:55 +03:00
"date": dateshort,
"summarization": "true" if enable_summarization else "false",
2023-05-22 09:39:55 +03:00
}
doc_with_metadata = Document(
page_content=doc.page_content, metadata=metadata)
2023-06-29 19:26:03 +03:00
2023-06-19 22:15:35 +03:00
neurons = Neurons(commons=commons)
2023-06-28 20:39:27 +03:00
created_vector = neurons.create_vector(doc_with_metadata, user_openai_api_key)
# add_usage(stats_db, "embedding", "audio", metadata={"file_name": file_meta_name,"file_type": ".txt", "chunk_size": chunk_size, "chunk_overlap": chunk_overlap})
2023-06-28 20:39:27 +03:00
created_vector_id = created_vector[0]
2023-05-22 09:39:55 +03:00
2023-06-28 20:39:27 +03:00
brain = Brain(id=brain_id)
2023-06-29 19:26:03 +03:00
brain.create_brain_vector(created_vector_id, file.file_sha1)
2023-06-06 01:38:15 +03:00
2023-06-28 20:39:27 +03:00
return