added qudrant inside the container

Dockerfile

@@ -4,12 +4,18 @@ RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
 
+# Install Qdrant
+RUN pip install qdrant-client && \
+    apt-get update && apt-get install -y qdrant
+
 WORKDIR /app
 
 COPY --chown=user * ./
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["sh", "-c", "qdrant & uvicorn main:app --host 0.0.0.0 --port 7860"]
+
 
 

dependencies.py

@@ -12,4 +12,14 @@ async def get_urls_content(body: TextModelRequest ) -> str:
             return urls_content
         except Exception as e:
             logger.warning(f"Failed to fetch one or several URls - Error: {e}")
-    return ""
+    return ""
+
+async def get_rag_content(body: TextModelRequest ) -> str: 
+    rag_content = await vector_service.search( 
+        "knowledgebase", embed(body.prompt), 3, 0.7
+    )
+    rag_content_str = "\n".join( 
+        [c.payload["original_text"] for c in rag_content]
+    )
+
+    return rag_content_str

main.py

@@ -1,5 +1,5 @@
 # main.py
-from fastapi import FastAPI,status,Response,Request,Depends,HTTPException,UploadFile, File
+from fastapi import FastAPI,status,Response,Request,Depends,HTTPException,UploadFile, File,BackgroundTasks
 from fastapi.responses import StreamingResponse,FileResponse
 from models import load_text_model,generate_text,load_audio_model,generate_audio,load_image_model, generate_image
 from schemas import VoicePresets
@@ -14,6 +14,8 @@ from dependencies import get_urls_content
 from schemas import TextModelResponse,TextModelRequest
 import shutil, uuid
 from upload import save_file
+from rag import pdf_text_extractor, vector_service
+
 
 models = {} 
 
@@ -115,17 +117,42 @@ def serve_text_to_image_model_controller(prompt: str):
     output = generate_image(models["text2image"], prompt)
     return Response(content=img_to_bytes(output), media_type="image/png") 
 
+# @app.post("/upload")
+# async def file_upload_controller(
+#     file: Annotated[UploadFile, File(description="Uploaded PDF documents")]
+# ):
+#     if file.content_type != "application/pdf":
+#         raise HTTPException(
+#             detail=f"Only uploading PDF documents are supported",
+#             status_code=status.HTTP_400_BAD_REQUEST,
+#         )
+#     try:
+#         await save_file(file)
+#     except Exception as e:
+#         raise HTTPException(
+#             detail=f"An error occurred while saving file - Error: {e}",
+#             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+#         )
+#     return {"filename": file.filename, "message": "File uploaded successfully"}
+
+
 @app.post("/upload")
 async def file_upload_controller(
-    file: Annotated[UploadFile, File(description="Uploaded PDF documents")]
+    file: Annotated[UploadFile, File(description="A file read as UploadFile")],
+    bg_text_processor: BackgroundTasks, 
 ):
-    if file.content_type != "application/pdf":
-        raise HTTPException(
-            detail=f"Only uploading PDF documents are supported",
-            status_code=status.HTTP_400_BAD_REQUEST,
-        )
+    ... # Raise an HTTPException if data upload is not a PDF file
     try:
