Added 3 model selection with env support

.env-example

@@ -31,4 +31,4 @@ MY_SECRET_TOKEN="SECRET_CODE_TOKEN"
 # LLM_TEMPERATURE=0.1
 
 # Maximum tokens for response
-# LLM_MAX_TOKENS=4096
+# LLM_MAX_TOKENS=4096

features/rag_chatbot/rag_pipeline.py

@@ -14,11 +14,32 @@ from langchain.chat_models import ChatOpenAI
 
 load_dotenv()
 
-CHROMA_HOST = os.getenv("CHROMA_HOST", "localhost")
+# ChromaDB configuration
+CHROMA_HOST = os.getenv("CHROMA_HOST", "localhost") # change in env in production when hosted
 COLLECTION_NAME = "company_docs_collection"
 
-# OpenRouter configuration
-OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
+# LLM Provider Configuration
+LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower()  
+LLM_API_KEY = os.getenv("LLM_API_KEY")
+LLM_MODEL = os.getenv("LLM_MODEL", "gpt-3.5-turbo")
+LLM_TEMPERATURE = float(os.getenv("LLM_TEMPERATURE", "0"))
+LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "2048"))
+
+# Provider-specific configurations
+PROVIDER_CONFIGS = {
+    "openai": {
+        "api_base": "https://api.openai.com/v1",
+        "default_model": "gpt-3.5-turbo"
+    },
+    "groq": {
+        "api_base": "https://api.groq.com/openai/v1",
+        "default_model": "llama-3.3-70b-versatile"
+    },
+    "openrouter": {
+        "api_base": "https://openrouter.ai/api/v1",
+        "default_model": "mistralai/mistral-small-3.2-24b-instruct:free"
+    }
+}
 
 vector_store = None
 company_qa_chain = None
@@ -26,36 +47,54 @@ query_router_chain = None
 cybersecurity_chain = None
 llm = None
 
+def get_llm_config():
+    """Get the appropriate LLM configuration based on the provider."""
+    if LLM_PROVIDER not in PROVIDER_CONFIGS:
+        raise ValueError(f"Unsupported LLM provider: {LLM_PROVIDER}. Supported: {list(PROVIDER_CONFIGS.keys())}")
+    
+    config = PROVIDER_CONFIGS[LLM_PROVIDER].copy()
+    
+    # Use provided model or fall back to default
+    model = LLM_MODEL if LLM_MODEL != "gpt-3.5-turbo" else config["default_model"]
+    
+    return {
+        "model": model,
+        "openai_api_key": LLM_API_KEY,
+        "openai_api_base": config["api_base"],
+        "temperature": LLM_TEMPERATURE,
+        "max_tokens": LLM_MAX_TOKENS,
+    }
+
+def initialize_llm():
+    """Initialize the LLM based on the configured provider."""
+    if not LLM_API_KEY:
+        raise ValueError(f"LLM_API_KEY environment variable is required for {LLM_PROVIDER}")
+    
+    config = get_llm_config()
+    
+    print(f"Initializing {LLM_PROVIDER.upper()} with model: {config['model']}")
+    
+    return ChatOpenAI(**config)
+
 def initialize_pipelines():
     """Initializes all required models, chains, and the vector store."""
     global vector_store, company_qa_chain, query_router_chain, cybersecurity_chain, llm
     
     try:
-        # Check for required API keys
-        if not OPENROUTER_API_KEY:
-            raise ValueError("OPENROUTER_API_KEY environment variable is required")
-        
-        
-        # Initialize LLM with OpenRouter
-        llm = ChatOpenAI(
-            model="meta-llama/llama-3.3-70b-instruct:free",
-            openai_api_key=OPENROUTER_API_KEY,
-            openai_api_base="https://openrouter.ai/api/v1",
-            temperature=0,
-            max_tokens=2048,
-        )
+        # Initialize LLM
+        llm = initialize_llm()
 
+        # Initialize embeddings
         embeddings = HuggingFaceEmbeddings(
             model_name="all-MiniLM-L6-v2",  
             model_kwargs={'device': 'cpu'}, 
-            encode_kwargs={'normalize_embeddings': True}  # Normalize embeddings for better similarity search
+            encode_kwargs={'normalize_embeddings': True}
         )
         
-        
         # Initialize ChromaDB client 
         try:
             chroma_client = chromadb.HttpClient(host=CHROMA_HOST, port=8000)
-            chroma_client.heartbeat() # Heartbeat check to confirm the connection
+            chroma_client.heartbeat()
         except Exception as e:
             raise ConnectionError("Failed to connect to ChromaDB.") from e
         
