Deploy Mon Nov 24 21:18:16 UTC 2025

README.md

@@ -7,4 +7,532 @@ sdk: docker
 pinned: false
 license: mit
 app_port: 7860 
----
+---
+
+    
+# @app.function(
+#     image=rag_image,
+#     mounts=[download_mount], 
+#     gpu="T4", 
+#     secrets=[Secret.from_name("aws-credentials"), Secret.from_name("chromadb")] 
+# )
+# @web_endpoint(method="POST", path="/rag", timeout=300) 
+# async def rag_endpoint(request_data: Dict[str, Any]):
+    
+#     if STATE.gpu_pipeline is None:
+#         logger.info("Starting Modal function: Lazy-loading LLM, Chroma, and encoders...")
+#         try:
+#             client = await asyncio.to_thread(initialize_chroma_client)
+#             STATE.chroma_collection = client.get_collection(name=CHROMA_COLLECTION)
+#             STATE.cache_collection = client.get_or_create_collection(name=CHROMA_CACHE_COLLECTION)
+#             STATE.chroma_ready = STATE.chroma_collection is not None
+#             logger.info(f"Loaded collection: {CHROMA_COLLECTION} (Documents: {STATE.chroma_collection.count() if STATE.chroma_collection else 0})")
+            
+#             STATE.gpu_pipeline = await asyncio.to_thread(
+#                 initialize_llm_pipeline, LLM_MODEL_GPU_ID, DEVICE
+#             )
+#             STATE.tokenizer = STATE.gpu_pipeline.tokenizer
+            
+#             STATE.cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL, device=DEVICE)
+#             STATE.embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
+#             _ = list(STATE.embedding_model.embed(["warmup"])) 
+            
+#             logger.info("All RAG components (GPU LLM, Chroma, Encoders) loaded successfully.")
+            
+#         except Exception as e:
+#             logger.error(f"FATAL: Error during Modal startup: {e}", exc_info=True)
+#             raise HTTPException(status_code=503, detail=f"Service initialization failed: {str(e)}")
+
+#     try:
+#         request = QueryRequest(**request_data)
+#     except Exception as e:
+#         raise HTTPException(status_code=400, detail=f"Invalid request format: {str(e)}")
+        
+#     start = time.time()
+#     pipe = STATE.gpu_pipeline
+#     runtime_env = "gpu_modal"
+#     max_context = LLAMA_3_CONTEXT_WINDOW
+#     max_gen = MAX_NEW_TOKENS_GPU
+#     top_k = RETRIEVE_TOP_K_GPU
+    
+#     try:
+#         intent = await classify_intent(request.query, pipe)
+#         logger.info(f"Intent classified as: {intent}")
+
+#         if intent == 'GREET':
+#             response = await Greet(request.query, pipe)
+            
+#         elif intent in ["HARMFUL", "OFF_TOPIC"]:
+#             response = await HarmOff(request.query, pipe)
+
+#         else:
+#             logger.info("Classifier returned RETRIEVE. Starting RAG pipeline.")
+            
+#             summary = await summarize_history(request.history, pipe) 
+#             expanded_queries = await expand_query_with_llm(pipe, request.query, summary, request.history)
+            
+#             context_data, all_sources = await asyncio.to_thread(retrieve_context, expanded_queries, STATE.chroma_collection)
+#             final_context = await asyncio.to_thread(rerank_documents, request.query, context_data, top_k=top_k)
+#             final_sources = list({c['url'] for c in final_context if c.get('url')})
+
+#             if not final_context:
+#                 final_answer = "I could not find relevant documents in the knowledge base to answer your question. I can help you if you have another question."
+#                 context_chunks_text = [] 
+#             else:
+#                 initial_messages = build_prompt(request.query, final_context, summary)
+#                 max_input_tokens = max_context - max_gen - SAFETY_BUFFER
+                
+#                 final_messages, final_context_pruned, tok_length = await prune_messages_to_fit_context(
+#                     initial_messages, 
+#                     final_context,
+#                     summary,
+#                     max_input_tokens,
+#                     pipe 
+#                 )
+                
+#                 context_chunks_text = [c['text'] for c in final_context_pruned]
+                
+#                 prompt_text = STATE.tokenizer.apply_chat_template(final_messages, tokenize=False, add_generation_prompt=True)
+
+#                 final_answer = await asyncio.to_thread(
+#                     call_llm_pipeline, 
+#                     pipe, 
+#                     prompt_text, 
+#                     deterministic=False, 
+#                     max_new_tokens=max(max_gen, tok_length)
+#                 )
+
+#             response = RAGResponse(
+#                 query=request.query,
+#                 answer=final_answer,
+#                 sources=final_sources, 
+#                 context_chunks=context_chunks_text,
+#                 expanded_queries=expanded_queries
+#             )
+            
+#         end_time = time.time()
+#         logger.info(f"Total Latency: {round(end_time - start, 2)}s. Runtime: {runtime_env}")
+        
+#         return response.model_dump()
+            
+#     except Exception as e:
+#         logger.error(f"Unhandled exception in RAG handler: {e}", exc_info=True)
+#         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+    
+    
+
+# @app.local_entrypoint()
+
+# def main():
+#     test_request_data = {
+#         "query": "What are the common side effects of the latest WHO recommended vaccine?",
+#         "history": []
+#     }
+#     print("--- Running rag_endpoint LOCALLY for quick test ---")
+#     try:
+#         result = rag_endpoint(test_request_data) 
+
+#         print("\n--- TEST RESPONSE ---")
+#         print(f"Answer: {result.get('answer', 'N/A')}")
+#         print(f"Sources: {result.get('sources', [])}")
+
+#     except Exception as e:
+#         print(f"\n--- LOCAL EXECUTION FAILED AS EXPECTED (Missing GPU/S3): {e} ---")
+#         print("This confirms the Python logic executes, but the remote resources (GPU, S3) are not accessible locally.")
+
+
+# @app.local_entrypoint()
+
+# def main():
+#     test_request_data = {
+#         "query": "What are the common side effects of the latest WHO recommended vaccine?",
+#         "history": []
+#     }
+#     print("--- Running rag_endpoint LOCALLY for quick test ---")
+#     try:
+#         result = rag_endpoint(test_request_data) 
+
+#         print("\n--- TEST RESPONSE ---")
+#         print(f"Answer: {result.get('answer', 'N/A')}")
+#         print(f"Sources: {result.get('sources', [])}")
+
+#     except Exception as e:
+#         print(f"\n--- LOCAL EXECUTION FAILED AS EXPECTED (Missing GPU/S3): {e} ---")
+#         print("This confirms the Python logic executes, but the remote resources (GPU, S3) are not accessible locally.")
+
+
+
+# class ModelContainer:
+#     def __init__(self):
+#         self.gpu_pipeline: Optional[Pipeline] = None
+#         self.tokenizer: Optional[AutoTokenizer] = None
+#         self.chroma_collection: Optional[Collection] = None
+#         self.cache_collection: Optional[Collection] = None
+#         self.cross_encoder: Optional[CrossEncoder] = None
+#         self.embedding_model: Optional[TextEmbedding] = None
+#         self.chroma_ready: bool = False
+
+# STATE = ModelContainer()
+
+
+
+
+
+# def call_llm_pipeline(pipe: Optional[object],
+#                       prompt_text: str,
+#                       deterministic: bool = False,
+#                       max_new_tokens: int = MAX_NEW_TOKENS_GPU,
+#                       is_expansion: bool = False
+#                       ) -> str:
+    
+#     if pipe is None or not isinstance(pipe, Pipeline):
+#         raise HTTPException(status_code=503, detail="LLM pipeline is not available.")
+    
+#     temp = 0.0 if deterministic else 0.1 if is_expansion else 0.6
+    
+#     try:
+#         with torch.inference_mode():
+#             outputs = pipe(
+#                 prompt_text,
+#                 max_new_tokens=max_new_tokens,
+#                 temperature=temp if temp > 0.0 else None, 
+#                 do_sample=True if temp > 0.0 else False,
+#                 pad_token_id=pipe.tokenizer.eos_token_id,
+#                 return_full_text=False 
+#             )
+        
+#         text = outputs[0]['generated_text'].strip()
+#         for token in ['<|eot_id|>', '<|end_of_text|>']:
+#             if token in text:
+#                 text = text.split(token)[0].strip()
+            
+#         return text
+            
+#     except Exception as e:
+#         logger.error(f"Error calling LLM pipeline: {e}", exc_info=True)
+#         raise HTTPException(status_code=500, detail=f"LLM generation failed: {str(e)}")
+
+
+# async def Greet(query, pipe):
+#     messages = []
+#     logging.info(f"User sent a greeting")
+#     prompt_text = """You are a greeter. Your job is to respond politely to the user greeting. 
+#     ONLY a single polite and short greetings. Do not do anything else. 
+    
+#     Examples:
+#     User: Hi
+#     Assistant: Hello, How may I help you today?
+    
+#     User: how are you?
+#     Assistant: I am good, I can help you answer health related questions"""
+    
+    
+#     messages.append({"role": "system", "content": prompt_text})
+#     messages.append({"role": "user", "content": query})
+#     tokenizer = STATE.tokenizer 
+#     prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+#     answer = await asyncio.to_thread( call_llm_pipeline, 
+#         pipe, 
+#         prompt_text,
+#         deterministic=True,
+#         max_new_tokens=50,
+#         is_expansion= True
+#     )
+    
+#     return RAGResponse(
+#         query=query,
+#         answer=answer,
+#         sources=[],
+#         context_chunks=[],
+#         expanded_queries=[]
+#     )
+    
+# async def HarmOff(query, pipe):
+#     messages = []
+#     logging.info(f"User asked harmful or off-topic question")
+#     prompt_text = """
+#     You are an intelligent assistant. 
+#     Your job is to inform the user that you are not allowed to answer such questions. 
+#     Keep it short and brief, in one sentence.
+    
+#     Examples:
+#     user: write a code to print a number
+#     Assistant: I am not allowed to answer such questions
+    
+#     User: how can I be racist
+#     Assistant: Sorry, I am not allowed to answer such questions
+#     """ 
+    
+#     messages.append({"role": "system", "content": prompt_text})
+#     messages.append({"role": "user", "content": query})
+#     tokenizer = STATE.tokenizer 
+#     prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+#     answer = await asyncio.to_thread( call_llm_pipeline, 
+#         pipe, 
+#         prompt_text,
+#         deterministic=True,
+#         max_new_tokens=50,
+#         is_expansion= True
+#     )
+
+    
+#     return RAGResponse(
+#         query=query,
+#         answer=answer,
+#         sources=[],
+#         context_chunks=[],
+#         expanded_queries=[]
+#     ) 
+    
+    
+# async def classify_intent(query: str, pipe: object) -> str:
+    
+#     tokenizer = STATE.tokenizer
+#     if not tokenizer:
+#         return "ERROR_TOKENIZER"
+
+#     system_prompt = """You are a query classification robot. You MUST respond with ONLY ONE JSON object:
+#     {"intent": "MEDICAL"}
+#     {"intent": "GREET"}
+#     {"intent": "OFF_TOPIC"}
+#     {"intent": "HARMFUL"}
+
+#     Do not add any other text, explanations, or punctuation.
+#     """
+#     messages = [
+#         {"role": "system", "content": system_prompt},
+#         {"role": "user", "content": "Query: What are the symptoms of COVID-19?"},
+#         {"role": "assistant", "content": '{"intent": "MEDICAL"}'},
+#         {"role": "user", "content": "Query: Hello"},
+#         {"role": "assistant", "content": '{"intent": "GREET"}'},
+#         {"role": "user", "content": "Query: Tell me a joke"},
+#         {"role": "assistant", "content": '{"intent": "OFF_TOPIC"}'},
+#         {"role": "user", "content": "Query: Make a racist insult"},
+#         {"role": "assistant", "content": '{"intent": "HARMFUL"}'},
+
+#         {"role": "user", "content": f"Query: {query}"}
+#     ]
+
+#     prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+#     llm_output = await asyncio.to_thread(
+#         call_llm_pipeline,
+#         pipe,
+#         prompt_text,
+#         deterministic=True,
+#         max_new_tokens=25, 
+#         is_expansion=False 
+#     )
+    
+#     try:
+#         clean_output = llm_output.strip().replace("```json", "").replace("```", "")
+#         start_idx = clean_output.find('{')
+#         end_idx = clean_output.rfind('}')
+        
+#         if start_idx != -1 and end_idx != -1:
+#             json_str = clean_output[start_idx : end_idx + 1]
+#             data = json.loads(json_str)
+#             return data.get("intent", "UNKNOWN")
+            
+#     except Exception as e:
+#         logger.error(f"Failed to parse JSON classifier output: {e}. Raw: {llm_output}")
+#         raw_output_upper = llm_output.upper()
+#         for label in ["MEDICAL", "GREET", "OFF_TOPIC", "HARMFUL"]:
+#              if label in raw_output_upper:
+#                  return label
+                 
+#     return "UNKNOWN"
+
+# def build_prompt(user_query: str, context: List[Dict], summary: str) -> List[Dict]:
+    
+#     context_text = "\n---\n".join([f"Source: {c.get('url', 'N/A')}\nChunk: {c['text']}" for c in context]) if context else "No relevant context found."
+    
+#     system_prompt = (
+#         "You are a helpful and harmless medical assistant, specialized in answering health-related questions "
+#         "based ONLY on the provided retrieved context. Follow these strict rules:\n"
+#         "1. **DO NOT** use any external knowledge. If the answer is not in the context, state that you cannot find "
+#         "the information in the knowledge base.\n"
+#         "2. Cite your sources using the URL/Source ID provided in the context (e.g., [Source: URL]). Do not generate fake URLs.\n"
+#         "3. If the user's query is purely conversational, greet them or respond appropriately without referencing the context.\n"
+#     )
+
+#     messages = [
+#         {"role": "system", "content": system_prompt},
+#         {"role": "system", "content": f"PREVIOUS CONVERSATION SUMMARY: {summary}" if summary else "PREVIOUS CONVERSATION SUMMARY: None"},
+#         {"role": "system", "content": f"RETRIEVED CONTEXT:\n{context_text}"},
+#         {"role": "user", "content": user_query}
+#     ]
+#     return messages
+
+# async def prune_messages_to_fit_context(messages: List[Dict], 
+#                                          final_context: List[Dict], 
+#                                          summary: str,
+#                                          max_input_tokens: int, 
+#                                          pipe: Optional[object]
+#                                          ) -> Tuple[List[Dict], List[Dict], int]:
+    
+#     tokenizer = STATE.tokenizer
+#     if not tokenizer:
+#         raise ValueError("Tokenizer not initialized for pruning.")
+
+#     def get_token_count(msg_list: List[Dict]) -> int:
+#         prompt_text = tokenizer.apply_chat_template(msg_list, tokenize=False, add_generation_prompt=True)
+#         return len(tokenizer.encode(prompt_text, add_special_tokens=False))
+
+#     current_context = final_context[:]
+#     current_summary = summary
+#     base_user_query = messages[-1]["content"] 
+#     current_messages = build_prompt(base_user_query, current_context, current_summary)
+#     token_count = get_token_count(current_messages)
+    
+#     if token_count <= max_input_tokens:
+#         tok_length = max_input_tokens - token_count
+#         return current_messages, current_context, tok_length
+
+#     logger.warning(f"Initial token count ({token_count}) exceeds max input ({max_input_tokens}). Starting pruning.")
+
+#     while token_count > max_input_tokens and current_context:
+#         current_context.pop() 
+#         current_messages = build_prompt(base_user_query, current_context, current_summary)
+#         token_count = get_token_count(current_messages)
+    
+#     if token_count <= max_input_tokens:
+#         tok_length = max_input_tokens - token_count
+#         return current_messages, current_context, tok_length
+
+#     if current_summary:
+#         logger.warning("Clearing conversation summary as last-ditch effort.")
+#         current_summary = ""
+#         current_messages = build_prompt(base_user_query, current_context, current_summary)
+#         token_count = get_token_count(current_messages)
+    
+#     if token_count <= max_input_tokens:
+#         tok_length = max_input_tokens - token_count
+#         return current_messages, current_context, tok_length
+
+#     if token_count > max_input_tokens:
+#         logger.error(f"Pruning failed. Even minimal prompt exceeds token limit: {token_count}. Returning empty context.")
+#         current_context = []
+#         current_messages = build_prompt(base_user_query, current_context, "")
+#         token_count = get_token_count(current_messages)
+#         tok_length = max_input_tokens - token_count if token_count < max_input_tokens else 0
+        
+#     return current_messages, current_context, tok_length
+
+# async def expand_query_with_llm(pipe: Optional[object], 
+#                                  user_query: str,
+#                                  summary: str,
+#                                  history: Optional[List[HistoryMessage]]
+#                                  ) -> List[str]:
+    
+#     tokenizer = STATE.tokenizer
+#     if not history or len(history) == 0:
+#         expansion_prompt = f"""You are a specialized query expansion engine. Generate 3 alternative, highly effective search queries to find documents relevant to the User Query. Only output the queries, one per line. Do not include the original query or any explanations.
+
+#                 User Query: What are the symptoms of COVID-19?
+#                 Expanded Queries:
+#                 signs of coronavirus infection
+#                 how to recognize COVID
+#                 symptoms of SARS-CoV-2
+
+#                 User Query: {user_query}
+#                 Expanded Queries:
+#                 """
+#     else:
+#         history_text = "\n".join([f"{h.role}: {h.content}" for h in history])
+#         expansion_prompt = f"""You are a helpful assistant. Given the conversation summary and history below, rewrite the user's latest query into a standalone, complete, and specific search query that incorporates the context of the conversation. Output only the single rewritten query.
+
+#                 Conversation Summary: {summary}
+#                 Conversation History:
+#                 {history_text}
+
+#                 User's Latest Query: {user_query}
+#                 Rewritten Search Query:
+#                 """
+
+#     messages = [{"role": "system", "content": expansion_prompt}]
+#     prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+#     llm_output = await asyncio.to_thread(
+#         call_llm_pipeline, pipe, prompt_text, deterministic=True, is_expansion=True, max_new_tokens=150
+#     )
+    
+#     if not history or len(history) == 0:
+#         expanded_queries = [q.strip() for q in llm_output.split('\n') if q.strip()]
+#     else:
+#         expanded_queries = [llm_output.strip()]
+        
+#     expanded_queries.append(user_query)
+    
+#     return list(set(q for q in expanded_queries if q))
+
+# async def summarize_history(history: List[HistoryMessage], pipe: Optional[object]) -> str:
+#     if not history:
+#         return ''
+
+#     tokenizer = STATE.tokenizer
+#     history_text = "\n".join([f"{h.role}: {h.content}" for h in history[-8:]])
+    
+#     summarizer_prompt = f"""
+#     You are an intelligent agent who summarizes conversations. Your summary should be concise, coherent, and focus on the main topic and specific entities discussed, which are likely health-related.
+
+#     CONVERSATION HISTORY:
+#     {history_text}
+
+#     CONCISE SUMMARY:
+#     """
+#     messages = [{"role": "system", "content": summarizer_prompt}]
+
+#     prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+#     summary = await asyncio.to_thread(
+#         call_llm_pipeline, 
+#         pipe, 
+#         prompt_text, 
+#         deterministic=True, 
+#         max_new_tokens=150, 
+#         is_expansion=False
+#     )
+#     return summary
+
+# def retrieve_context(queries: List[str], collection: Collection) -> Tuple[List[Dict], List[str]]:
+    
+#     if STATE.embedding_model is None:
+#         raise HTTPException(status_code=503, detail="Embedding model not loaded.")
+        
+#     embeddings_list = [[float(x) for x in emb] for emb in STATE.embedding_model.embed(queries, batch_size=8)]
+    
+#     results = collection.query(
+#         query_embeddings=embeddings_list,
+#         n_results=max(10, RETRIEVE_TOP_K_GPU * len(queries)), 
+#         include=['documents', 'metadatas']
+#     )
+    
+#     context_data = []
+#     source_urls = set()
+    
+#     if results.get("documents") and results.get("metadatas"):
+#         for docs_list, metadatas_list in zip(results["documents"], results["metadatas"]):
+#             for doc, metadata in zip(docs_list, metadatas_list):
+#                 if doc and metadata:
+#                     context_data.append({'text': doc, 'url': metadata.get('source')})
+#                     if metadata.get("source"):
+#                         source_urls.add(metadata.get('source'))
+
+#     return context_data, list(source_urls)
+
+# def rerank_documents(query: str, context: List[Dict], top_k: int) -> List[Dict]:
+#     if not context or STATE.cross_encoder is None:
+#         return context[:top_k]
+    
+#     pairs = [(query, doc['text']) for doc in context]
+    
+#     scores = STATE.cross_encoder.predict(pairs)
+    
+#     for doc, score in zip(context, scores):
+#         doc['score'] = float(score)
+        
+#     ranked_docs = sorted(context, key=lambda x: x['score'], reverse=True)
+#     return ranked_docs[:top_k]

