Authenticate users and use userids to embed documents

Author: edkaya

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,4 @@
-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 +15,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 +69,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 +94,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 +103,7 @@ 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 +135,19 @@ 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 +156,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,54 @@
+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]
+    
+    user_service_urls = [
+    "http://localhost:8081/user/info",   # dev
+    "http://user-service:8081/user/info" # prod      
+    ]
+
+    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,4 @@ 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 genai.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
 
@@ -31,9 +36,16 @@ 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"
 
@@ -45,11 +57,18 @@ def test_upload_file_invalid_type():
     assert response.status_code == 400
     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()