feat: :lock: Implement API Key authentication and CORS configuration for enhanced security. Don`t returns thread ID

app.py

@@ -4,6 +4,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from functools import partial
 from fastapi.responses import JSONResponse
+from fastapi import Security, Depends
+from fastapi.security.api_key import APIKeyHeader, APIKey
+from fastapi.middleware.cors import CORSMiddleware
 
 from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.checkpoint.memory import MemorySaver
@@ -13,6 +16,19 @@ import os
 from dotenv import load_dotenv
 load_dotenv()
 
+# Configuración de API Key
+API_KEY_NAME = "X-API-Key"
+API_KEY = os.getenv("API_KEY")
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+
+async def get_api_key(api_key_header: str = Security(api_key_header)):
+    if api_key_header == API_KEY:
+        return api_key_header
+    raise HTTPException(
+        status_code=403,
+        detail="Could not validate API KEY"
+    )
+
 # Initialize the model and tokenizer
 print("Loading model and tokenizer...")
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -120,7 +136,45 @@ class SummaryRequest(BaseModel):
     max_length: int = 200
 
 # Create the FastAPI application
-app = FastAPI(title="LangChain FastAPI", description="API to generate text using LangChain and LangGraph - Máximo Fernández Núñez IriusRisk test challenge")
+app = FastAPI(
+    title="LangChain FastAPI",
+    description="API to generate text using LangChain and LangGraph - Máximo Fernández Núñez IriusRisk test challenge",
+    version="1.0.0",
+    openapi_tags=[
+        {
+            "name": "Authentication",
+            "description": "Endpoints require API Key authentication via X-API-Key header"
+        }
+    ]
+)
+
+# Configure the security scheme in the OpenAPI documentation
+app.openapi_tags = [
+    {"name": "Authentication", "description": "Protected endpoints that require API Key"}
+]
+
+# Import and configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Configure the security scheme
+app.openapi_components = {
+    "securitySchemes": {
+        "api_key": {
+            "type": "apiKey",
+            "name": API_KEY_NAME,
+            "in": "header",
+            "description": "Enter your API key"
+        }
+    }
+}
+
+app.openapi_security = [{"api_key": []}]
 
 # Add general exception handler
 @app.exception_handler(Exception)
@@ -138,7 +192,10 @@ async def api_home():
 
 # Generate endpoint
 @app.post("/generate")
-async def generate(request: QueryRequest):
+async def generate(
+    request: QueryRequest,
+    api_key: APIKey = Depends(get_api_key)
+):
     """
     Endpoint to generate text using the language model
     
@@ -147,9 +204,10 @@ async def generate(request: QueryRequest):
             query: str
             thread_id: str = "default"
             system_prompt: str = DEFAULT_SYSTEM_PROMPT
+        api_key: APIKey - API key for authentication
 
     Returns:
-        dict: A dictionary containing the generated text and the thread ID
+        dict: A dictionary containing the generated text
     """
     try:
         # Configure the thread ID
@@ -175,21 +233,22 @@ async def generate(request: QueryRequest):
         response = output["messages"][-1].content
         
         return {
-            "generated_text": response,
-            "thread_id": request.thread_id
+            "generated_text": response
         }
     except Exception as e:
         return JSONResponse(
             status_code=500,
             content={
                 "error": f"Error generando texto: {str(e)}", 
-                "type": type(e).__name__,
-                "thread_id": request.thread_id
+                "type": type(e).__name__
             }
         )
 
 @app.post("/summarize")
-async def summarize(request: SummaryRequest):
+async def summarize(
+    request: SummaryRequest,
+    api_key: APIKey = Depends(get_api_key)
+):
     """
     Endpoint to generate a summary using the language model
     
@@ -198,9 +257,10 @@ async def summarize(request: SummaryRequest):
             text: str - The text to summarize
             thread_id: str = "default"
             max_length: int = 200 - Maximum summary length
+        api_key: APIKey - API key for authentication
 
     Returns:
-        dict: A dictionary containing the summary and the thread ID
+        dict: A dictionary containing the summary
     """
     try:
         # Configure the thread ID
@@ -228,16 +288,14 @@ async def summarize(request: SummaryRequest):
         response = output["messages"][-1].content
         
         return {
-            "summary": response,
-            "thread_id": request.thread_id
+            "summary": response
         }
     except Exception as e:
         return JSONResponse(
             status_code=500,
             content={
                 "error": f"Error generando resumen: {str(e)}", 
-                "type": type(e).__name__,
-                "thread_id": request.thread_id
+                "type": type(e).__name__
             }
         )