x_api_key new end point

app.py

@@ -9,6 +9,7 @@ from utils.auth import token_required
 from dotenv import load_dotenv
 import os
 import torch
+from utils.auth_x import x_api_key_auth
 
 # Load environment variables from .env file
 load_dotenv()
@@ -78,7 +79,11 @@ class SearchDocumentsRequest(BaseModel):
 
 class GenerateRAGRequest(BaseModel):
     search_query: str
-    
+
+class XApiKeyRequest(BaseModel):
+    organization_id: str
+    user_id: str
+    search_query: str  
 
 
 @app.get("/")
@@ -160,6 +165,84 @@ async def generate_rag_response_api(
         logging.error(f"Unexpected error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/api/search-documents/v1")
+async def search_documents_x_api_key(
+    body: XApiKeyRequest,
+    authorized: bool = Depends(x_api_key_auth)
+):
+    if not authorized:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
+    organization_id = body.organization_id
+    user_id = body.user_id
+    logging.info(f'search query {body.search_query}')
+    logging.info(f"organization_id: {organization_id}, user_id: {user_id}")
+    logging.info("Received request to search documents with x-api-key auth")
+    try:
+        logging.info("Starting document search")
+
+        # Encode the query using the custom embedding function
+        query_embedding = embed_text(body.search_query)
+        collection_name = "embed"  # Use the collection name where the embeddings are stored   
+
+        # Perform search using the precomputed embeddings
+        hits, error = searcher.search_documents(collection_name, query_embedding, user_id, limit=3)
+        
+        if error:
+            logging.error(f"Search documents error: {error}")
+            raise HTTPException(status_code=500, detail=error)
+
+        logging.info(f"Document search completed with {len(hits)} hits")
+        return hits
+    except Exception as e:
+        logging.error(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/generate-rag-response/v1")
+async def generate_rag_response_x_api_key(
+    body: XApiKeyRequest,
+    authorized: bool = Depends(x_api_key_auth)
+):
+    # Assuming x_api_key_auth validates the key
+    if not authorized:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
+    organization_id = body.organization_id
+    user_id = body.user_id
+
+    logging.info(f'search query {body.search_query}')
+    logging.info(f"organization_id: {organization_id}, user_id: {user_id}")
+    logging.info("Received request to generate RAG response with x-api-key auth")
+    try:
+        logging.info("Starting document search")
+
+        # Encode the query using the custom embedding function
+        query_embedding = embed_text(body.search_query)
+        collection_name = "embed"  # Use the collection name where the embeddings are stored   
+
+        # Perform search using the precomputed embeddings
+        hits, error = searcher.search_documents(collection_name, query_embedding, user_id)
+        
+        if error:
+            logging.error(f"Search documents error: {error}")
+            raise HTTPException(status_code=500, detail=error)
+
+        logging.info("Generating RAG response")
+
+        # Generate the RAG response using the retrieved documents
+        response, error = generate_rag_response(hits, body.search_query)
+        
+        if error:
+            logging.error(f"Generate RAG response error: {error}")
+            raise HTTPException(status_code=500, detail=error)
+
+        return {"response": response}
+    except Exception as e:
+        logging.error(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+
 if __name__ == '__main__':
     import uvicorn
     uvicorn.run(app, host='0.0.0.0', port=8000)

utils/auth_x.py

@@ -0,0 +1,7 @@
+from fastapi import Header, HTTPException
+
+def x_api_key_auth(x_api_key: str = Header(...)):
+    expected_api_key = "your_expected_api_key"  # Replace with your actual API key
+    if x_api_key != expected_api_key:
+        raise HTTPException(status_code=401, detail="Invalid API Key")
+    return True