First commit

Dockerfile

@@ -0,0 +1,30 @@
+FROM python:3.10.8
+
+WORKDIR /app
+
+COPY ./app/requirements.txt /app/requirements.txt
+
+# Create cache and milvus_data directories and set permissions
+RUN mkdir -p /app/cache /app/milvus_data && chmod -R 777 /app/cache /app/milvus_data
+
+# Install dependencies
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+USER user
+
+# Set environment variables for Hugging Face cache and Milvus data
+ENV HF_HOME=/app/cache \
+    HF_MODULES_CACHE=/app/cache/hf_modules \
+    MILVUS_DATA_DIR=/app/milvus_data \
+    HF_WORKER_COUNT=1
+
+# Copy the application code
+COPY ./app /app
+
+# Expose the port Uvicorn will run on
+EXPOSE 7860
+
+# Start Uvicorn
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

Dockerfile.milvus

@@ -0,0 +1,32 @@
+FROM python:3.11
+
+# Install Milvus dependencies
+USER root
+RUN apt-get update && apt-get install -y wget ffmpeg libsm6 libxext6 libaio1
+
+# Download and install Milvus
+RUN wget https://github.com/milvus-io/milvus/releases/download/v2.3.7/milvus_2.3.7-1_amd64.deb && \
+    dpkg -i milvus_2.3.7-1_amd64.deb && \
+    apt-get -f install && \
+    apt-get clean && \
+    rm milvus_2.3.7-1_amd64.deb
+
+# Create a directory for Milvus data
+RUN mkdir -p /milvus/data
+
+# Set up Milvus user
+RUN useradd -m -u 1000 milvus
+USER milvus
+
+# Set Milvus environment variables
+ENV MILVUS_HOME=/home/milvus
+ENV PATH=$MILVUS_HOME/bin:$PATH
+
+# Set working directory
+WORKDIR $MILVUS_HOME
+
+# Expose Milvus ports
+EXPOSE 19530
+
+# Start Milvus server
+CMD ["milvus", "run", "standalone", "-d", "/milvus/data"]

app/main.py

@@ -0,0 +1,135 @@
+from io import BytesIO
+from fastapi import FastAPI, Form, Depends, Request, File, UploadFile
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from pymilvus import connections, utility, Collection, CollectionSchema, FieldSchema, DataType
+import os
+import pypdf
+from uuid import uuid4
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from sentence_transformers import SentenceTransformer
+import torch
+from app.milvus_singleton import MilvusClientSingleton
+
+# Set environment variables for Hugging Face cache
+os.environ['HF_HOME'] = '/app/cache'
+os.environ['HF_MODULES_CACHE'] = '/app/cache/hf_modules'
+
+# Embedding model
+embedding_model = SentenceTransformer('Alibaba-NLP/gte-large-en-v1.5',
+                                       trust_remote_code=True,
+                                       device='cuda' if torch.cuda.is_available() else 'cpu',
+                                      cache_folder='/app/cache')
+
+# Milvus connection details
+collection_name="rag"
+milvus_uri = os.getenv("MILVUS_URI", "sqlite:///$MILVUS_DATA_DIR/milvus_demo.db")
+
+# Initialize Milvus client using singleton
+milvus_client = MilvusClientSingleton.get_instance(uri=milvus_uri)
+
+def document_to_embeddings(content:str) -> list:
+    return embedding_model.encode(content, show_progress_bar=True)
+
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Replace with allowed origins for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+def split_documents(document_data):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=10)
+    return splitter.split_text(document_data)
+
+def create_a_collection(milvus_client, collection_name):
+    # Define the fields for the collection
+    id_field = FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=40, is_primary=True)
+    content_field = FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=4096)
+    vector_field = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024)
+    # Define the schema for the collection
+    schema = CollectionSchema(fields=[id_field, content_field, vector_field])
+    # Create the collection
+    milvus_client.create_collection(
+        collection_name=collection_name,
+        schema=schema
+    )
+    connections.connect(uri=milvus_uri)
+    collection = Collection(name=collection_name)
+    # Create an index for the collection
+    # IVF_FLAT index is used here, with metric_type COSINE
+    index_params = {
+        "index_type": "FLAT",
+        "metric_type": "COSINE",
+        "params": {
+            "nlist": 128
+        }
+    }
+    # Create the index on the vector field
+    collection.create_index(
+        field_name="vector",
+        index_params=index_params
+    )
+
+@app.get("/")
+async def root():
+    return {"message": "Hello World"}
+
+@app.post("/insert")
+async def insert(file: UploadFile = File(...)):
+    contents = await file.read()
+    if not milvus_client.has_collection(collection_name):
+        create_a_collection(milvus_client, collection_name)
+    contents = pypdf.PdfReader(BytesIO(contents))
+    extracted_text = ""
+    for page_num in range(len(contents.pages)):
+        page = contents.pages[page_num]
+        extracted_text += page.extract_text()
+        splitted_document_data = split_documents(extracted_text)
+    print(splitted_document_data)
+    data_objects = []
+    for doc in splitted_document_data:
+        data = {
+            "id": str(uuid4()),
+            "vector": document_to_embeddings(doc),
+            "content": doc,
+        }
+        data_objects.append(data)
+    print(data_objects)
+    try:
+        milvus_client.insert(collection_name=collection_name, data=data_objects)
+    except Exception as e:
+        raise JSONResponse(status_code=500, content={"error": str(e)})
+    else:
+        return JSONResponse(status_code=200, content={"result": 'good'})
+
+class RAGRequest(BaseModel):
+    question: str
+
+@app.post("/rag")
+async def rag(request: RAGRequest):
+    question = request.question
+    if not question:
+        return JSONResponse(status_code=400, content={"message": "Please a question!"})
+    try:
+        search_res = milvus_client.search(
+            collection_name=collection_name,
+            data=[
+                document_to_embeddings(question)
+            ],
+             limit=5,  # Return top 3 results
+            # search_params={"metric_type": "COSINE"},  # Inner product distance
+            output_fields=["content"],  # Return the text field
+        )
+        retrieved_lines_with_distances = [
+            (res["entity"]["content"]) for res in search_res[0]
+        ]
+        return JSONResponse(status_code=200, content={"result": retrieved_lines_with_distances})
+    except Exception as e:
+        return JSONResponse(status_code=400, content={"error": str(e)})

app/milvus_singleton.py

@@ -0,0 +1,25 @@
+from pymilvus import Milvus, connections
+from pymilvus.exceptions import ConnectionConfigException
+
+class MilvusClientSingleton:
+    _instance = None
+
+    @staticmethod
+    def get_instance(uri):
+        if MilvusClientSingleton._instance is None:
+            MilvusClientSingleton(uri)
+        return MilvusClientSingleton._instance
+
+    def __init__(self, uri):
+        if MilvusClientSingleton._instance is not None:
+            raise Exception("This class is a singleton!")
+        try:
+            # Use the regular Milvus client (not MilvusClient)
+            self._instance = Milvus(uri=uri)
+            print(f"Successfully connected to Milvus at {uri}")
+        except ConnectionConfigException as e:
+            print(f"Error connecting to Milvus: {e}")
+            raise  # Re-raise the exception to stop initialization
+
+    def __getattr__(self, name):
+        return getattr(self._instance, name)

app/requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn
+pypdf
+python-multipart
+langchain
+sentence-transformers
+torch
+pymilvu