Update app.py

app.py

@@ -55,13 +55,13 @@ is_initialized = False
 
 # Configuration
 class Config:
-    # OpenAI Compatible API Configuration
-    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
-    # REMOVED: OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
+    # API Keys - separate for each service
+    OPENAI_LLM_API_KEY = os.getenv("OPENAI_LLM_API_KEY", os.getenv("OPENAI_API_KEY", ""))
+    OPENAI_EMBEDDING_API_KEY = os.getenv("OPENAI_EMBEDDING_API_KEY", os.getenv("OPENAI_API_KEY", ""))
     
-    # ADDED: Separate base URLs for LLM and Embeddings
-    LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://api.openai.com/v1")
-    EMBEDDING_BASE_URL = os.getenv("EMBEDDING_BASE_URL", "https://api.openai.com/v1")
+    # Base URLs - separate for each service
+    OPENAI_LLM_BASE_URL = os.getenv("OPENAI_LLM_BASE_URL", os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"))
+    OPENAI_EMBEDDING_BASE_URL = os.getenv("OPENAI_EMBEDDING_BASE_URL", os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"))
     
     # Model Configuration
     LLM_MODEL = os.getenv("LLM_MODEL", "gpt-3.5-turbo")
@@ -104,18 +104,21 @@ class SystemStatus(BaseModel):
     is_initialized: bool
     model_name: str
     embedding_model: str
-    # CHANGED: Use separate URLs
     llm_base_url: str
     embedding_base_url: str
     vector_store_ready: bool
     total_chunks: int = 0
-    api_key_configured: bool
+    llm_api_key_configured: bool
+    embedding_api_key_configured: bool
 
 class InitializeRequest(BaseModel):
-    api_key: str = Field(..., min_length=1)
-    # CHANGED: Accept separate URLs
-    llm_base_url: Optional[str] = Field(default=None, description="LLM (text generation) API base URL")
-    embedding_base_url: Optional[str] = Field(default=None, description="Embedding model API base URL")
+    llm_api_key: Optional[str] = Field(default=None, description="API key for LLM service")
+    embedding_api_key: Optional[str] = Field(default=None, description="API key for embedding service")
+    # Backward compatibility - if provided, will be used for both services if individual keys not specified
+    api_key: Optional[str] = Field(default=None, description="Fallback API key for both services")
+    
+    llm_base_url: Optional[str] = Field(default=None, description="Base URL for LLM/text generation API")
+    embedding_base_url: Optional[str] = Field(default=None, description="Base URL for embedding API")
     llm_model: Optional[str] = Field(default=None, description="LLM model name")
     embedding_model: Optional[str] = Field(default=None, description="Embedding model name")
 
@@ -133,7 +136,7 @@ def estimate_tokens(text: str) -> int:
     except:
         return len(text.split()) * 1.3  # Rough estimate
 
-# Rate limiting helper functions (No changes needed here)
+# Rate limiting helper functions
 async def rate_limited_embedding_creation(chunks, embeddings):
     """Create embeddings with rate limiting to avoid API limits."""
     logger.info(f"Creating embeddings for {len(chunks)} chunks with rate limiting...")
@@ -196,7 +199,7 @@ async def rate_limited_embedding_creation(chunks, embeddings):
     logger.info("Successfully created and merged all embeddings")
     return final_vector_store
 
-# Custom Callback Handler for OpenAI (No changes needed here)
+# Custom Callback Handler for OpenAI
 class TokenUsageCallbackHandler(BaseCallbackHandler):
     """Callback handler to track token usage in OpenAI calls."""
     
@@ -247,29 +250,40 @@ class TokenUsageCallbackHandler(BaseCallbackHandler):
         }
 
 # RAG System Functions
-# CHANGED: Function signature to accept separate URLs
 async def initialize_rag_system(
-    api_key: str = None, 
+    llm_api_key: str = None, 
+    embedding_api_key: str = None,
+    api_key: str = None,  # Fallback for backward compatibility
     llm_base_url: str = None, 
     embedding_base_url: str = None, 
     llm_model: str = None, 
     embedding_model: str = None
 ):
-    """Initialize or reinitialize the RAG system with OpenAI compatible API."""
+    """Initialize or reinitialize the RAG system with separate OpenAI compatible APIs and keys."""
     global vector_store, qa_chain, token_callback_handler, is_initialized, config
     
     try:
-        # Update configuration
-        if api_key:
-            config.OPENAI_API_KEY = api_key
-        elif not config.OPENAI_API_KEY:
-            raise ValueError("OpenAI API key not provided")
+        # Handle API key configuration with fallback logic
+        if llm_api_key:
+            config.OPENAI_LLM_API_KEY = llm_api_key
+        elif api_key:
+            config.OPENAI_LLM_API_KEY = api_key
+        elif not config.OPENAI_LLM_API_KEY:
+            raise ValueError("LLM API key not provided")
+            
+        if embedding_api_key:
+            config.OPENAI_EMBEDDING_API_KEY = embedding_api_key
+        elif api_key:
+            config.OPENAI_EMBEDDING_API_KEY = api_key
+        elif not config.OPENAI_EMBEDDING_API_KEY:
+            raise ValueError("Embedding API key not provided")
         
-        # CHANGED: Update separate base URLs
+        # Update base URLs
         if llm_base_url:
-            config.LLM_BASE_URL = llm_base_url
+            config.OPENAI_LLM_BASE_URL = llm_base_url
+            
         if embedding_base_url:
-            config.EMBEDDING_BASE_URL = embedding_base_url
+            config.OPENAI_EMBEDDING_BASE_URL = embedding_base_url
             
         if llm_model:
             config.LLM_MODEL = llm_model
@@ -277,10 +291,11 @@ async def initialize_rag_system(
         if embedding_model:
             config.EMBEDDING_MODEL = embedding_model
         
-        # CHANGED: Update logging
         logger.info(f"Initializing RAG system with:")
-        logger.info(f"  - LLM Base URL: {config.LLM_BASE_URL}")
-        logger.info(f"  - Embedding Base URL: {config.EMBEDDING_BASE_URL}")
+        logger.info(f"  - LLM Base URL: {config.OPENAI_LLM_BASE_URL}")
+        logger.info(f"  - LLM API Key: {'*' * (len(config.OPENAI_LLM_API_KEY) - 8) + config.OPENAI_LLM_API_KEY[-8:] if len(config.OPENAI_LLM_API_KEY) > 8 else '*' * len(config.OPENAI_LLM_API_KEY)}")
+        logger.info(f"  - Embedding Base URL: {config.OPENAI_EMBEDDING_BASE_URL}")
+        logger.info(f"  - Embedding API Key: {'*' * (len(config.OPENAI_EMBEDDING_API_KEY) - 8) + config.OPENAI_EMBEDDING_API_KEY[-8:] if len(config.OPENAI_EMBEDDING_API_KEY) > 8 else '*' * len(config.OPENAI_EMBEDDING_API_KEY)}")
         logger.info(f"  - LLM Model: {config.LLM_MODEL}")
         logger.info(f"  - Embedding Model: {config.EMBEDDING_MODEL}")
         
@@ -302,14 +317,15 @@ async def initialize_rag_system(
         chunks = text_splitter.split_documents(documents)
         logger.info(f"Document split into {len(chunks)} chunks")
         
+        # Check if we have too many chunks that might cause rate limiting
         if len(chunks) > 200:
             logger.warning(f"Large number of chunks ({len(chunks)}). Consider increasing chunk_size to reduce API calls.")
         
-        # CHANGED: Initialize OpenAI embeddings with its specific base URL
+        # Initialize OpenAI embeddings with separate API key and base URL
         embeddings = OpenAIEmbeddings(
             model=config.EMBEDDING_MODEL,
-            openai_api_key=config.OPENAI_API_KEY,
-            openai_api_base=config.EMBEDDING_BASE_URL,
+            openai_api_key=config.OPENAI_EMBEDDING_API_KEY,  # Use embedding-specific API key
+            openai_api_base=config.OPENAI_EMBEDDING_BASE_URL,
             chunk_size=1000
         )
         
@@ -343,11 +359,11 @@ async def initialize_rag_system(
             vector_store.save_local(config.INDEX_PATH)
             logger.info(f"Created new FAISS index at '{config.INDEX_PATH}'")
         
-        # CHANGED: Initialize OpenAI LLM with its specific base URL
+        # Initialize OpenAI LLM with separate API key and base URL
         llm = ChatOpenAI(
             model_name=config.LLM_MODEL,
-            openai_api_key=config.OPENAI_API_KEY,
-            openai_api_base=config.LLM_BASE_URL,
+            openai_api_key=config.OPENAI_LLM_API_KEY,  # Use LLM-specific API key
+            openai_api_base=config.OPENAI_LLM_BASE_URL,
             temperature=config.TEMPERATURE,
             max_tokens=config.MAX_OUTPUT_TOKENS,
             callbacks=[token_callback_handler],
@@ -356,9 +372,6 @@ async def initialize_rag_system(
         
         # Test LLM connection
         try:
-            # Note: The os.environ is not strictly needed if passing params directly,
-            # but setting it can be a good practice for other potential library uses.
-            # We'll rely on direct parameter passing which is cleaner.
             test_response = llm.invoke("Test connection")
             logger.info("Successfully connected to LLM API")
         except Exception as e:
@@ -407,10 +420,9 @@ Answer:"""
 # API Endpoints
 @app.on_event("startup")
 async def startup_event():
-    """Initialize the system on startup if API key is available."""
-    if config.OPENAI_API_KEY:
+    """Initialize the system on startup if API keys are available."""
+    if config.OPENAI_LLM_API_KEY and config.OPENAI_EMBEDDING_API_KEY:
         try:
-            # This will use the URLs from environment variables by default
             await initialize_rag_system()
         except Exception as e:
             logger.warning(f"Could not initialize on startup: {str(e)}")
@@ -445,35 +457,37 @@ async def get_status():
         is_initialized=is_initialized,
         model_name=config.LLM_MODEL,
         embedding_model=config.EMBEDDING_MODEL,
-        # CHANGED: Return separate URLs
-        llm_base_url=config.LLM_BASE_URL,
-        embedding_base_url=config.EMBEDDING_BASE_URL,
+        llm_base_url=config.OPENAI_LLM_BASE_URL,
+        embedding_base_url=config.OPENAI_EMBEDDING_BASE_URL,
         vector_store_ready=vector_store is not None,
         total_chunks=len(vector_store.docstore._dict) if vector_store else 0,
-        api_key_configured=bool(config.OPENAI_API_KEY)
+        llm_api_key_configured=bool(config.OPENAI_LLM_API_KEY),
+        embedding_api_key_configured=bool(config.OPENAI_EMBEDDING_API_KEY)
     )
 
 @app.post("/api/initialize", response_model=Dict[str, Any])
 async def initialize_system(request: InitializeRequest):
-    """Initialize the RAG system with provided API key and configuration."""
+    """Initialize the RAG system with provided API keys and configuration."""
     try:
-        # CHANGED: Pass separate URLs to the initialization function
         await initialize_rag_system(
+            llm_api_key=request.llm_api_key,
+            embedding_api_key=request.embedding_api_key,
             api_key=request.api_key,
             llm_base_url=request.llm_base_url,
             embedding_base_url=request.embedding_base_url,
             llm_model=request.llm_model,
             embedding_model=request.embedding_model
         )
-        # CHANGED: Return separate URLs in the response
         return {
             "success": True,
             "message": "System initialized successfully",
             "config": {
-                "llm_base_url": config.LLM_BASE_URL,
-                "embedding_base_url": config.EMBEDDING_BASE_URL,
+                "llm_base_url": config.OPENAI_LLM_BASE_URL,
+                "embedding_base_url": config.OPENAI_EMBEDDING_BASE_URL,
                 "llm_model": config.LLM_MODEL,
-                "embedding_model": config.EMBEDDING_MODEL
+                "embedding_model": config.EMBEDDING_MODEL,
+                "llm_api_key_configured": bool(config.OPENAI_LLM_API_KEY),
+                "embedding_api_key_configured": bool(config.OPENAI_EMBEDDING_API_KEY)
             }
         }
     except Exception as e:
@@ -485,7 +499,7 @@ async def process_query(request: QueryRequest):
     if not is_initialized:
         raise HTTPException(
             status_code=503, 
-            detail="System not initialized. Please provide API key and configuration."
+            detail="System not initialized. Please provide API keys and configuration."
         )
     
     try:
@@ -539,8 +553,6 @@ async def process_query(request: QueryRequest):
         logger.error(f"Error processing query: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-# (No changes needed in the remaining endpoints)
-
 @app.get("/api/token-stats", response_model=Dict[str, Any])
 async def get_token_stats():
     """Get token usage statistics."""
@@ -568,7 +580,7 @@ async def upload_document(file: UploadFile = File(...)):
         logger.info(f"Uploaded new document: {file.filename}")
         
         # Reinitialize the system with new data
-        if config.OPENAI_API_KEY:
+        if config.OPENAI_LLM_API_KEY and config.OPENAI_EMBEDDING_API_KEY:
             # Remove old index to force recreation
             if os.path.exists(config.INDEX_PATH):
                 import shutil
@@ -592,22 +604,23 @@ async def health_check():
         "status": "healthy",
         "timestamp": datetime.now().isoformat(),
         "system_initialized": is_initialized,
-        "api_configured": bool(config.OPENAI_API_KEY)
+        "llm_api_configured": bool(config.OPENAI_LLM_API_KEY),
+        "embedding_api_configured": bool(config.OPENAI_EMBEDDING_API_KEY)
     }
 
 # Configuration endpoint
 @app.get("/api/config")
 async def get_config():
     """Get current configuration."""
-    # CHANGED: Return separate URLs
     return {
-        "llm_base_url": config.LLM_BASE_URL,
-        "embedding_base_url": config.EMBEDDING_BASE_URL,
+        "llm_base_url": config.OPENAI_LLM_BASE_URL,
+        "embedding_base_url": config.OPENAI_EMBEDDING_BASE_URL,
         "llm_model": config.LLM_MODEL,
         "embedding_model": config.EMBEDDING_MODEL,
         "chunk_size": config.CHUNK_SIZE,
         "retriever_k": config.RETRIEVER_K,
-        "api_key_configured": bool(config.OPENAI_API_KEY)
+        "llm_api_key_configured": bool(config.OPENAI_LLM_API_KEY),
+        "embedding_api_key_configured": bool(config.OPENAI_EMBEDDING_API_KEY)
     }
 
 # Mount static files