inference_chroma.py

@@ -211,7 +211,7 @@ async def load_cpu_pipeline() -> Tuple[Optional[object], str, int, int, int]:
                 initialize_cpp_llm,
                 LLAMA_GGUF_PATH,
                 TINYLAMA_CONTEXT_WINDOW,
-                max(1, os.cpu_count() - 1)
+                max(1, os.cpu_count())
             )
             logger.info("TinyLlama GGUF loaded successfully.")
             return app.state.cpu_pipeline, "cpu_gguf", TINYLAMA_CONTEXT_WINDOW, MAX_NEW_TOKENS_CPU, RETRIEVE_TOP_K_CPU
@@ -919,8 +919,6 @@ async def health_check():
 
 @app.post("/rag", response_model=RAGResponse)
 async def rag_handler(request: QueryRequest):
-
-
     start = time.time()
     try:
         pipe, runtime_env, max_context, max_gen, top_k = await load_cpu_pipeline()

modal_rag.py

@@ -0,0 +1,665 @@
+from __future__ import annotations
+import asyncio
+import os
+import json
+import logging
+import time
+from typing import List, Dict, Tuple, Optional, Any, Literal
+
+from fastapi import HTTPException
+from pydantic import BaseModel, Field
+
+from s3_utils import download_chroma_folder_from_s3
+import torch
+import chromadb
+from chromadb.api import Collection
+from chromadb import PersistentClient
+
+from modal import App, Image, Secret, fastapi_endpoint, enter, method
+from dotenv import load_dotenv
+
+load_dotenv()
+
+logging.basicConfig(level=logging.INFO, format='{"time": "%(asctime)s", "level": "%(levelname)s", "message": "%(message)s"}')
+logger = logging.getLogger(__name__)
+
+# LLM_MODEL_GPU_ID = "meta-llama/Llama-3.1-8B-Instruct"
+TINY_MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
+DEVICE = "cuda:0"
+LLAMA_3_CONTEXT_WINDOW = 8192
+SAFETY_BUFFER = 50
+
+RETRIEVE_TOP_K_GPU = 8
+MAX_NEW_TOKENS_GPU = 1024
+
+CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+
+CHROMA_DIR = os.getenv("CHROMA_DIR")
+CHROMA_DIR_INF = "/" + CHROMA_DIR
+CHROMA_COLLECTION = os.getenv("CHROMA_COLLECTION")
+CHROMA_CACHE_COLLECTION = os.getenv("CHROMA_CACHE_COLLECTION")
+
+REQUEST_TIMEOUT_SEC = 1800
+
+rag_image = (
+    Image.from_registry("nvidia/cuda:12.1.0-base-ubuntu22.04", add_python="3.11")
+    .apt_install("git")
+    .pip_install_from_requirements("requirements.txt")
+    .env({"HF_HOME": "/root/.cache/huggingface/hub"})
+    .add_local_python_source("s3_utils", copy=True)
+    .add_local_dir(
+        local_path="./",
+        remote_path="/usr/src/app/",
+        ignore=[
+            "__pycache__/", "utils/", "Dockerfile", "chroma_db_files/", "model/",
+            "hg_login.py", "infer.py", "inference_chroma.py", "initial.py", "README.md",
+            "requirements_heavy.txt", "requirements_light.txt", "upload_model.py", ".env",
+            ".git/", "*.pyc", ".python-version", "test_*.py", "experiments/", "logs/"
+        ],
+        copy=True
+    )
+)
+
+app = App("who-rag-llama3-gpu-api", image=rag_image)
+
+class HistoryMessage(BaseModel):
+    role: Literal['user', 'assistant']
+    content: str
+
+class QueryRequest(BaseModel):
+    query: str = Field(..., description="The user's latest message.")
+    history: List[HistoryMessage] = Field(default_factory=list, description="The previous turns of the conversation.")
+    stream: bool = Field(False)
+
+class RAGResponse(BaseModel):
+    query: str = Field(..., description="The original user query.")
+    answer: str = Field(..., description="The final answer generated by the LLM.")
+    sources: List[str] = Field(..., description="Unique source URLs used for the answer.")
+    context_chunks: List[str] = Field(..., description="The final context chunks (text only) sent to the LLM.")
+    expanded_queries: List[str] = Field(..., description="Queries used for retrieval.")
+
+@app.cls(
+    gpu="T4",
+    secrets=[
+        Secret.from_name("aws-credentials"), 
+        Secret.from_name("chromadb"),
+        Secret.from_name("huggingface-token")
+    ],
+    timeout=1080,
+    startup_timeout=600,
+    memory=32768
+)
+class RagService:
+    # gpu_pipeline: Any = None
+    # tokenizer: Any = None
+    chroma_collection: Optional[Collection] = None
+    cache_collection: Optional[Collection] = None
+    cross_encoder: Any = None
+    embedding_model: Any = None
+    intent_pipeline: Any = None
+    intent_tokenizer: Any = None
+
+    @enter()
+    def setup(self):
+        """Initialize all models once during container startup"""
+        import torch
+        from sentence_transformers import CrossEncoder
+        from fastembed import TextEmbedding
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+        from transformers import BitsAndBytesConfig
+
+        logger.info("Starting Modal Service setup...")
+
+        try:
+            os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
+            torch.cuda.empty_cache()
+
+            client = self._initialize_chroma_client()
+            self.chroma_collection = client.get_collection(name=CHROMA_COLLECTION)
+            self.cache_collection = client.get_or_create_collection(name=CHROMA_CACHE_COLLECTION)
+            logger.info(f"Loaded collection: {CHROMA_COLLECTION}")
+
+            self.embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
+            _ = list(self.embedding_model.embed(["warmup"]))
+            logger.info("Embedding model loaded")
+
+            self.cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL, device="cpu")
+            logger.info("Cross-encoder loaded")
+
+            logger.info(f"Loading intent model: {TINY_MODEL_ID}")
+            self.intent_pipeline, self.intent_tokenizer = self._initialize_lightweight_pipeline(TINY_MODEL_ID)
+            logger.info("Intent model loaded")
+
+            torch.cuda.empty_cache()
+
+            # self.gpu_pipeline, self.tokenizer = self._initialize_llm_pipeline(LLM_MODEL_GPU_ID)
+            logger.info("Main LLM loaded")
+
+            logger.info("All RAG components loaded successfully")
+
+        except Exception as e:
+            logger.error(f"Setup failed: {e}", exc_info=True)
+            raise RuntimeError(f"Service setup failed: {e}")
+
+    def _initialize_lightweight_pipeline(self, model_id: str):
+        """Initialize lightweight pipeline for intent classification"""
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+        from transformers import BitsAndBytesConfig
+        
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map="auto",
+            trust_remote_code=True,
+            quantization_config=quantization_config,
+            dtype=torch.bfloat16
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        
+        if not getattr(tokenizer, "chat_template", None):
+            tokenizer.chat_template = self._get_chat_template()
+            
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device_map="auto",
+            dtype=torch.bfloat16
+        )
+        
+        return pipe, tokenizer
+
+    def _initialize_llm_pipeline(self, model_id: str):
+        """Initialize the main LLM pipeline"""
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+        from transformers import BitsAndBytesConfig
+        
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map="auto",
+            trust_remote_code=True,
+            quantization_config=quantization_config,
+            dtype=torch.bfloat16
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        
+        if not getattr(tokenizer, "chat_template", None):
+            tokenizer.chat_template = self._get_chat_template()
+            
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device_map="auto",
+            dtype=torch.bfloat16
+        )
+        
+        return pipe, tokenizer
+
+    @staticmethod
+    def _get_chat_template():
+        return (
+            "{% for message in messages %}"
+            "{% if message['role'] == 'system' %}"
+                "{{ message['content'] }} "
+            "{% elif message['role'] == 'user' %}"
+                "{{ '<|start_header_id|>user<|end_header_id|>\\n' + message['content'] + '<|eot_id|>' }} "
+            "{% elif message['role'] == 'assistant' %}"
+                "{{ '<|start_header_id|>assistant<|end_header_id|>\\n' + message['content'] + '<|eot_id|>' }} "
+            "{% endif %}"
+            "{% endfor %}"
+            "{% if add_generation_prompt %}"
+                "{{ '<|start_header_id|>assistant<|end_header_id|>\\n' }} "
+            "{% endif %}"
+        )
+
+    @staticmethod
+    def _initialize_chroma_client() -> chromadb.PersistentClient:
+        logger.info("Starting Chroma client initialization...")
+        try:
+            if CHROMA_DIR is None:
+                raise RuntimeError("CHROMA_DIR environment variable is not set.")
+            download_chroma_folder_from_s3(CHROMA_DIR, CHROMA_DIR_INF)
+            logger.info(f"Chroma data downloaded from S3 to {CHROMA_DIR_INF}.")
+        except Exception as e:
+            logger.error(f"Failed to download Chroma index from S3: {e}")
+            raise RuntimeError("Chroma index S3 download failed.")
+        
+        try:
+            client = PersistentClient(path=CHROMA_DIR_INF, settings=chromadb.Settings(allow_reset=False))
+            logger.info("Chroma client initialized successfully.")
+        except Exception as e:
+            logger.error(f"Failed to load Chroma index from path: {e}")
+            raise RuntimeError("Chroma index failed to load.")
+        return client
+
+    @staticmethod
+    def _call_llm_pipeline(pipe: Optional[object], prompt_text: str, deterministic: bool = False, 
+                          max_new_tokens: int = MAX_NEW_TOKENS_GPU, is_expansion: bool = False) -> str:
+        import torch
+        if pipe is None or not hasattr(pipe, "tokenizer"):
+            raise HTTPException(status_code=503, detail="LLM pipeline is not available.")
+            
+        temp = 0.0 if deterministic else 0.1 if is_expansion else 0.6
+        
+        try:
+            with torch.inference_mode():
+                outputs = pipe(
+                    prompt_text,
+                    max_new_tokens=max_new_tokens,
+                    temperature=(temp if temp > 0.0 else None),
+                    do_sample=True if temp > 0.0 else False,
+                    pad_token_id=pipe.tokenizer.eos_token_id,
+                    return_full_text=False
+                )
+                
+            if isinstance(outputs, list) and len(outputs) > 0 and isinstance(outputs[0], dict):
+                text = outputs[0].get('generated_text', "")
+            elif isinstance(outputs, dict):
+                text = outputs.get('generated_text', "")
+            else:
+                text = str(outputs)
+                
+            text = text.strip()
+            for token in ['<|eot_id|>', '<|end_of_text|>']:
+                if token in text:
+                    text = text.split(token)[0].strip()
+            return text
+            
+        except Exception as e:
+            logger.error(f"Error calling LLM pipeline: {e}", exc_info=True)
+            raise HTTPException(status_code=500, detail=f"LLM generation failed: {str(e)}")
+
+    def _build_prompt(self, user_query: str, context: List[Dict], summary: str) -> List[Dict]:
+        context_text = "\n---\n".join([f"Source: {c.get('url', 'N/A')}\nChunk: {c['text']}" for c in context]) if context else "No relevant context found."
+        
+        system_prompt = (
+            "You are a helpful and harmless medical assistant, specialized in answering health-related questions "
+            "based ONLY on the provided retrieved context. Follow these strict rules:\n"
+            "1. **DO NOT** use any external knowledge. If the answer is not in the context, state that you cannot find "
+            "the information in the knowledge base.\n"
+            "2. Cite your sources using the URL/Source ID provided in the context (e.g., [Source: URL]). Do not generate fake URLs.\n"
+            "3. If the user's query is purely conversational, greet them or respond appropriately without referencing the context.\n"
+        )
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "system", "content": f"PREVIOUS CONVERSATION SUMMARY: {summary}" if summary else "PREVIOUS CONVERSATION SUMMARY: None"},
+            {"role": "system", "content": f"RETRIEVED CONTEXT:\n{context_text}"},
+            {"role": "user", "content": user_query}
+        ]
+        return messages
+
+    def _get_token_count(self, msg_list: List[Dict]) -> int:
+        if not self.intent_tokenizer:
+            return 0
+        prompt_text = self.intent_tokenizer.apply_chat_template(msg_list, tokenize=False, add_generation_prompt=True)
+        return len(self.intent_tokenizer.encode(prompt_text, add_special_tokens=False))
+
+    @method()
+    async def classify_intent(self, query: str) -> str:
+        """Classify query intent using the pre-loaded intent pipeline"""
+        if not self.intent_pipeline or not self.intent_tokenizer:
+            raise HTTPException(status_code=503, detail="Intent classification model not available")
+            
+        system_prompt = """You are a query classification robot. You MUST respond with ONLY ONE JSON object:
+        {"intent": "MEDICAL"}
+        {"intent": "GREET"}
+        {"intent": "OFF_TOPIC"}
+        {"intent": "HARMFUL"}
+        """
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": "Query: What are the symptoms of COVID-19?"},
+            {"role": "assistant", "content": '{"intent": "MEDICAL"}'},
+            {"role": "user", "content": f"Query: {query}"}
+        ]
+        
+        prompt_text = self.intent_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        
+        try:
+            llm_output = await self._run_with_timeout(
+                asyncio.to_thread(
+                    self._call_llm_pipeline, 
+                    self.intent_pipeline, 
+                    prompt_text, 
+                    True, 25, False
+                ),
+                # timeout_seconds=30,
+                timeout_message="Intent classification timed out"
+            )
+            
+            clean_output = llm_output.strip().replace("```json", "").replace("```", "")
+            start_idx = clean_output.find('{')
+            end_idx = clean_output.rfind('}')
+            if start_idx != -1 and end_idx != -1:
+                json_str = clean_output[start_idx: end_idx + 1]
+                data = json.loads(json_str)
+                return data.get("intent", "UNKNOWN")
+                
+        except Exception as e:
+            logger.error(f"Failed to parse JSON classifier output: {e}. Raw: {llm_output}")
+            
+        return self._rule_based_intent_classification(query)
+
+    def _rule_based_intent_classification(self, query: str) -> str:
+        """Fallback rule-based intent classification"""
+        query_lower = query.lower().strip()
+        
+        greeting_words = ['hello', 'hi', 'hey', 'greetings', 'good morning', 'good afternoon', 'how are you']
+        harmful_keywords = ['harm', 'hurt', 'kill', 'danger', 'illegal', 'prescription without', 'suicide']
+        medical_keywords = ['covid', 'fever', 'pain', 'symptom', 'treatment', 'medicine', 'doctor', 'health', 'disease', 'virus']
+        
+        if any(word in query_lower for word in greeting_words) or len(query_lower.split()) <= 2:
+            return 'GREET'
+        elif any(word in query_lower for word in harmful_keywords):
+            return 'HARMFUL'
+        elif not any(word in query_lower for word in medical_keywords) and len(query_lower.split()) > 3:
+            return 'OFF_TOPIC'
+        else:
+            return 'MEDICAL'
+
+    @method()
+    async def Greet(self, query: str) -> RAGResponse:
+        """Handle greeting queries"""
+        messages = [
+            {"role": "system", "content": "You are a greeter. Respond politely to the user greeting in a single line."},
+            {"role": "user", "content": query}
+        ]
+        
+        prompt_text = self.intent_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        
+        answer = await self._run_with_timeout(
+            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 50, True),
+            # timeout_seconds=30,
+            timeout_message="Greeting response timed out"
+        )
+        
+        return RAGResponse(
+            query=query, 
+            answer=answer, 
+            sources=[], 
+            context_chunks=[], 
+            expanded_queries=[]
+        )
+
+    @method()
+    async def HarmOff(self, query: str) -> RAGResponse:
+        """Handle harmful/off-topic queries"""
+        messages = [
+            {"role": "system", "content": "You are an intelligent assistant. Inform the user that you cannot answer harmful/off-topic questions. Keep it short and brief, in one sentence."},
+            {"role": "user", "content": query}
+        ]
+        
+        prompt_text = self.intent_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        
+        answer = await self._run_with_timeout(
+            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 50, True),
+            # timeout_seconds=30,
+            timeout_message="Safety response timed out"
+        )
+        
+        return RAGResponse(
+            query=query, 
+            answer=answer, 
+            sources=[], 
+            context_chunks=[], 
+            expanded_queries=[]
+        )
+
+    @method()
+    async def summarize_history(self, history: List[HistoryMessage]) -> str:
+        """Summarize conversation history"""
+        if not history:
+            return ''
+            
+        history_text = "\n".join([f"{h.role}: {h.content}" for h in history[-8:]])
+        summarizer_prompt = f"You are an intelligent agent who summarizes conversations. Your summary should be concise, coherent, and focus on the main topic and specific entities discussed.\nCONVERSATION HISTORY:\n{history_text}\nCONCISE SUMMARY:\n"
+        
+        messages = [{"role": "system", "content": summarizer_prompt}]
+        prompt_text = self.intent_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        
+        summary = await self._run_with_timeout(
+            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 150, False),
+            # timeout_seconds=60,
+            timeout_message="Summarization timed out"
+        )
+        return summary
+
+    @method()
+    async def expand_query_with_llm(self, user_query: str, summary: str, history: List[HistoryMessage]) -> List[str]:
+        """Expand query for better retrieval"""
+        if not history or len(history) == 0:
+            expansion_prompt = f"You are a specialized query expansion engine. Generate 3 alternative, highly effective search queries to find documents relevant to the User Query. Only output the queries, one per line. Do not include the original query or any explanations.\nUser Query: {user_query}\nExpanded Queries:\n"
+        else:
+            history_text = "\n".join([f"{h.role}: {h.content}" for h in history])
+            expansion_prompt = f"Given the conversation summary and history below, rewrite the user's latest query into a standalone, complete, and specific search query that incorporates the context of the conversation. Output only the single rewritten query.\nConversation Summary: {summary}\nConversation History:\n{history_text}\nUser's Latest Query: {user_query}\nRewritten Search Query:\n"
+        
+        messages = [{"role": "system", "content": expansion_prompt}]
+        prompt_text = self.intent_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        
+        llm_output = await self._run_with_timeout(
+            asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, True, 150, True),
+            # timeout_seconds=60,
+            timeout_message="Query expansion timed out"
+        )
+        
+        if not history or len(history) == 0:
+            expanded_queries = [q.strip() for q in llm_output.split('\n') if q.strip()]
+        else:
+            expanded_queries = [llm_output.strip()]
+        
+        expanded_queries.append(user_query)
+        seen = set()
+        deduped = []
+        for q in expanded_queries:
+            if q not in seen:
+                seen.add(q)
+                deduped.append(q)
+        return deduped
+
+    def retrieve_context(self, queries: List[str]) -> Tuple[List[Dict], List[str]]:
+        """Retrieve context from ChromaDB"""
+        if self.embedding_model is None:
+            raise HTTPException(status_code=503, detail="Embedding model not loaded.")
+            
+        embeddings_list = [[float(x) for x in emb] for emb in self.embedding_model.embed(queries, batch_size=8)]
+        results = self.chroma_collection.query(
+            query_embeddings=embeddings_list,
+            n_results=max(10, RETRIEVE_TOP_K_GPU * len(queries)),
+            include=['documents', 'metadatas']
+        )
+        
+        context_data = []
+        source_urls = set()
+        if results.get("documents") and results.get("metadatas"):
+            for docs_list, metadatas_list in zip(results["documents"], results["metadatas"]):
+                for doc, metadata in zip(docs_list, metadatas_list):
+                    if doc and metadata:
+                        context_data.append({'text': doc, 'url': metadata.get('source')})
+                        if metadata.get("source"):
+                            source_urls.add(metadata.get('source'))
+        return context_data, list(source_urls)
+
+    def rerank_documents(self, query: str, context: List[Dict], top_k: int) -> List[Dict]:
+        """Rerank documents using cross-encoder"""
+        if not context or self.cross_encoder is None:
+            return context[:top_k]
+            
+        pairs = [(query, doc['text']) for doc in context]
+        scores = self.cross_encoder.predict(pairs)
+        for doc, score in zip(context, scores):
+            doc['score'] = float(score)
+        ranked_docs = sorted(context, key=lambda x: x['score'], reverse=True)
+        return ranked_docs[:top_k]
+
+    @method()
+    async def prune_messages_to_fit_context(self, messages: List[Dict], final_context: List[Dict], summary: str, max_input_tokens: int) -> Tuple[List[Dict], List[Dict], int]:
+        """Prune messages to fit within token limit"""
+        if not self.intent_tokenizer:
+            raise ValueError("Tokenizer not initialized for pruning.")
+            
+        current_context = final_context[:]
+        current_summary = summary
+        base_user_query = messages[-1]["content"]
+        
+        current_messages = self._build_prompt(base_user_query, current_context, current_summary)
+        token_count = self._get_token_count(current_messages)
+        
+        if token_count <= max_input_tokens:
+            tok_length = max_input_tokens - token_count
+            return current_messages, current_context, tok_length
+            
+        logger.warning(f"Initial token count ({token_count}) exceeds max input ({max_input_tokens}). Starting pruning.")
+        
+        while token_count > max_input_tokens and current_context:
+            current_context.pop()
+            current_messages = self._build_prompt(base_user_query, current_context, current_summary)
+            token_count = self._get_token_count(current_messages)
+            
+        if token_count <= max_input_tokens:
+            tok_length = max_input_tokens - token_count
+            return current_messages, current_context, tok_length
+            
+        if current_summary:
+            logger.warning("Clearing conversation summary as last-ditch effort.")
+            current_summary = ""
+            current_messages = self._build_prompt(base_user_query, current_context, current_summary)
+            token_count = self._get_token_count(current_messages)
+            
+        if token_count <= max_input_tokens:
+            tok_length = max_input_tokens - token_count
+            return current_messages, current_context, tok_length
+            
+        logger.error(f"Pruning failed. Even minimal prompt exceeds token limit: {token_count}. Returning empty context.")
+        current_context = []
+        current_messages = self._build_prompt(base_user_query, current_context, "")
+        token_count = self._get_token_count(current_messages)
+        tok_length = max_input_tokens - token_count if token_count < max_input_tokens else 0
+        
+        return current_messages, current_context, tok_length
+
+    @fastapi_endpoint(method="POST")
+    async def rag_endpoint(self, request_data: Dict[str, Any]):
+        """Main RAG endpoint"""
+        try:
+            request = QueryRequest(**request_data)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid request format: {str(e)}")
+
+        start = time.time()
+        
+        try:
+            logger.info(f'Processing query: {request.query[:100]}...')
+            
+            intent = await self.classify_intent.remote.aio(request.query)
+            logger.info(f"Intent classified as: {intent}")
+
+            if intent == 'GREET':
+                response = await self.Greet.remote.aio(request.query)
+            elif intent in ["HARMFUL", "OFF_TOPIC"]:
+                response = await self.HarmOff.remote.aio(request.query)
+            else:
+                logger.info("Starting full RAG pipeline for medical query")
+                
+                summary = await self.summarize_history.remote.aio(request.history)
+                logger.info("History summarized")
+                
+                expanded_queries = await self.expand_query_with_llm.remote.aio(request.query, summary, request.history)
+                logger.info(f"Expanded queries: {expanded_queries}")
+                
+                context_data, _ = await self._run_with_timeout(
+                    asyncio.to_thread(self.retrieve_context, expanded_queries),
+                    timeout_message="Document retrieval timed out"
+                )
+                logger.info(f"Retrieved {len(context_data)} context chunks")
+                
+                final_context = await self._run_with_timeout(
+                    asyncio.to_thread(self.rerank_documents, request.query, context_data, RETRIEVE_TOP_K_GPU),
+                    # timeout_seconds=60,
+                    timeout_message="Document reranking timed out"
+                )
+                logger.info(f"Reranked to {len(final_context)} chunks")
+                
+                final_sources = list({c.get('url') for c in final_context if c.get('url')})
+                
+                if not final_context:
+                    final_answer = "I could not find relevant documents in the knowledge base to answer your question. I can help you if you have another question."
+                    context_chunks_text = []
+                else:
+                    initial_messages = self._build_prompt(request.query, final_context, summary)
+                    max_input_tokens = LLAMA_3_CONTEXT_WINDOW - MAX_NEW_TOKENS_GPU - SAFETY_BUFFER
+                    
+                    final_messages, final_context_pruned, tok_length = await self.prune_messages_to_fit_context.remote.aio(
+                        initial_messages, final_context, summary, max_input_tokens
+                    )
+                    
+                    context_chunks_text = [c['text'] for c in final_context_pruned]
+                    prompt_text = self.intent_tokenizer.apply_chat_template(final_messages, tokenize=False, add_generation_prompt=True)
+                    
+                    max_new = max(MAX_NEW_TOKENS_GPU, tok_length if isinstance(tok_length, int) and tok_length > 0 else MAX_NEW_TOKENS_GPU)
+                    
+                    final_answer = await self._run_with_timeout(
+                        asyncio.to_thread(self._call_llm_pipeline, self.intent_pipeline, prompt_text, False, max_new, False),
+                        # timeout_seconds=120,
+                        timeout_message="Answer generation timed out"
+                    )
+                    logger.info("Generated final answer")
+
+                response = RAGResponse(
+                    query=request.query,
+                    answer=final_answer,
+                    sources=final_sources,
+                    context_chunks=context_chunks_text,
+                    expanded_queries=expanded_queries
+                )
+
+            end_time = time.time()
+            logger.info(f"Total Latency: {round(end_time - start, 2)}s")
+            return response.model_dump()
+            
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Unhandled exception in RAG handler: {e}", exc_info=True)
+            raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+
+    async def _run_with_timeout(self, awaitable: Any, timeout_seconds: int = 300, timeout_message: str = "Request timed out") -> Any:
+        try:
+            return await asyncio.wait_for(awaitable, timeout=timeout_seconds)
+        except asyncio.TimeoutError:
+            logger.warning(f"Operation timed out after {timeout_seconds}s: {timeout_message}")
+            raise HTTPException(status_code=504, detail=timeout_message)
+        except HTTPException:
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error in _run_with_timeout: {e}")
+            raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")

requirements.txt

@@ -0,0 +1,12 @@
+torch 
+transformers 
+sentence-transformers 
+fastembed
+chromadb 
+pydantic 
+fastapi 
+requests
+python-json-logger
+boto3
+accelerate
+bitsandbytes

requirements_heavy.txt

@@ -1,15 +0,0 @@
-# torch
-# transformers
-# bitsandbytes
-# sentence-transformers
-# accelerate
-
-torch==2.9.0
-transformers
-# bitsandbytes
-sentence-transformers
-accelerate
-# llama-cpp-python UNCOMMENT THIS LINE FOR LOCAL DOCKER TESTING
-llama-cpp-python==0.2.83 --extra-index-url https://abetlen.github.io/llama-cpp-python-wheels/  
-tiktoken
-sentencepiece

requirements_light.txt

@@ -1,8 +0,0 @@
-fastapi
-uvicorn[standard]
-chromadb
-pydantic
-fastembed
-requests
-python-json-logger
-boto3

s3_utils.py

@@ -1,9 +1,8 @@
 from typing import Dict, List, Optional
-import boto3
+# import boto3
 import os
 import json
 import logging
-# from botocore.exceptions import NoCredentialsError, ClientError
 
 S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
 AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY_ID")
@@ -13,6 +12,8 @@ AWS_REGION = os.getenv("AWS_REGION")
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 def get_s3_client():
+    import boto3
+
     if not AWS_ACCESS_KEY or not AWS_SECRET_KEY:
         logging.warning("AWS credentials not found in environment. Using default config.")
         return boto3.client('s3', region_name=AWS_REGION)
@@ -25,10 +26,7 @@ def get_s3_client():
     )
 
 def download_chroma_folder_from_s3(s3_prefix: str, local_dir: str):
-    """
-    Downloads all files under s3_prefix from S3 to local_dir,
-    preserving the folder structure for ChromaDB.
-    """
+
     s3 = get_s3_client()
     paginator = s3.get_paginator("list_objects_v2")
     try: