Update app.py

app.py

@@ -57,11 +57,15 @@ is_initialized = False
 class Config:
     # OpenAI Compatible API Configuration
     OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
-    OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")  # Can be changed to compatible APIs
+    # REMOVED: OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
+    
+    # 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")
     
     # Model Configuration
-    LLM_MODEL = os.getenv("LLM_MODEL", "gpt-3.5-turbo")  # Can be changed to any compatible model
-    EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "text-embedding-ada-002")  # Can be changed to compatible embedding model
+    LLM_MODEL = os.getenv("LLM_MODEL", "gpt-3.5-turbo")
+    EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "text-embedding-ada-002")
     
     # Document Processing
     CHUNK_SIZE = 500
@@ -70,12 +74,12 @@ class Config:
     # Rate Limiting
     MAX_RETRIES = 5
     RATE_LIMIT_DELAY = 2.0
-    EMBEDDING_BATCH_SIZE = 10  # OpenAI allows more requests
-    EMBEDDING_DELAY = 1.0  # Lower delay for OpenAI
+    EMBEDDING_BATCH_SIZE = 10
+    EMBEDDING_DELAY = 1.0
     
     # Model Parameters
     TEMPERATURE = 0.5
-    MAX_OUTPUT_TOKENS = 120000
+    MAX_OUTPUT_TOKENS = 2000
     RETRIEVER_K = 10
     
     # Paths
@@ -86,7 +90,7 @@ config = Config()
 
 # Request/Response Models
 class QueryRequest(BaseModel):
-    query: str = Field(..., min_length=1, max_length=120000)
+    query: str = Field(..., min_length=1, max_length=10000)
     
 class QueryResponse(BaseModel):
     answer: str
@@ -100,14 +104,18 @@ class SystemStatus(BaseModel):
     is_initialized: bool
     model_name: str
     embedding_model: str
-    base_url: 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
 
 class InitializeRequest(BaseModel):
     api_key: str = Field(..., min_length=1)
-    base_url: Optional[str] = Field(default=None, description="OpenAI compatible API base URL")
+    # 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_model: Optional[str] = Field(default=None, description="LLM model name")
     embedding_model: Optional[str] = Field(default=None, description="Embedding model name")
 
@@ -125,7 +133,7 @@ def estimate_tokens(text: str) -> int:
     except:
         return len(text.split()) * 1.3  # Rough estimate
 
-# Rate limiting helper functions
+# Rate limiting helper functions (No changes needed here)
 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...")
@@ -188,7 +196,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
+# Custom Callback Handler for OpenAI (No changes needed here)
 class TokenUsageCallbackHandler(BaseCallbackHandler):
     """Callback handler to track token usage in OpenAI calls."""
     
@@ -239,7 +247,14 @@ class TokenUsageCallbackHandler(BaseCallbackHandler):
         }
 
 # RAG System Functions
