Spaces:

areen09
/

test_app

Paused

Upload 2 files

8b0215d verified over 1 year ago

5.25 kB

import warnings
warnings.filterwarnings("ignore", category=FutureWarning)

import os
import logging
from contextlib import asynccontextmanager
from typing import List, Dict
from dotenv import load_dotenv
from fastapi import FastAPI, Request, Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel, HttpUrl

load_dotenv()
# --- Configuration ---
# Set up basic logging to see output in Hugging Face logs
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load the token directly from environment variables (set in HF Secrets)
# This is the correct way for server deployments.
EXPECTED_TOKEN = os.getenv("BEARER_TOKEN")
if not EXPECTED_TOKEN:
    logger.warning("BEARER_TOKEN environment variable not set. Using a default for local testing.")
    EXPECTED_TOKEN = "default-secret-for-testing"

http_bearer_scheme = HTTPBearer()

# This dictionary will act as a cache for our heavy models.
# It will be populated during the application's startup.
model_cache: Dict[str, any] = {}

# --- Lifespan Management (The Core Fix) ---
# This function runs ONCE when the application starts up.
# It loads the models before the API starts accepting requests.
@asynccontextmanager
async def lifespan(app: FastAPI):
    # --- Startup ---
    logger.info("--- Application Startup: Loading models... ---")
    
    # IMPORTANT: Import your agent and model-loading modules INSIDE the lifespan
    # to prevent them from being loaded at the top level.
    from agent import agent_app, set_retriever_for_tool
    
    # Store the initialized agent app in our cache so endpoints can access it.
    model_cache["agent_app"] = agent_app
    model_cache["set_retriever_for_tool"] = set_retriever_for_tool
    
    logger.info("--- Models Loaded. Application is ready to serve requests. ---")
    
    yield  # The application is now running and serving requests
    
    # --- Shutdown ---
    logger.info("--- Application Shutdown: Clearing model cache. ---")
    model_cache.clear()


# --- FastAPI App Initialization ---
# We pass our lifespan manager to the FastAPI app.
app = FastAPI(
    title="Intelligent Query-Retrieval System",
    description="An API for processing documents and answering questions using a LangGraph ReAct agent.",
    version="1.0.0",
    lifespan=lifespan
)


# --- Authentication Dependency ---
async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(http_bearer_scheme)):
    if not credentials or credentials.scheme != "Bearer":
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authentication scheme.")
    if credentials.credentials != EXPECTED_TOKEN:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token.")
    return credentials.credentials


# --- Pydantic Models ---
class QueryRequest(BaseModel):
    documents: HttpUrl
    questions: List[str]

class QueryResponse(BaseModel):
    answers: List[str]


# --- API Endpoints ---
@app.api_route("/", methods=["GET", "HEAD"])
def read_root():
    return {"message": "API is running. Use /hackrx/run for document processing."}

