configured with openrouter

features/rag_chatbot/rag_pipeline.py

@@ -3,38 +3,60 @@ import chromadb
 from dotenv import load_dotenv
 from langchain_core.documents import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_openai import OpenAIEmbeddings, OpenAI
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms import OpenAI
 from langchain.chains.question_answering import load_qa_chain
 from langchain_community.vectorstores import Chroma
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
+from langchain.chat_models import ChatOpenAI
+
 
 load_dotenv()
 
 CHROMA_HOST = os.getenv("CHROMA_HOST", "localhost")
 COLLECTION_NAME = "company_docs_collection"
 
+# OpenRouter configuration
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
+
 vector_store = None
 company_qa_chain = None
 query_router_chain = None
 cybersecurity_chain = None
-llm = OpenAI(temperature=0)
+llm = None
 
 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:
-        embeddings = OpenAIEmbeddings()
+        # 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,
+        )
+
+        embeddings = HuggingFaceEmbeddings(
+            model_name="all-MiniLM-L6-v2",  
+            model_kwargs={'device': 'cpu'}, 
+            encode_kwargs={'normalize_embeddings': True}  # Normalize embeddings for better similarity search
+        )
+        
         
         # Initialize ChromaDB client 
         try:
             chroma_client = chromadb.HttpClient(host=CHROMA_HOST, port=8000)
             chroma_client.heartbeat() # Heartbeat check to confirm the connection
-            print("Successfully connected to ChromaDB.")
         except Exception as e:
-            print(f"FATAL: Could not connect to ChromaDB at {CHROMA_HOST}:8000. Please ensure the ChromaDB server is running.")
-            print(f"Error details: {e}")
             raise ConnectionError("Failed to connect to ChromaDB.") from e
         
         # Initialize vector store
@@ -45,16 +67,14 @@ def initialize_pipelines():
         )
 
         # Query Router Chain
-        router_template = """
-        You are a query classifier. Classify the following query into one of these categories:
-        - COMPANY: Questions about company policies, procedures, documents, or internal information
-        - CYBERSECURITY: Questions about cybersecurity, security threats, best practices, or vulnerabilities
-        - OFF_TOPIC: Questions that don't fit the above categories
-        
-        Query: {query}
-        
-        Respond with only the category name (COMPANY, CYBERSECURITY, or OFF_TOPIC):
-        """
+        router_template = """You are a query classifier. Classify the following query into one of these categories:
+- COMPANY: Questions about company policies, procedures, documents, or internal information
+- CYBERSECURITY: Questions about cybersecurity, security threats, best practices, or vulnerabilities
+- OFF_TOPIC: Questions that don't fit the above categories
+
+Query: {query}
+
+Respond with only the category name (COMPANY, CYBERSECURITY, or OFF_TOPIC):"""
         
         router_prompt = PromptTemplate(
             input_variables=["query"],
@@ -70,13 +90,11 @@ def initialize_pipelines():
         company_qa_chain = load_qa_chain(llm, chain_type="stuff")
         
         # Cybersecurity Chain
-        cybersecurity_template = """
-        You are a cybersecurity expert. Answer the following cybersecurity question based on your knowledge:
-        
-        Question: {question}
-        
-        Provide a comprehensive and accurate answer about cybersecurity:
-        """
+        cybersecurity_template = """You are a cybersecurity expert. Answer the following cybersecurity question based on your knowledge without claiming yourself as expert:
+
+Question: {question}
+
+Provide a comprehensive and accurate answer about cybersecurity:"""
         
         cybersecurity_prompt = PromptTemplate(
             input_variables=["question"],
@@ -88,8 +106,7 @@ def initialize_pipelines():
             prompt=cybersecurity_prompt
         )
         
-        print("All pipelines initialized successfully!")
-        
+
     except Exception as e:
         print(f"Error initializing pipelines: {e}")
         raise
@@ -112,7 +129,6 @@ def add_document_to_rag(text: str, metadata: dict):
             print("Document was empty after splitting, not adding to ChromaDB.")
             return False
 
-        print(f"Adding {len(docs)} document chunks to ChromaDB...")
         vector_store.add_documents(docs)
         print("Successfully added documents.")
         return True
@@ -133,7 +149,6 @@ def route_and_process_query(query: str):
         route_result = query_router_chain.run(query)
         route = route_result.strip().upper()
         
-        print(f"Query routed to: {route}")
 
         # 2. Route to appropriate logic
         if "CYBERSECURITY" in route:
@@ -147,8 +162,6 @@ def route_and_process_query(query: str):
         elif "COMPANY" in route:
             # Perform similarity search on ChromaDB
             docs = vector_store.similarity_search(query, k=3)
-            print(f"Found {len(docs)} relevant documents.")
-            print(f"Documents: {[doc.metadata.get('source', 'Unknown') for doc in docs]}")
 
             if not docs:
                 return {
@@ -195,7 +208,8 @@ def check_system_health():
             "vector_store": vector_store is not None,
             "company_qa_chain": company_qa_chain is not None,
             "query_router_chain": query_router_chain is not None,
-            "cybersecurity_chain": cybersecurity_chain is not None
+            "cybersecurity_chain": cybersecurity_chain is not None,
+            "llm": llm is not None
         }
         
         return {
@@ -209,6 +223,20 @@ def check_system_health():
             "error": str(e)
         }
 
+# Test function to verify OpenRouter connection
+def test_openrouter_connection():
+    """Test the OpenRouter API connection."""
+    try:
+        if not llm:
+            initialize_pipelines()
+        
+        # Simple test query
+        test_response = llm("Say 'Hello, OpenRouter is working!'")
+        return True
+    except Exception as e:
+        print(f"OpenRouter connection test failed: {e}")
+        return False
+
 # Initialize pipelines on module import
 try:
     initialize_pipelines()

requirements.txt

@@ -28,4 +28,6 @@ faiss-cpu
 PyPDF2
 tiktoken
 chromadb
-langchain_chroma
+langchain_chroma
+sentence-transformers
+tf-keras