Merge pull request #142 from AET-DevOps25/feature/improve-rag-structure

Improve rag structure and add auth service

Author: edkaya

README.md

@@ -295,8 +295,12 @@ The LLM service will be available at [http://localhost:8000](http://localhost:80
 
 ##### Integration
 
-- The client UI sends user requests to the server, which forwards them to the GenAI service along with the user’s query and chat history to support multi-turn conversations. GenAI service then makes a similarity search in the vector database with the given query, and generates a respective answer. GenAI service is able to provide a proper answer altough no similar context is found in the vector database. (Endpoint: POST - `genai/generate`)
-- If the user wants to upload a recipe file, client UI sends the file content directly to the GenAI service, where the content of the file is chunked, embedded, and stored in the vector database. (Endpoint: POST - `genai/upload`)
+- The client UI sends user requests to the API gateway, which forwards them to the chat service. Chat service forwards them to the GenAI service along with the user’s query and chat history to support multi-turn conversations. GenAI service then makes a similarity search in the vector database with the given query, and generates a respective answer. GenAI service is able to provide a proper answer altough no similar context is found in the vector database. (Endpoint: POST - `genai/generate`)
+- If the user wants to upload a recipe file, client UI sends the file content directly to the API gateway, which forwards to the GenAI service, where the content of the file is chunked, embedded, and stored in the vector database. (Endpoint: POST - `genai/upload`)
+- For using/testing the upload functionality, you can find some recipe PDFs to test the upload under `recipe_pdfs` folder. If you want, you can also modify the content of the script to generate your own recipe PDFs which can be found under `recipe_pdfs/scripts` folder. You can run the script from the root folder like:
+   - ```bash
+      python recipe_pdfs/scripts/basic_recipes.py
+      ```
 
 ##### Vector Database - Qdrant
 

genai/requirements.txt

@@ -31,4 +31,7 @@ python-multipart==0.0.20
 # Testing
 pytest==8.4.1
 fpdf==1.7.2
-pypdf==5.6.0
+pypdf==5.6.0
+
+# Authentication
+python-jose[cryptography]==3.3.0

genai/routes/routes.py

@@ -1,4 +1,11 @@
-from fastapi import APIRouter, UploadFile, File, HTTPException, Request
+from fastapi import (
+    APIRouter,
+    UploadFile,
+    File,
+    HTTPException,
+    Request,
+    Depends
+)
 from fastapi.responses import JSONResponse
 import os
 
@@ -15,6 +22,7 @@
     prepare_prompt,
     process_raw_messages
 )
+from service.auth_service import get_current_user, UserInfo
 from metrics import (
     file_upload_request_counter,
     file_upload_successfully_counter,
@@ -68,12 +76,16 @@
 
 
 @router.post("/upload")
-async def upload_file(file: UploadFile = File(...)):
+async def upload_file(
+    file: UploadFile = File(...),
+    current_user: UserInfo = Depends(get_current_user)
+):
     file_upload_request_counter.inc()
     start_time = perf_counter()
     logger.info(
-        "Upload endpoint is called in genai for the file %s",
-        file.filename
+        "Upload endpoint is called in genai for the file %s by user %s",
+        file.filename,
+        current_user.username
     )
 
     if not file.filename.endswith(".pdf"):
@@ -89,7 +101,7 @@ async def upload_file(file: UploadFile = File(...)):
         with open(file_path, "wb") as buffer:
             buffer.write(await file.read())
 
-        collection_name = "recipes"
+        collection_name = f"recipes_{current_user.user_id}"
         if (
             qdrant.client.collection_exists(collection_name)
             and qdrant.collection_contains_file(
@@ -98,7 +110,10 @@ async def upload_file(file: UploadFile = File(...)):
                 filename
             )
         ):
-            logger.info("File already exists in qdrant")
+            logger.info(
+                "File already exists in qdrant for user %s",
+                current_user.username
+                )
             return {"message": f"File '{filename}' already uploaded."}
 
         vector_store = qdrant.create_and_get_vector_storage(collection_name)
@@ -130,16 +145,21 @@ async def generate(request: Request):
     logger.info("Generate endpoint is called in genai")
 
     body = await request.json()
-    if "query" not in body or "messages" not in body:
-        logger.error("Missing 'query' or 'messages' in the request body")
+    if "query" not in body or "messages" not in body or "user_id" not in body:
+        logger.error(
+            "Missing 'query', 'messages', or 'user_id' in the request body"
+            )
         raise HTTPException(
             status_code=400,
-            detail="Missing 'query' or 'messages'"
+            detail="Missing 'query', 'messages', or 'user_id'"
         )
 
     query = body["query"]
     messages_raw = body["messages"]
-    collection_name = "recipes"
+    user_id = body["user_id"]
+    collection_name = f"recipes_{user_id}"
+
+    logger.info("Generate endpoint called for user_id: %s", user_id)
 
     try:
         retrieved_docs = ""
@@ -148,8 +168,9 @@ async def generate(request: Request):
                 collection_name
             )
             logger.info(
-                "Vector store is created for the collection %s",
-                collection_name
+                "Vector store is created for the collection %s for user_id %s",
+                collection_name,
+                user_id
             )
             retrieved_docs = retrieve_similar_docs(vector_store, query)
             logger.info("Similar docs retrieved from the vector store")