-        await save_file(file)
+        filepath = await save_file(file)
+        bg_text_processor.add_task(pdf_text_extractor, filepath) 
+        bg_text_processor.add_task( 
+            vector_service.store_file_content_in_db,
+            filepath.replace("pdf", "txt"),
+            512,
+            "knowledgebase",
+            768,
+        )
+
     except Exception as e:
         raise HTTPException(
             detail=f"An error occurred while saving file - Error: {e}",

rag/__init__.py

@@ -0,0 +1,4 @@
+from .extractor import *
+from .repository import *
+from .service import *
+from .transform import *

rag/repository.py

@@ -0,0 +1,84 @@
+from loguru import logger
+from qdrant_client import AsyncQdrantClient
+from qdrant_client.http import models
+from qdrant_client.http.models import ScoredPoint
+
+class VectorRepository: 
+    def __init__(self, host: str = "https://ahmed-eisa-qdrant-db.hf.space", port: int = 6333) -> None:
+        self.db_client = AsyncQdrantClient(host=host, port=port)
+
+    async def create_collection(self, collection_name: str, size: int) -> bool: 
+        vectors_config = models.VectorParams(
+            size=size, distance=models.Distance.COSINE 
+        )
+        response = await self.db_client.get_collections()
+
+        collection_exists = any(
+            collection.name == collection_name
+            for collection in response.collections
+        )
+        if collection_exists: 
+            logger.debug(
+                f"Collection {collection_name} already exists - recreating it"
+            )
+            await self.db_client.delete_collection(collection_name)
+            return await self.db_client.create_collection(
+                collection_name,
+                vectors_config=vectors_config,
+            )
+
+        logger.debug(f"Creating collection {collection_name}")
+        return await self.db_client.create_collection(
+            collection_name=collection_name,
+            vectors_config=models.VectorParams(
+                size=size, distance=models.Distance.COSINE
+            ),
+        )
+
+    async def delete_collection(self, name: str) -> bool:
+        logger.debug(f"Deleting collection {name}")
+        return await self.db_client.delete_collection(name)
+
+    async def create(
+        self,
+        collection_name: str,
+        embedding_vector: list[float],
+        original_text: str,
+        source: str,
+    ) -> None:
+        response = await self.db_client.count(collection_name=collection_name)
+        logger.debug(
+            f"Creating a new vector with ID {response.count} "
+            f"inside the {collection_name}"
+        )
+        await self.db_client.upsert(
+            collection_name=collection_name,
+            points=[
+                models.PointStruct(
+                    id=response.count,
+                    vector=embedding_vector,
+                    payload={
+                        "source": source,
+                        "original_text": original_text,
+                    },
+                )
+            ],
+        )
+
+    async def search(
+        self,
+        collection_name: str,
+        query_vector: list[float],
+        retrieval_limit: int,
+        score_threshold: float, 
+    ) -> list[ScoredPoint]:
+        logger.debug(
+            f"Searching for relevant items in the {collection_name} collection"
+        )
+        response = await self.db_client.query_points(
+            collection_name=collection_name,
+            query_vector=query_vector,
+            limit=retrieval_limit,
+            score_threshold=score_threshold,
+        )
+        return response.points

rag/service.py

@@ -0,0 +1,31 @@
+import os
+
+from loguru import logger
+from .repository import VectorRepository
+from .transform import clean, embed, load
+
+
+class VectorService(VectorRepository): 
+    def __init__(self):
+        super().__init__()
+
+    async def store_file_content_in_db( 
+        self,
+        filepath: str,
+        chunk_size: int = 512,
+        collection_name: str = "knowledgebase",
+        collection_size: int = 768,
+    ) -> None:
+        await self.create_collection(collection_name, collection_size)
+        logger.debug(f"Inserting {filepath} content into database")
+        async for chunk in load(filepath, chunk_size): 
+            logger.debug(f"Inserting '{chunk[0:20]}...' into database")
+
+            embedding_vector = embed(clean(chunk))
+            filename = os.path.basename(filepath)
+            await self.create(
+                collection_name, embedding_vector, chunk, filename
+            )
+
+
+vector_service = VectorService()

rag/transform.py

@@ -10,9 +10,9 @@ embedder = AutoModel.from_pretrained(
     "jinaai/jina-embeddings-v2-base-en", trust_remote_code=True 
 )
 
-async def load(filepath: str) -> AsyncGenerator[str, Any]:
+async def load(filepath: str, chunksize:int=DEFAULT_CHUNK_SIZE) -> AsyncGenerator[str, Any]:
     async with aiofiles.open(filepath, "r", encoding="utf-8") as f: 
-        while chunk := await f.read(DEFAULT_CHUNK_SIZE): 
+        while chunk := await f.read(chunksize): 
             yield chunk 
 
 def clean(text: str) -> str: