updates

Dockerfile

@@ -2,30 +2,25 @@ FROM python:3.10.8
 
 WORKDIR /app
 
-# Copy only requirements.txt first to leverage Docker caching
-COPY requirements.txt /app/
+COPY 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 (now main.py is at the root)
 COPY . /app
 
-# Expose the port Uvicorn will run on
+# Expose port for Uvicorn
 EXPOSE 7860
 
-# Start Uvicorn (main:app is correct now)
+# Command to run Uvicorn
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -1,44 +1,44 @@
 from io import BytesIO
-from fastapi import FastAPI, File, UploadFile
+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 utility, Collection, CollectionSchema, FieldSchema, DataType
+from pymilvus import connections
+
+
 import os
 import pypdf
 from uuid import uuid4
+
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from pymilvus import MilvusClient, db, utility, Collection, CollectionSchema, FieldSchema, DataType
 from sentence_transformers import SentenceTransformer
 import torch
 from 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 = 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'
+)
+collection_name="rag"
 
-# 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", "http://localhost:19530")  # Correct URI for Milvus
 
-# Initialize Milvus client using singleton
-milvus_client = MilvusClientSingleton.get_instance(uri=milvus_uri)
+# milvus_client = MilvusClientSingleton.get_instance(uri="/app/milvus_data/milvus_demo.db")
+milvus_client = MilvusClient(uri="/app/milvus_data/milvus_demo.db")
 
-def document_to_embeddings(content: str) -> list:
+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_origins=["*"],  # Replace with the list of allowed origins for production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -53,14 +53,20 @@ def create_a_collection(milvus_client, collection_name):
     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="/app/milvus_data/milvus_demo.db")
+
     collection = Collection(name=collection_name)
+
     # Create an index for the collection
     # IVF_FLAT index is used here, with metric_type COSINE
     index_params = {
@@ -70,10 +76,11 @@ def create_a_collection(milvus_client, collection_name):
             "nlist": 128
         }
     }
+
     # Create the index on the vector field
     collection.create_index(
         field_name="vector",
-        index_params=index_params
+        index_params=index_params  # Pass the dictionary, not a string
     )
 
 @app.get("/")
@@ -83,15 +90,21 @@ async def root():
 @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)
+    
+    splitted_document_data = split_documents(extracted_text)
+
     print(splitted_document_data)
+
     data_objects = []
     for doc in splitted_document_data:
         data = {
@@ -100,32 +113,37 @@ async def insert(file: UploadFile = File(...)):
             "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 5 results
+            ], 
+            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]
         ]

milvus_singleton.py

@@ -7,21 +7,16 @@ class MilvusClientSingleton:
     @staticmethod
     def get_instance(uri):
         if MilvusClientSingleton._instance is None:
-            MilvusClientSingleton(uri)
+            MilvusClientSingleton()
+            # Initialize the client here
+            try:
+                MilvusClientSingleton._instance = connections.connect(uri=uri)
+            except ConnectionConfigException as e:
+                print(f"Error connecting to Milvus: {e}")
+                # Handle error appropriately
         return MilvusClientSingleton._instance
 
-    def __init__(self, uri):
+    def __init__(self):
         if MilvusClientSingleton._instance is not None:
             raise Exception("This class is a singleton!")
-        try:
-            # Use connections.connect() to establish the connection
-            connections.connect(uri=uri)
-            self._instance = connections  # Store the connections object
-            print(f"Successfully connected to Milvus at {uri}")
-        except ConnectionConfigException as e:
-            print(f"Error connecting to Milvus: {e}")
-            raise
-
-    def __getattr__(self, name):
-        # Delegate attribute access to the default connection
-        return getattr(connections, name)
+        self._instance = None