@@ -86,8 +125,8 @@ Respond with only the category name (COMPANY, CYBERSECURITY, or OFF_TOPIC):"""
             prompt=router_prompt
         )
         
-        # Custom Company QA Chain with natural prompt
-        company_qa_template = """You are a helpful assistant for CyberAlertNepal. Answer the following question about our company using the information provided and links if only available. Give a natural, direct and polite response .
+        # Custom Company QA Chain
+        company_qa_template = """You are a helpful assistant for CyberAlertNepal. Answer the following question about our company using the information provided and links if only available. Give a natural, direct and polite response.
 
 Question: {question}
 
@@ -109,8 +148,7 @@ Answer:"""
         # Cybersecurity Chain
         cybersecurity_template = """You are a cybersecurity professional. Answer the following question truthfully and concisely. 
 If you are not 100% sure about the answer, simply respond with: "I am not sure about the answer." 
-Do not add extra explanations or assumptions. Do not provide false or speculative information. 
-
+Do not add extra explanations or assumptions. Do not provide false or speculative information.
 
 Question: {question}
 
@@ -126,6 +164,7 @@ Provide a comprehensive and accurate answer about cybersecurity:"""
             prompt=cybersecurity_prompt
         )
         
+        print(f"Successfully initialized pipelines with {LLM_PROVIDER.upper()}")
 
     except Exception as e:
         print(f"Error initializing pipelines: {e}")
@@ -176,7 +215,9 @@ def route_and_process_query(query: str):
             return {
                 "answer": answer, 
                 "source": "Cybersecurity Knowledge Base",
-                "route": "CYBERSECURITY"
+                "route": "CYBERSECURITY",
+                "provider": LLM_PROVIDER.upper(),
+                "model": get_llm_config()["model"]
             }
             
         elif "COMPANY" in route:
@@ -187,7 +228,9 @@ def route_and_process_query(query: str):
                 return {
                     "answer": "I could not find any relevant information to answer your question.",
                     "source": "Company Documents",
-                    "route": "COMPANY"
+                    "route": "COMPANY",
+                    "provider": LLM_PROVIDER.upper(),
+                    "model": get_llm_config()["model"]
                 }
             
             # Combine document content for context
@@ -201,14 +244,18 @@ def route_and_process_query(query: str):
                 "answer": answer,
                 "source": "Company Documents",
                 "documents": sources,
-                "route": "COMPANY"
+                "route": "COMPANY",
+                "provider": LLM_PROVIDER.upper(),
+                "model": get_llm_config()["model"]
             }
 
         else:  # OFF_TOPIC
             return {
                 "answer": "I am a specialized assistant of CyberAlertNepal. I cannot answer questions outside of cybersecurity topics.",
                 "source": "N/A",
-                "route": "OFF_TOPIC"
+                "route": "OFF_TOPIC",
+                "provider": LLM_PROVIDER.upper(),
+                "model": get_llm_config()["model"]
             }
             
     except Exception as e:
@@ -216,6 +263,9 @@ def route_and_process_query(query: str):
         return {
             "answer": "I encountered an error while processing your query. Please try again.",
             "source": "Error",
+            "route": None,
+            "documents": None,
+            "provider": LLM_PROVIDER.upper(),
             "error": str(e)
         }
 
@@ -237,28 +287,38 @@ def check_system_health():
         
         return {
             "status": "healthy" if all(components.values()) else "unhealthy",
-            "components": components
+            "components": components,
+            "provider": LLM_PROVIDER.upper(),
+            "model": get_llm_config()["model"] if llm else "Not initialized"
         }
         
     except Exception as e:
         return {
             "status": "unhealthy",
-            "error": str(e)
+            "error": str(e),
+            "provider": LLM_PROVIDER.upper()
         }
 
-# Test function to verify OpenRouter connection
-def test_openrouter_connection():
-    """Test the OpenRouter API connection."""
+def test_llm_connection():
+    """Test the LLM API connection."""
     try:
         if not llm:
             initialize_pipelines()
         
         # Simple test query
-        test_response = llm("Say 'Hello, OpenRouter is working!'")
-        return True
+        test_response = llm("Say 'Hello, LLM is working!'")
+        return {
+            "success": True,
+            "provider": LLM_PROVIDER.upper(),
+            "model": get_llm_config()["model"],
+            "response": str(test_response)
+        }
     except Exception as e:
-        print(f"OpenRouter connection test failed: {e}")
-        return False
+        return {
+            "success": False,
+            "provider": LLM_PROVIDER.upper(),
+            "error": str(e)
+        }
 
 # Initialize pipelines on module import
 try: