diff --git a/.gitignore b/.gitignore index 4649d05a..73631b69 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ uploads/ myenv/ venv/ *.pyc +env/ +env2/ diff --git a/.vscode/launch.json b/.vscode/launch.json index 2c241927..2e0640fb 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -13,6 +13,19 @@ "--reload" ], "jinja": true - } + }, + { + "name": "Debug FastAPI (debugpy)", + "type": "debugpy", + "request": "launch", + "module": "uvicorn", + "args": [ + "main:app", + "--port", + "8000", + "--reload" + ], + "jinja": true + } ] } diff --git a/config.py b/config.py index e4263612..dbbaf8a0 100644 --- a/config.py +++ b/config.py @@ -4,6 +4,7 @@ import logging from enum import Enum from datetime import datetime +from threading import TIMEOUT_MAX from dotenv import find_dotenv, load_dotenv from langchain_community.embeddings import ( HuggingFaceEmbeddings, @@ -12,7 +13,7 @@ ) from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings from starlette.middleware.base import BaseHTTPMiddleware -from store_factory import get_vector_store +from store_factory import QdrantConfig, get_vector_store load_dotenv(find_dotenv()) @@ -20,6 +21,7 @@ class VectorDBType(Enum): PGVECTOR = "pgvector" ATLAS_MONGO = "atlas-mongo" + QDRANT = "qdrant" class EmbeddingsProvider(Enum): @@ -63,9 +65,22 @@ def get_env_variable( MONGO_VECTOR_COLLECTION = get_env_variable( "MONGO_VECTOR_COLLECTION", "vector_collection" ) +QDRANT_HOST= get_env_variable("QDRANT_HOST", "127.0.0.1:6333") +COLLECTION_NAME = get_env_variable( + "COLLECTION_NAME", "vector_collection" +) +EMBEDDINGS_DIMENSION = get_env_variable("EMBEDDINGS_DIMENSION", "768") +QDRANT_API_KEY = get_env_variable("QDRANT_API_KEY") CHUNK_SIZE = int(get_env_variable("CHUNK_SIZE", "1500")) CHUNK_OVERLAP = int(get_env_variable("CHUNK_OVERLAP", "100")) +maxChunks = get_env_variable("MAX_CHUNKS") +MAX_CHUNKS = int(maxChunks) if maxChunks else None +embeddingTimeout = get_env_variable("EMBEDDING_TIMEOUT") +EMBEDDING_TIMEOUT = int(get_env_variable("EMBEDDING_TIMEOUT",100000))#default 100 second timeout + +BATCH_SIZE = int(get_env_variable("BATCH_SIZE","75")) +CONCURRENT_LIMIT = int(get_env_variable("CONCURRENT_LIMIT","10")) env_value = get_env_variable("PDF_EXTRACT_IMAGES", "False").lower() PDF_EXTRACT_IMAGES = True if env_value == "true" else False @@ -228,12 +243,13 @@ def init_embeddings(provider, model): logger.info(f"Initialized embeddings of type: {type(embeddings)}") # Vector store +print(VECTOR_DB_TYPE) if VECTOR_DB_TYPE == VectorDBType.PGVECTOR: vector_store = get_vector_store( connection_string=CONNECTION_STRING, embeddings=embeddings, collection_name=COLLECTION_NAME, - mode="async", + mode="async-PGVector", ) elif VECTOR_DB_TYPE == VectorDBType.ATLAS_MONGO: logger.warning("Using Atlas MongoDB as vector store is not fully supported yet.") @@ -243,9 +259,21 @@ def init_embeddings(provider, model): collection_name=MONGO_VECTOR_COLLECTION, mode="atlas-mongo", ) +elif VECTOR_DB_TYPE == VectorDBType.QDRANT: + vector_store = get_vector_store( + connection_string=QDRANT_HOST, + embeddings=embeddings, + collection_name=COLLECTION_NAME, + mode="qdrant", + additional_kwargs = QdrantConfig( + qdrant_api_key=QDRANT_API_KEY, + embeddings_dimension=EMBEDDINGS_DIMENSION + ) +) else: raise ValueError(f"Unsupported vector store type: {VECTOR_DB_TYPE}") + retriever = vector_store.as_retriever() known_source_ext = [ diff --git a/main.py b/main.py index c82b550a..1811a49e 100644 --- a/main.py +++ b/main.py @@ -1,4 +1,5 @@ import os +import time import hashlib import aiofiles import aiofiles.os @@ -12,6 +13,7 @@ from fastapi.middleware.cors import CORSMiddleware from langchain_core.runnables.config import run_in_executor from langchain.text_splitter import RecursiveCharacterTextSplitter +from uuid import uuid1 from fastapi import ( File, Form, @@ -49,7 +51,8 @@ from middleware import security_middleware from mongo import mongo_health_check from constants import ERROR_MESSAGES -from store import AsyncPgVector +from store import AsyncPgVector, AsyncQdrant, async_DB +from process_docs import store_documents load_dotenv(find_dotenv()) @@ -58,6 +61,10 @@ debug_mode, CHUNK_SIZE, CHUNK_OVERLAP, + EMBEDDING_TIMEOUT, + MAX_CHUNKS, + BATCH_SIZE, + CONCURRENT_LIMIT, vector_store, RAG_UPLOAD_DIR, known_source_ext, @@ -65,11 +72,10 @@ LogMiddleware, RAG_HOST, RAG_PORT, - VectorDBType, + VECTOR_DB_TYPE, # RAG_EMBEDDING_MODEL, # RAG_EMBEDDING_MODEL_DEVICE_TYPE, # RAG_TEMPLATE, - VECTOR_DB_TYPE, ) @@ -85,6 +91,7 @@ async def lifespan(app: FastAPI): app = FastAPI(lifespan=lifespan) + app.add_middleware( CORSMiddleware, allow_origins=["*"], @@ -105,7 +112,7 @@ async def lifespan(app: FastAPI): @app.get("/ids") async def get_all_ids(): try: - if isinstance(vector_store, AsyncPgVector): + if isinstance(vector_store, async_DB): ids = await vector_store.get_all_ids() else: ids = vector_store.get_all_ids() @@ -116,12 +123,7 @@ async def get_all_ids(): def isHealthOK(): - if VECTOR_DB_TYPE == VectorDBType.PGVECTOR: - return pg_health_check() - if VECTOR_DB_TYPE == VectorDBType.ATLAS_MONGO: - return mongo_health_check() - else: - return True + return pg_health_check() @app.get("/health") @@ -135,7 +137,7 @@ async def health_check(): @app.get("/documents", response_model=list[DocumentResponse]) async def get_documents_by_ids(ids: list[str] = Query(...)): try: - if isinstance(vector_store, AsyncPgVector): + if isinstance(vector_store, async_DB): existing_ids = await vector_store.get_all_ids() documents = await vector_store.get_documents_by_ids(ids) else: @@ -145,13 +147,6 @@ async def get_documents_by_ids(ids: list[str] = Query(...)): # Ensure all requested ids exist if not all(id in existing_ids for id in ids): raise HTTPException(status_code=404, detail="One or more IDs not found") - - # Ensure documents list is not empty - if not documents: - raise HTTPException( - status_code=404, detail="No documents found for the given IDs" - ) - return documents except HTTPException as http_exc: raise http_exc @@ -160,19 +155,17 @@ async def get_documents_by_ids(ids: list[str] = Query(...)): @app.delete("/documents") -async def delete_documents(document_ids: List[str] = Body(...)): +async def delete_documents(ids: list[str]): try: - if isinstance(vector_store, AsyncPgVector): + if isinstance(vector_store, async_DB): existing_ids = await vector_store.get_all_ids() - await vector_store.delete(ids=document_ids) + await vector_store.delete(ids=ids) + if not all(id in existing_ids for id in ids): + raise HTTPException(status_code=404, detail="One or more IDs not found") else: existing_ids = vector_store.get_all_ids() - vector_store.delete(ids=document_ids) - - if not all(id in existing_ids for id in document_ids): - raise HTTPException(status_code=404, detail="One or more IDs not found") - - file_count = len(document_ids) + vector_store.delete(ids=ids) + file_count = len(ids) return { "message": f"Documents for {file_count} file{'s' if file_count > 1 else ''} deleted successfully" } @@ -188,9 +181,8 @@ async def query_embeddings_by_file_id(body: QueryRequestBody, request: Request): authorized_documents = [] try: - embedding = vector_store.embedding_function.embed_query(body.query) - - if isinstance(vector_store, AsyncPgVector): + if isinstance(vector_store, async_DB): + embedding = vector_store.embedding_function.embed_query(body.query) documents = await run_in_executor( None, vector_store.similarity_search_with_score_by_vector, @@ -199,6 +191,7 @@ async def query_embeddings_by_file_id(body: QueryRequestBody, request: Request): filter={"file_id": body.file_id}, ) else: + embedding = vector_store.embedding_function.embed_query(body.query) documents = vector_store.similarity_search_with_score_by_vector( embedding, k=body.k, filter={"file_id": body.file_id} ) @@ -234,12 +227,18 @@ async def store_data_in_vector_db( file_id: str, user_id: str = "", clean_content: bool = False, -) -> bool: +) -> dict: text_splitter = RecursiveCharacterTextSplitter( chunk_size=app.state.CHUNK_SIZE, chunk_overlap=app.state.CHUNK_OVERLAP ) + ids = [] documents = text_splitter.split_documents(data) + if MAX_CHUNKS and len(documents) > MAX_CHUNKS: + raise HTTPException( + status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + detail=f"Too big file, Attempted to process {len(documents)} chunks, but MAX_CHUNKS is set to {MAX_CHUNKS} chunks", + ) # If `clean_content` is True, clean the page_content of each document (remove null bytes) if clean_content: for doc in documents: @@ -258,21 +257,21 @@ async def store_data_in_vector_db( ) for doc in documents ] + try: - if isinstance(vector_store, AsyncPgVector): - ids = await vector_store.aadd_documents( - docs, ids=[file_id] * len(documents) - ) + uuids = [str(uuid1()) for _ in docs] + if isinstance(vector_store, async_DB): + ids = await store_documents(docs, ids=uuids) else: - ids = vector_store.add_documents(docs, ids=[file_id] * len(documents)) + ids = vector_store.add_documents(docs, ids=uuids) return {"message": "Documents added successfully", "ids": ids} except Exception as e: logger.error(e) return {"message": "An error occurred while adding documents.", "error": str(e)} - + def get_loader(filename: str, file_content_type: str, filepath: str): file_ext = filename.split(".")[-1].lower() @@ -303,8 +302,6 @@ def get_loader(filename: str, file_content_type: str, filepath: str): "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", ] or file_ext in ["xls", "xlsx"]: loader = UnstructuredExcelLoader(filepath) - elif file_content_type == "application/json" or file_ext == "json": - loader = TextLoader(filepath, autodetect_encoding=True) elif file_ext in known_source_ext or ( file_content_type and file_content_type.find("text/") >= 0 ): @@ -375,7 +372,6 @@ async def embed_file( user_id = "public" else: user_id = request.state.user.get("id") - temp_base_path = os.path.join(RAG_UPLOAD_DIR, user_id) os.makedirs(temp_base_path, exist_ok=True) temp_file_path = os.path.join(RAG_UPLOAD_DIR, user_id, file.filename) @@ -443,7 +439,7 @@ async def embed_file( async def load_document_context(id: str): ids = [id] try: - if isinstance(vector_store, AsyncPgVector): + if isinstance(vector_store, async_DB): existing_ids = await vector_store.get_all_ids() documents = await vector_store.get_documents_by_ids(ids) else: @@ -531,14 +527,20 @@ async def query_embeddings_by_file_ids(body: QueryMultipleBody): if isinstance(vector_store, AsyncPgVector): documents = await run_in_executor( None, - vector_store.similarity_search_with_score_by_vector, + vector_store.asimilarity_search_with_score_by_vector, embedding, k=body.k, - filter={"file_id": {"$in": body.file_ids}}, + filter={"custom_id": {"$in": body.file_ids}}, ) + elif isinstance(vector_store, AsyncQdrant): + documents = await vector_store.similarity_search_many( + embedding=embedding, + k=body.k, + ids=body.file_ids) + else: documents = vector_store.similarity_search_with_score_by_vector( - embedding, k=body.k, filter={"file_id": {"$in": body.file_ids}} + embedding, k=body.k, filter={"custom_id": {"$in": body.file_ids}} ) # Ensure documents list is not empty @@ -556,4 +558,4 @@ async def query_embeddings_by_file_ids(body: QueryMultipleBody): app.include_router(router=pgvector_router) if __name__ == "__main__": - uvicorn.run(app, host=RAG_HOST, port=RAG_PORT, log_config=None) + uvicorn.run(app, host=RAG_HOST, port=RAG_PORT, log_config=None) \ No newline at end of file diff --git a/process_docs.py b/process_docs.py new file mode 100644 index 00000000..49d1536a --- /dev/null +++ b/process_docs.py @@ -0,0 +1,33 @@ +from langchain.schema import Document +import asyncio +from config import ( + CONCURRENT_LIMIT, + BATCH_SIZE, + EMBEDDING_TIMEOUT, + vector_store +) + + +#Prepare documents to be async added to vectorstore async in batches +async def store_documents( + docs: list[Document], ids:list[str] +): + semaphore = asyncio.Semaphore(CONCURRENT_LIMIT) + tasks = [] + #logger.info(f"Processing list of documents of length: {len(docs)}") + for i in range(0, len(docs), BATCH_SIZE): + batch = docs[i : min(i + BATCH_SIZE, len(docs))] + #logger.info(f"Sending batch {i} to {i+len(batch)} / {len(docs)}") + task = asyncio.create_task(process_batch(batch, ids, semaphore)) + tasks.append(task) + try: + idList = await asyncio.wait_for(asyncio.gather(*tasks), timeout=(EMBEDDING_TIMEOUT/1000)) + except asyncio.TimeoutError: + raise Exception(f"TIMEOUT: embedding process took over the time limit of {EMBEDDING_TIMEOUT}ms. Partially added to database") + return [id for sublist in idList for id in sublist] + +#Helper for process_documents +async def process_batch(batch: list[Document], ids: list[str], semaphore): + async with semaphore: + return await vector_store.aadd_documents(batch,ids=ids) + diff --git a/requirements.txt b/requirements.txt index edf53ce5..e7794821 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,8 +23,10 @@ sentence_transformers==2.5.1 aiofiles==23.2.1 rapidocr-onnxruntime==1.3.17 opencv-python-headless==4.9.0.80 +qdrant-client==1.9.1 pymongo==4.6.3 langchain-mongodb==0.1.3 cryptography==42.0.7 +qdrant-client==1.9.1 python-magic==0.4.27 python-pptx==0.6.23 diff --git a/store.py b/store.py index f6fd1e0e..3f74ae61 100644 --- a/store.py +++ b/store.py @@ -1,9 +1,15 @@ from typing import Any, Optional -from sqlalchemy import delete +from aiohttp import Payload +from sqlalchemy import delete, func from langchain_community.vectorstores.pgvector import PGVector from langchain_core.documents import Document from langchain_core.runnables.config import run_in_executor from sqlalchemy.orm import Session +from langchain_community.vectorstores import Qdrant +import qdrant_client as client +from qdrant_client.http import models + + from langchain_mongodb import MongoDBAtlasVectorSearch from langchain_core.embeddings import Embeddings @@ -14,25 +20,24 @@ ) import copy - class ExtendedPgVector(PGVector): + def get_all_ids(self) -> list[str]: with Session(self._bind) as session: - results = session.query(self.EmbeddingStore.custom_id).all() - return [result[0] for result in results if result[0] is not None] + results = session.query( + func.json_extract_path_text(self.EmbeddingStore.cmetadata,'file_id').label('file_id') + ).distinct().all() + return [id.file_id for id in results] def get_documents_by_ids(self, ids: list[str]) -> list[Document]: with Session(self._bind) as session: - results = ( - session.query(self.EmbeddingStore) - .filter(self.EmbeddingStore.custom_id.in_(ids)) - .all() - ) - return [ - Document(page_content=result.document, metadata=result.cmetadata or {}) - for result in results - if result.custom_id in ids + results = session.query(self.EmbeddingStore).filter( + func.json_extract_path_text(self.EmbeddingStore.cmetadata, 'file_id').in_(ids) + ).all() + return [ + Document(page_content=result.document, metadata=result.cmetadata or {}) + for result in results ] def _delete_multiple( @@ -42,7 +47,7 @@ def _delete_multiple( if ids is not None: self.logger.debug( "Trying to delete vectors by ids (represented by the model " - "using the custom ids field)" + "using the custom file_id field in cmetadata)" ) stmt = delete(self.EmbeddingStore) @@ -56,8 +61,8 @@ def _delete_multiple( stmt = stmt.where( self.EmbeddingStore.collection_id == collection.uuid ) - - stmt = stmt.where(self.EmbeddingStore.custom_id.in_(ids)) + + stmt = stmt.where(func.json_extract_path_text(self.EmbeddingStore.cmetadata, 'file_id').in_(ids)) session.execute(stmt) session.commit() @@ -71,11 +76,109 @@ async def get_documents_by_ids(self, ids: list[str]) -> list[Document]: return await run_in_executor(None, super().get_documents_by_ids, ids) async def delete( - self, ids: Optional[list[str]] = None, collection_only: bool = False - ) -> None: - await run_in_executor(None, self._delete_multiple, ids, collection_only) + self, + ids: Optional[list[str]] = None, + collection_only: bool = False + ) -> None: + await run_in_executor(None, self._delete_multiple, ids, collection_only) + +class ExtendedQdrant(Qdrant): + @property + def embedding_function(self) -> Embeddings: + return self.embeddings + + def delete_vectors_by_source_document(self, source_document_ids: list[str]) -> None: + points_selector = models.Filter( + must=[ + models.FieldCondition( + key="metadata.file_id", + match=models.MatchAny(any=source_document_ids), + ), + ], + ) + response = self.client.delete(collection_name=self.collection_name, points_selector=points_selector) + status = response.status.name + return status + + + def get_all_ids(self) -> list[str]: + collection_info = self.client.get_collection(self.collection_name) + total_points = collection_info.points_count + + # Scroll through all points in the collection + unique_values = set() + pointsRec = 0 + limit = 500 #How much to load each time + next_offset = None + while pointsRec < total_points: + points, next_offset = self.client.scroll( + collection_name=self.collection_name, + limit=limit, + with_payload=True, + offset=next_offset, + ) + for point in points: + unique_values.add(point.payload['metadata']['file_id']) + pointsRec += limit + return list(unique_values) + + def get_documents_by_ids(self, ids: list[str]): + filter = models.Filter( + must=[ + models.FieldCondition( + key="metadata.file_id", + match=models.MatchAny(any=ids) + ) + ] + ) + limit = 500 + next_offset = None + docList = [] + while True: + results = self.client.scroll( + collection_name=self.collection_name, + scroll_filter=filter, + limit=limit, + offset=next_offset + ) + points, next_offset = results + if points: + docList.extend([Document(page_content=point.payload['page_content'], + metadata=point.payload['metadata'] + ) + for point in points + ]) + if not next_offset: + break + return docList + + +class AsyncQdrant(ExtendedQdrant): + + async def get_all_ids(self) -> list[str]: + return await run_in_executor(None, super().get_all_ids) + async def get_documents_by_ids(self, ids: list[str]) -> list[Document]: + return await run_in_executor(None, super().get_documents_by_ids, ids) + async def delete( + self, + ids: list[str] + ) -> None: + # Garantir que o argumento correto está sendo passado + await run_in_executor(None, self.delete_vectors_by_source_document, ids) + + async def similarity_search_many(self, embedding, k:int, ids:list[str])-> List[Tuple[Document, float]]: + filter = models.Filter( + must=[ + models.FieldCondition( + key="metadata.file_id", + match=models.MatchAny(any=ids) + ) + ] + ) + results = await self.asimilarity_search_with_score_by_vector(embedding=embedding, k=k, filter=filter) + return results class AtlasMongoVector(MongoDBAtlasVectorSearch): @property def embedding_function(self) -> Embeddings: @@ -142,3 +245,6 @@ def delete(self, ids: Optional[list[str]] = None) -> None: # implement the deletion of documents by file_id in self._collection if ids is not None: self._collection.delete_many({"file_id": {"$in": ids}}) + + +async_DB = (AsyncPgVector, AsyncQdrant) #Add if async database implementation diff --git a/store_factory.py b/store_factory.py index 16b81ef6..9d1be44f 100644 --- a/store_factory.py +++ b/store_factory.py @@ -1,27 +1,69 @@ +from typing import Optional, TypedDict from langchain_community.embeddings import OpenAIEmbeddings - -from store import AsyncPgVector, ExtendedPgVector -from store import AtlasMongoVector from pymongo import MongoClient +import qdrant_client +from store import AsyncPgVector, ExtendedPgVector, AsyncQdrant, AtlasMongoVector + + +class QdrantConfig(TypedDict, total=False): + qdrant_api_key: Optional[str] + embeddings_dimension: Optional[int] def get_vector_store( connection_string: str, embeddings: OpenAIEmbeddings, collection_name: str, - mode: str = "sync", + mode: str = "sync-PGVector", + *, + additional_kwargs: Optional[QdrantConfig] = None ): - if mode == "sync": + + if mode == "sync-PGVector": return ExtendedPgVector( connection_string=connection_string, embedding_function=embeddings, collection_name=collection_name, ) - elif mode == "async": + elif mode == "async-PGVector": return AsyncPgVector( connection_string=connection_string, embedding_function=embeddings, collection_name=collection_name, ) + elif mode == "atlas-mongo": + mongo_db = MongoClient(connection_string).get_database() + mong_collection = mongo_db[collection_name] + return AtlasMongoVector(collection=mong_collection, embedding=embeddings) + + elif mode == "qdrant": + if additional_kwargs is None: + additional_kwargs = {} + qdrant_api_key = additional_kwargs['qdrant_api_key'] + qdrant_embeddings_dimension = additional_kwargs['embeddings_dimension'] + client = qdrant_client.QdrantClient( + url=connection_string, + api_key=qdrant_api_key + ) + collection_config = qdrant_client.http.models.VectorParams( + size=qdrant_embeddings_dimension, + distance=qdrant_client.http.models.Distance.COSINE + ) + if not client.collection_exists(collection_name): + collection_config = qdrant_client.http.models.VectorParams( + size=qdrant_embeddings_dimension, + distance=qdrant_client.http.models.Distance.COSINE + ) + client.create_collection( + collection_name=collection_name, + vectors_config=collection_config + ) + return AsyncQdrant( + client=client, + collection_name=collection_name, + embeddings=embeddings + + ) + elif mode == "atlas-mongo": mongo_db = MongoClient(connection_string).get_database() mong_collection = mongo_db[collection_name]