-async def initialize_rag_system(api_key: str = None, base_url: str = None, llm_model: str = None, embedding_model: str = None):
+# CHANGED: Function signature to accept separate URLs
+async def initialize_rag_system(
+    api_key: str = None, 
+    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."""
     global vector_store, qa_chain, token_callback_handler, is_initialized, config
     
@@ -247,13 +262,14 @@ async def initialize_rag_system(api_key: str = None, base_url: str = None, llm_m
         # Update configuration
         if api_key:
             config.OPENAI_API_KEY = api_key
-            os.environ["OPENAI_API_KEY"] = api_key
         elif not config.OPENAI_API_KEY:
             raise ValueError("OpenAI API key not provided")
         
-        if base_url:
-            config.OPENAI_BASE_URL = base_url
-            os.environ["OPENAI_BASE_URL"] = base_url
+        # CHANGED: Update separate base URLs
+        if llm_base_url:
+            config.LLM_BASE_URL = llm_base_url
+        if embedding_base_url:
+            config.EMBEDDING_BASE_URL = embedding_base_url
             
         if llm_model:
             config.LLM_MODEL = llm_model
@@ -261,8 +277,10 @@ async def initialize_rag_system(api_key: str = None, base_url: str = None, llm_m
         if embedding_model:
             config.EMBEDDING_MODEL = embedding_model
         
+        # CHANGED: Update logging
         logger.info(f"Initializing RAG system with:")
-        logger.info(f"  - Base URL: {config.OPENAI_BASE_URL}")
+        logger.info(f"  - LLM Base URL: {config.LLM_BASE_URL}")
+        logger.info(f"  - Embedding Base URL: {config.EMBEDDING_BASE_URL}")
         logger.info(f"  - LLM Model: {config.LLM_MODEL}")
         logger.info(f"  - Embedding Model: {config.EMBEDDING_MODEL}")
         
@@ -284,16 +302,15 @@ async def initialize_rag_system(api_key: str = None, base_url: str = None, llm_m
         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.")
         
-        # Initialize OpenAI embeddings
+        # CHANGED: Initialize OpenAI embeddings with its specific base URL
         embeddings = OpenAIEmbeddings(
             model=config.EMBEDDING_MODEL,
             openai_api_key=config.OPENAI_API_KEY,
-            openai_api_base=config.OPENAI_BASE_URL,
-            chunk_size=1000  # Embedding batch size
+            openai_api_base=config.EMBEDDING_BASE_URL,
+            chunk_size=1000
         )
         
         # Test embedding connection
@@ -326,11 +343,11 @@ async def initialize_rag_system(api_key: str = None, base_url: str = None, llm_m
             vector_store.save_local(config.INDEX_PATH)
             logger.info(f"Created new FAISS index at '{config.INDEX_PATH}'")
         
-        # Initialize OpenAI LLM
+        # CHANGED: Initialize OpenAI LLM with its specific base URL
         llm = ChatOpenAI(
             model_name=config.LLM_MODEL,
             openai_api_key=config.OPENAI_API_KEY,
-            openai_api_base=config.OPENAI_BASE_URL,
+            openai_api_base=config.LLM_BASE_URL,
             temperature=config.TEMPERATURE,
             max_tokens=config.MAX_OUTPUT_TOKENS,
             callbacks=[token_callback_handler],
@@ -339,6 +356,9 @@ async def initialize_rag_system(api_key: str = None, base_url: str = None, llm_m
         
         # 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:
@@ -390,6 +410,7 @@ async def startup_event():
     """Initialize the system on startup if API key is available."""
     if config.OPENAI_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)}")
@@ -424,7 +445,9 @@ async def get_status():
         is_initialized=is_initialized,
         model_name=config.LLM_MODEL,
         embedding_model=config.EMBEDDING_MODEL,
-        base_url=config.OPENAI_BASE_URL,
+        # CHANGED: Return separate URLs
+        llm_base_url=config.LLM_BASE_URL,
+        embedding_base_url=config.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)
@@ -434,17 +457,21 @@ async def get_status():
 async def initialize_system(request: InitializeRequest):
     """Initialize the RAG system with provided API key and configuration."""
     try:
+        # CHANGED: Pass separate URLs to the initialization function
         await initialize_rag_system(
             api_key=request.api_key,
-            base_url=request.base_url,
+            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": {
-                "base_url": config.OPENAI_BASE_URL,
+                "llm_base_url": config.LLM_BASE_URL,
+                "embedding_base_url": config.EMBEDDING_BASE_URL,
                 "llm_model": config.LLM_MODEL,
                 "embedding_model": config.EMBEDDING_MODEL
             }
@@ -512,6 +539,8 @@ 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."""
@@ -570,8 +599,10 @@ async def health_check():
 @app.get("/api/config")
 async def get_config():
     """Get current configuration."""
+    # CHANGED: Return separate URLs
     return {
-        "base_url": config.OPENAI_BASE_URL,
+        "llm_base_url": config.LLM_BASE_URL,
+        "embedding_base_url": config.EMBEDDING_BASE_URL,
         "llm_model": config.LLM_MODEL,
         "embedding_model": config.EMBEDDING_MODEL,
         "chunk_size": config.CHUNK_SIZE,