@app.post("/hackrx/run", response_model=QueryResponse, tags=["Query"])
async def run_submission(request: QueryRequest, token: str = Depends(verify_token)):
    try:
        # Retrieve the pre-loaded agent from our cache
        agent_app = model_cache.get("agent_app")
        set_retriever_for_tool = model_cache.get("set_retriever_for_tool")

        if not agent_app or not set_retriever_for_tool:
            raise HTTPException(status_code=503, detail="Models are not available. The application might be starting up.")

        logger.info(f"Processing document from URL: {request.documents}")
        # Process the document ONCE per request
        from utils import process_documents
        vector_store = process_documents(str(request.documents))
        retriever = vector_store.as_retriever(search_kwargs={"k": 5})
        
        # Configure the agent's tool with the newly created retriever for this specific request
        set_retriever_for_tool(retriever)
        
        responses = []
        logger.info(f"Answering {len(request.questions)} questions...")
        for i, question in enumerate(request.questions):
            logger.info(f"Invoking agent for question {i+1}/{len(request.questions)}: '{question}'")
            
            final_state = agent_app.invoke({
                "question": question,
                "messages": [("user", question)]
            })
            
            answer_str = final_state.get("answer", "Error: Could not generate an answer.")
            responses.append(answer_str)
            logger.info(f"Generated answer for question {i+1}")

        logger.info("All questions answered successfully.")
        return QueryResponse(answers=responses)
        
    except Exception as e:
        logger.error(f"An unhandled error occurred in /hackrx/run: {e}", exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))

	import warnings
	warnings.filterwarnings("ignore", category=FutureWarning)

	import os
	import logging
	from contextlib import asynccontextmanager
	from typing import List, Dict
	from dotenv import load_dotenv
	from fastapi import FastAPI, Request, Depends, HTTPException, status
	from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
	from pydantic import BaseModel, HttpUrl

	load_dotenv()
	# --- Configuration ---
	# Set up basic logging to see output in Hugging Face logs
	logging.basicConfig(level=logging.INFO)
	logger = logging.getLogger(__name__)

	# Load the token directly from environment variables (set in HF Secrets)
	# This is the correct way for server deployments.
	EXPECTED_TOKEN = os.getenv("BEARER_TOKEN")
	if not EXPECTED_TOKEN:
	logger.warning("BEARER_TOKEN environment variable not set. Using a default for local testing.")
	EXPECTED_TOKEN = "default-secret-for-testing"

	http_bearer_scheme = HTTPBearer()

	# This dictionary will act as a cache for our heavy models.
	# It will be populated during the application's startup.
	model_cache: Dict[str, any] = {}

	# --- Lifespan Management (The Core Fix) ---
	# This function runs ONCE when the application starts up.
	# It loads the models before the API starts accepting requests.
	@asynccontextmanager
	async def lifespan(app: FastAPI):
	# --- Startup ---
	logger.info("--- Application Startup: Loading models... ---")

	# IMPORTANT: Import your agent and model-loading modules INSIDE the lifespan
	# to prevent them from being loaded at the top level.
	from agent import agent_app, set_retriever_for_tool

	# Store the initialized agent app in our cache so endpoints can access it.
	model_cache["agent_app"] = agent_app
	model_cache["set_retriever_for_tool"] = set_retriever_for_tool

	logger.info("--- Models Loaded. Application is ready to serve requests. ---")

	yield # The application is now running and serving requests

	# --- Shutdown ---
	logger.info("--- Application Shutdown: Clearing model cache. ---")
	model_cache.clear()


	# --- FastAPI App Initialization ---
	# We pass our lifespan manager to the FastAPI app.
	app = FastAPI(
	title="Intelligent Query-Retrieval System",
	description="An API for processing documents and answering questions using a LangGraph ReAct agent.",
	version="1.0.0",
	lifespan=lifespan
	)


	# --- Authentication Dependency ---
	async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(http_bearer_scheme)):
	if not credentials or credentials.scheme != "Bearer":
	raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authentication scheme.")
	if credentials.credentials != EXPECTED_TOKEN:
	raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token.")
	return credentials.credentials


	# --- Pydantic Models ---
	class QueryRequest(BaseModel):
	documents: HttpUrl
	questions: List[str]

	class QueryResponse(BaseModel):
	answers: List[str]


	# --- API Endpoints ---
	@app.api_route("/", methods=["GET", "HEAD"])
	def read_root():
	return {"message": "API is running. Use /hackrx/run for document processing."}

	@app.post("/hackrx/run", response_model=QueryResponse, tags=["Query"])
	async def run_submission(request: QueryRequest, token: str = Depends(verify_token)):
	try:
	# Retrieve the pre-loaded agent from our cache
	agent_app = model_cache.get("agent_app")
	set_retriever_for_tool = model_cache.get("set_retriever_for_tool")

	if not agent_app or not set_retriever_for_tool:
	raise HTTPException(status_code=503, detail="Models are not available. The application might be starting up.")

	logger.info(f"Processing document from URL: {request.documents}")
	# Process the document ONCE per request
	from utils import process_documents
	vector_store = process_documents(str(request.documents))
	retriever = vector_store.as_retriever(search_kwargs={"k": 5})

	# Configure the agent's tool with the newly created retriever for this specific request
	set_retriever_for_tool(retriever)

	responses = []
	logger.info(f"Answering {len(request.questions)} questions...")
	for i, question in enumerate(request.questions):
	logger.info(f"Invoking agent for question {i+1}/{len(request.questions)}: '{question}'")

	final_state = agent_app.invoke({
	"question": question,
	"messages": [("user", question)]
	})

	answer_str = final_state.get("answer", "Error: Could not generate an answer.")
	responses.append(answer_str)
	logger.info(f"Generated answer for question {i+1}")

	logger.info("All questions answered successfully.")
	return QueryResponse(answers=responses)

	except Exception as e:
	logger.error(f"An unhandled error occurred in /hackrx/run: {e}", exc_info=True)
	raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))