Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException
 import os
-from typing import List, Dict
+from typing import List, Dict, Any
 from dotenv import load_dotenv
 import logging
 from pathlib import Path
@@ -21,9 +21,12 @@ from langgraph.graph import END
 from langgraph.prebuilt import tools_condition
 from langgraph.checkpoint.memory import MemorySaver
 
-logging.basicConfig(level=logging.INFO)
+# Configure logging
+logging.basicConfig(level=logging.INFO, 
+                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# Load environment variables
 load_dotenv()
 GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
 GROQ_API_KEY = os.getenv('GROQ_API_KEY')
@@ -41,36 +44,73 @@ class QASystem:
         self.embeddings = None
         self.client = None
         self.pdf_dir = "pdfss"
+        self.is_initialized = False
 
     def load_pdf_documents(self):
+        """Load and process PDF documents from the pdf directory"""
         documents = []
         pdf_dir = Path(self.pdf_dir)
         
         if not pdf_dir.exists():
             raise FileNotFoundError(f"PDF directory not found: {self.pdf_dir}")
         
-        for pdf_path in pdf_dir.glob("*.pdf"):
+        pdf_files = list(pdf_dir.glob("*.pdf"))
+        if not pdf_files:
+            logger.warning(f"No PDF files found in directory: {self.pdf_dir}")
+            return []
+            
+        logger.info(f"Found {len(pdf_files)} PDF files to process")
+        
+        for pdf_path in pdf_files:
             try:
+                logger.info(f"Processing PDF: {pdf_path}")
                 loader = PyPDFLoader(str(pdf_path))
-                documents.extend(loader.load())
-                logger.info(f"Loaded PDF: {pdf_path}")
+                pdf_documents = loader.load()
+                
+                # Add source information to metadata
+                for doc in pdf_documents:
+                    if not hasattr(doc, 'metadata'):
+                        doc.metadata = {}
+                    doc.metadata['source'] = str(pdf_path.name)
+                
+                documents.extend(pdf_documents)
+                logger.info(f"Loaded PDF: {pdf_path} - {len(pdf_documents)} pages/sections")
             except Exception as e:
                 logger.error(f"Error loading PDF {pdf_path}: {str(e)}")
 
+        if not documents:
+            logger.warning("No documents were loaded from PDFs. Check the PDF directory and file formats.")
+            return []
+            
+        # Split documents into smaller chunks for better retrieval
         text_splitter = RecursiveCharacterTextSplitter(
             chunk_size=1000,
-            chunk_overlap=100
+            chunk_overlap=200
         )
         split_docs = text_splitter.split_documents(documents)
-        logger.info(f"Split documents into {len(split_docs)} chunks")
+        logger.info(f"Split {len(documents)} documents into {len(split_docs)} chunks")
+        
+        # Verify content of the first few chunks
+        for i, doc in enumerate(split_docs[:3]):
+            if i >= len(split_docs):
+                break
+            logger.info(f"Sample chunk {i+1} content preview: {doc.page_content[:100]}...")
+            
         return split_docs
 
     def initialize_system(self):
+        """Initialize the RAG system with vector store and LLM"""
         try:
+            logger.info("Initializing QA System...")
+            
+            # Initialize Qdrant client
             self.client = QdrantClient(":memory:")
+            logger.info("Qdrant client initialized (in-memory)")
             
+            # Create or get collection
             try:
-                self.client.get_collection("pdf_data")
+                collection_info = self.client.get_collection("pdf_data")
+                logger.info(f"Using existing collection: pdf_data")
             except Exception:
                 self.client.create_collection(
                     collection_name="pdf_data",
@@ -78,22 +118,32 @@ class QASystem:
                 )
                 logger.info("Created new collection: pdf_data")
             
+            # Initialize embeddings model
             self.embeddings = GoogleGenerativeAIEmbeddings(
                 model="models/embedding-001",
                 google_api_key=GOOGLE_API_KEY
             )
+            logger.info("Google AI Embeddings initialized")
             
+            # Initialize vector store
             self.vector_store = QdrantVectorStore(
                 client=self.client,
                 collection_name="pdf_data",
                 embeddings=self.embeddings,
             )
+            logger.info("Qdrant vector store initialized")
             
+            # Load documents
             documents = self.load_pdf_documents()
+            if not documents:
+                logger.warning("No documents loaded. The system will continue but may not provide relevant responses.")
+            
+            # Clear existing vectors if any
             if documents:
                 try:
-                    points = self.client.scroll(collection_name="pdf_data", limit=100)[0]
+                    points = self.client.scroll(collection_name="pdf_data", limit=1000)[0]
                     if points:
+                        logger.info(f"Clearing {len(points)} existing vectors from collection")
                         self.client.delete(
                             collection_name="pdf_data",
                             points_selector=PointIdsList(
@@ -103,65 +153,104 @@ class QASystem:
                 except Exception as e:
                     logger.error(f"Error clearing vectors: {str(e)}")
                 
+                # Add documents to vector store
+                logger.info(f"Adding {len(documents)} documents to vector store")
                 self.vector_store.add_documents(documents)
-                logger.info(f"Added {len(documents)} documents to vector store")
+                logger.info(f"Successfully added documents to vector store")
+                
+                # Verify vector store has documents
+                try:
+                    count = len(self.client.scroll(collection_name="pdf_data", limit=1)[0])
+                    logger.info(f"Vector store contains points: {count > 0}")
+                except Exception as e:
+                    logger.error(f"Error verifying vector store: {str(e)}")
 
+            # Initialize LLM
             llm = ChatGroq(
                 model="llama3-8b-8192", 
                 api_key=GROQ_API_KEY,
                 temperature=0.7
             )
+            logger.info("Groq LLM initialized")
             
+            # Create LangGraph
             graph_builder = StateGraph(MessagesState)
+            logger.info("Creating LangGraph for conversation flow")
 
-            # Define a retrieval node that fetches relevant docs
+            # Define retrieval node (self reference for vector_store access)
+            vector_store_ref = self.vector_store
+            
             def retrieve_docs(state: MessagesState):
+                """Node that retrieves relevant documents from the vector store"""
                 # Get the most recent human message
                 human_messages = [m for m in state["messages"] if m.type == "human"]
                 if not human_messages:
+                    logger.warning("No human messages found in state")
                     return {"messages": state["messages"]}
                 
                 user_query = human_messages[-1].content
-                logger.info(f"Retrieving documents for query: {user_query}")
+                logger.info(f"Retrieving documents for query: '{user_query}'")
+                
+                # Check if vector store exists
+                if not vector_store_ref:
+                    logger.error("Vector store not initialized or empty")
+                    return {"messages": state["messages"]}
                 
                 # Query the vector store
                 try:
-                    retrieved_docs = self.vector_store.similarity_search(user_query, k=3)
+                    retrieved_docs = vector_store_ref.similarity_search(user_query, k=3)
+                    
+                    if not retrieved_docs:
+                        logger.warning(f"No documents retrieved for query: '{user_query}'")
+                        return {"messages": state["messages"]}
+                    
+                    # Log what was actually retrieved
+                    for i, doc in enumerate(retrieved_docs):
+                        source = doc.metadata.get('source', 'Unknown') if hasattr(doc, 'metadata') else 'Unknown'
+                        content_preview = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
+                        logger.info(f"Retrieved doc {i+1} from {source}, preview: {content_preview}")
                     
-                    # Create tool messages for each retrieved document
+                    # Create tool messages with more detailed content
                     tool_messages = []
                     for i, doc in enumerate(retrieved_docs):
+                        # Include source information if available
+                        source_info = f" (Source: {doc.metadata.get('source', 'Unknown')})" if hasattr(doc, 'metadata') else ""
+                        
                         tool_messages.append(
                             ToolMessage(
-                                content=f"Document {i+1}: {doc.page_content}",
+                                content=f"Document {i+1}{source_info}: {doc.page_content}",
                                 tool_call_id=f"retrieval_{i}"
                             )
                         )
                     
-                    logger.info(f"Retrieved {len(tool_messages)} relevant documents")
+                    logger.info(f"Created {len(tool_messages)} tool messages with retrieved content")
                     return {"messages": state["messages"] + tool_messages}
                 
                 except Exception as e:
                     logger.error(f"Error retrieving documents: {str(e)}")
                     return {"messages": state["messages"]}
 
-            # Updated generate function that uses retrieved documents
+            # Generate response using retrieved documents
             def generate(state: MessagesState):
+                """Node that generates a response using the LLM and retrieved documents"""
                 # Extract retrieved documents (tool messages)
                 tool_messages = [m for m in state["messages"] if m.type == "tool"]
                 
                 # Collect context from retrieved documents
                 if tool_messages:
-                    context = "\n".join([m.content for m in tool_messages])
+                    context = "\n\n".join([m.content for m in tool_messages])
                     logger.info(f"Using context from {len(tool_messages)} retrieved documents")
                 else:
-                    context = "No specific mountain bicycle documentation available."
-                    logger.info("No relevant documents retrieved, using default context")
+                    context = "No specific mountain bicycle documentation available for this query."
+                    logger.warning("No relevant documents retrieved, using default context")
 
                 system_prompt = (
                     "You are an AI assistant embedded within the Interactive Electronic Technical Manual (IETM) for Mountain Cycles. "
-                    "Always provide accurate responses with references to provided data. "
-                    "If the user query is not technical-specific, still respond from a IETM perspective."
+                    "Your primary role is to provide accurate technical information about mountain bicycles. "
+                    "Always base your responses on the provided documentation. "
+                    "If you don't find specific information in the provided context, clearly state that the information "
+                    "is not available in the current documentation instead of making up details. "
+                    "When responding, reference specific parts of the documentation."
                     f"\n\nContext from mountain bicycle documentation:\n{context}"
                 )
 
@@ -174,8 +263,14 @@ class QASystem:
                 logger.info(f"Sending query to LLM with {len(messages)} messages")
                 
                 # Generate the response
-                response = llm.invoke(messages)
-                return {"messages": state["messages"] + [response]}
+                try:
+                    response = llm.invoke(messages)
+                    logger.info(f"LLM generated response successfully")
+                    return {"messages": state["messages"] + [response]}
+                except Exception as e:
+                    logger.error(f"Error generating response: {str(e)}")
+                    error_message = SystemMessage(content=f"Error generating response: {str(e)}")
+                    return {"messages": state["messages"] + [error_message]}
 
             # Add nodes to the graph
             graph_builder.add_node("retrieve_docs", retrieve_docs)
@@ -186,22 +281,35 @@ class QASystem:
             graph_builder.add_edge("retrieve_docs", "generate")
             graph_builder.add_edge("generate", END)
             
+            # Initialize memory
             self.memory = MemorySaver()
             self.graph = graph_builder.compile(checkpointer=self.memory)
+            logger.info("Graph compiled successfully")
+            
+            self.is_initialized = True
             return True
 
         except Exception as e:
             logger.error(f"System initialization error: {str(e)}")
+            self.is_initialized = False
             return False
 
-    def process_query(self, query: str) -> Dict[str, str]:
+    def process_query(self, query: str) -> Dict[str, Any]:
         """Process a query and return a single final response"""
         try:
-            # Generate a unique thread ID for production use
-            # For simplicity, using a fixed ID here
+            if not self.is_initialized:
+                logger.error("System not initialized. Cannot process query.")
+                return {
+                    'content': "Error: QA System not initialized properly",
+                    'type': 'error'
+                }
+                
+            logger.info(f"Processing query: '{query}'")
+            
+            # Generate a thread ID (use a more sophisticated method for production)
             thread_id = "abc123"
             
-            # Use invoke instead of stream to get only the final result
+            # Use invoke to get only the final result
             final_state = self.graph.invoke(
                 {"messages": [HumanMessage(content=query)]},
                 config={"configurable": {"thread_id": thread_id}}
@@ -211,31 +319,35 @@ class QASystem:
             ai_messages = [m for m in final_state["messages"] if m.type == "ai"]
             
             if ai_messages:
+                logger.info("Successfully generated response")
                 # Return only the last AI message
                 return {
                     'content': ai_messages[-1].content,
                     'type': ai_messages[-1].type
                 }
+            
+            logger.warning("No AI message generated in response")
             return {
-                'content': "No response generated",
+                'content': "No response could be generated for your query. Please try a different question.",
                 'type': 'error'
             }
             
         except Exception as e:
             logger.error(f"Query processing error: {str(e)}")
             return {
-                'content': f"Query processing error: {str(e)}",
+                'content': f"Error processing your query: {str(e)}",
                 'type': 'error'
             }
 
+# Initialize the QA system
 qa_system = QASystem()
-if qa_system.initialize_system():
-    logger.info("QA System Initialized Successfully")
-else:
-    raise RuntimeError("Failed to initialize QA System")
+initialization_success = qa_system.initialize_system()
 
 @app.post("/query")
 async def query_api(query: str):
     """API endpoint that returns a single response for a query"""
+    if not qa_system.is_initialized:
+        raise HTTPException(status_code=500, detail="QA System not initialized properly")
+        
     response = qa_system.process_query(query)
     return {"response": response}