genai/service/auth_service.py

@@ -0,0 +1,68 @@
+import requests
+from fastapi import HTTPException, Header
+from typing import Optional
+from logger import logger
+
+
+class UserInfo:
+    def __init__(self, user_id: int, username: str):
+        self.user_id = user_id
+        self.username = username
+
+
+async def get_current_user(
+    authorization: Optional[str] = Header(None)
+) -> UserInfo:
+    """
+    Extract user information from the Authorization header.
+
+    This function validates the OAuth token by calling the user service
+    and returns the user information including user_id.
+    """
+    if not authorization:
+        logger.error("Authorization header is missing")
+        raise HTTPException(
+            status_code=401,
+            detail="Authorization header required"
+            )
+
+    if not authorization.startswith("Bearer "):
+        logger.error("Invalid authorization header format")
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid authorization header format"
+            )
+
+    token = authorization.split(" ")[1]
+
+    # use dev and prod URLs for user service, where prod is the fallback
+    user_service_urls = [
+        "http://user-service:8081/user/info",
+        "http://localhost:8081/user/info",
+    ]
+
+    for url in user_service_urls:
+        try:
+            response = requests.get(
+                url,
+                headers={"Authorization": f"Bearer {token}"},
+                timeout=10
+            )
+            if response.status_code == 200:
+                user_data = response.json()
+                logger.info(
+                    f"User authenticated: {user_data.get('username')}"
+                    )
+                return UserInfo(
+                    user_id=user_data.get("id"),
+                    username=user_data.get("username")
+                )
+            else:
+                logger.error(f"{url} returned status {response.status_code}")
+        except requests.exceptions.RequestException as e:
+            logger.warning(f"Failed to reach {url}: {e}")
+
+    raise HTTPException(
+        status_code=500,
+        detail="Authentication service unavailable"
+        )

genai/tests/integration/generation_test.py

@@ -17,7 +17,8 @@ def test_generate_endpoint_success(_mock_exists, mock_invoke):
         "messages": [
             {"role": "USER", "content": "I have rice and eggs"},
             {"role": "ASSISTANT", "content": "How about fried rice?"}
-        ]
+        ],
+        "user_id": 123
     }
 
     response = client.post("/genai/generate", json=payload)
@@ -39,7 +40,8 @@ def test_generate_endpoint_empty_messages(_mock_exists, mock_invoke):
 
     payload = {
         "query": "Can I cook with lentils?",
-        "messages": []
+        "messages": [],
+        "user_id": 123
     }
 
     response = client.post("/genai/generate", json=payload)
@@ -59,4 +61,6 @@ def test_generate_endpoint_missing_fields():
     response = client.post("/genai/generate", json=payload)
 
     assert response.status_code == 400
-    assert response.json() == {"detail": "Missing 'query' or 'messages'"}
+    assert response.json() == {
+        "detail": "Missing 'query', 'messages', or 'user_id'"
+        }

genai/tests/integration/upload_test.py

@@ -2,6 +2,7 @@
 from unittest.mock import patch, MagicMock
 from fastapi.testclient import TestClient
 from main import app
+from service.auth_service import UserInfo, get_current_user
 
 client = TestClient(app)
 
@@ -14,6 +15,10 @@ def test_upload_file_success(
         _mock_vector_store,
         _mock_exists
 ):
+    # Mock the authentication by overriding the dependency
+    mock_user = UserInfo(user_id=123, username="test_user")
+    app.dependency_overrides[get_current_user] = lambda: mock_user
+
     mock_pipeline = MagicMock()
     mock_pipeline_class.return_value = mock_pipeline
 
@@ -32,8 +37,15 @@ def test_upload_file_success(
     mock_pipeline_class.assert_called_once()
     mock_pipeline.ingest.assert_called_once()
 
+    # Clean up the dependency override
+    app.dependency_overrides.clear()
+
 
 def test_upload_file_invalid_type():
+    # Mock the authentication by overriding the dependency
+    mock_user = UserInfo(user_id=123, username="test_user")
+    app.dependency_overrides[get_current_user] = lambda: mock_user
+
     file = io.BytesIO(b"just some text")
     file.name = "notes.txt"
 
@@ -46,10 +58,17 @@ def test_upload_file_invalid_type():
     assert (response.json()["detail"] ==
             "Invalid file type. Only PDF files are allowed.")
 
+    # Clean up the dependency override
+    app.dependency_overrides.clear()
+
 
 @patch("routes.routes.qdrant.client.collection_exists", return_value=True)
 @patch("routes.routes.qdrant.collection_contains_file", return_value=True)
 def test_upload_file_already_exists(_mock_contains, _mock_exists):
+    # Mock the authentication
+    mock_user = UserInfo(user_id=123, username="test_user")
+    app.dependency_overrides[get_current_user] = lambda: mock_user
+
     file = io.BytesIO(b"%PDF-1.4")
     file.name = "existing.pdf"
 
@@ -60,3 +79,6 @@ def test_upload_file_already_exists(_mock_contains, _mock_exists):
 
     assert response.status_code == 200
     assert "already uploaded" in response.json()["message"]
+
+    # Clean up the dependency override
+    app.dependency_overrides.clear()