Fix OpenAI client initialization and add robust error handling for Hugging Face compatibility

streamlit_app.py

@@ -95,148 +95,252 @@ Use these tools to provide the best possible answer.
 @st.cache_resource
 def load_document_chunks():
     """Load pre-processed document chunks from disk."""
-    with open(CHUNKS_FILE, 'rb') as f:
-        return pickle.load(f)
+    try:
+        print(f"Attempting to load chunks from: {CHUNKS_FILE}")
+        if not os.path.exists(CHUNKS_FILE):
+            print(f"ERROR: Chunks file not found at {CHUNKS_FILE}")
+            return []
+            
+        with open(CHUNKS_FILE, 'rb') as f:
+            chunks = pickle.load(f)
+        print(f"Successfully loaded {len(chunks)} document chunks")
+        return chunks
+    except Exception as e:
+        print(f"Error loading document chunks: {str(e)}")
+        # Try a direct load without caching
+        try:
+            print("Attempting direct load without caching")
+            with open(CHUNKS_FILE, 'rb') as f:
+                chunks = pickle.load(f)
+            print(f"Direct load successful: {len(chunks)} chunks")
+            return chunks
+        except Exception as e2:
+            print(f"Direct load also failed: {str(e2)}")
+            return []
 
 @st.cache_resource
 def get_chat_model():
     """Get the chat model for initial RAG."""
+    from openai import OpenAI
     import os
-    # Use openai_api_key and model_name for maximum compatibility
-    return ChatOpenAI(
-        openai_api_key=os.environ.get("OPENAI_API_KEY"),
-        model_name="gpt-4.1-mini",
-        temperature=0
-    )
+    
+    api_key = os.environ.get("OPENAI_API_KEY")
+    client = OpenAI(api_key=api_key)
+    
+    # Create a function with the same interface as ChatOpenAI.invoke
+    class SimpleOpenAIWrapper:
+        def __init__(self, client, model):
+            self.client = client
+            self.model = model
+            
+        def invoke(self, messages):
+            # Convert LangChain messages to OpenAI format
+            openai_messages = []
+            for msg in messages:
+                openai_messages.append({
+                    "role": msg.type if hasattr(msg, "type") else "user",
+                    "content": msg.content
+                })
+            
+            # Call the OpenAI API directly
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=openai_messages,
+                temperature=0
+            )
+            
+            # Create a simple object with a content attribute to match LangChain interface
+            class SimpleResponse:
+                def __init__(self, content):
+                    self.content = content
+            
+            return SimpleResponse(response.choices[0].message.content)
+    
+    # Return wrapper that matches the LangChain interface
+    return SimpleOpenAIWrapper(client, "gpt-4.1-mini")
 
 @st.cache_resource
 def get_agent_model():
     """Get the more powerful model for agent and evaluation."""
+    from openai import OpenAI
     import os
-    # Use openai_api_key and model_name for maximum compatibility
-    return ChatOpenAI(
-        openai_api_key=os.environ.get("OPENAI_API_KEY"),
-        model_name="gpt-4.1",
-        temperature=0
-    )
+    
+    api_key = os.environ.get("OPENAI_API_KEY")
+    client = OpenAI(api_key=api_key)
+    
+    # Create a function with the same interface as ChatOpenAI.invoke
+    class SimpleOpenAIWrapper:
+        def __init__(self, client, model):
+            self.client = client
+            self.model = model
+            
+        def invoke(self, messages):
+            # Convert LangChain messages to OpenAI format
+            openai_messages = []
+            for msg in messages:
+                openai_messages.append({
+                    "role": msg.type if hasattr(msg, "type") else "user",
+                    "content": msg.content
+                })
+            
+            # Call the OpenAI API directly
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=openai_messages,
+                temperature=0
+            )
+            
+            # Create a simple object with a content attribute to match LangChain interface
+            class SimpleResponse:
+                def __init__(self, content):
+                    self.content = content
+            
+            return SimpleResponse(response.choices[0].message.content)
+    
+    # Return wrapper that matches the LangChain interface
+    return SimpleOpenAIWrapper(client, "gpt-4.1")
 
 @st.cache_resource
 def get_embedding_model():
     """Get the embedding model."""
+    from openai import OpenAI
     import os
-    # Use openai_api_key and model_name for maximum compatibility
-    return OpenAIEmbeddings(
-        openai_api_key=os.environ.get("OPENAI_API_KEY"),
-        model_name="text-embedding-3-small"
-    )
+    import numpy as np
+    
+    api_key = os.environ.get("OPENAI_API_KEY")
+    client = OpenAI(api_key=api_key)
+    
+    # Create a wrapper class that matches the LangChain interface
+    class SimpleEmbeddings:
+        def __init__(self, client):
+            self.client = client
+        
+        def embed_query(self, text):
+            print(f"Embedding query: {text[:50]}...")
+            response = self.client.embeddings.create(
+                model="text-embedding-3-small",
+                input=text
+            )
+            return response.data[0].embedding
+        
+        def embed_documents(self, texts):
+            return [self.embed_query(text) for text in texts]
+    
+    return SimpleEmbeddings(client)
 
 @st.cache_resource
 def setup_qdrant_client():
     """Set up the Qdrant client."""
     import os
     
-    # DEBUG START - HF Compatibility fix
-    print(f"DEBUG: Setting up Qdrant client with path: {str(QDRANT_DIR)}")
-    print(f"DEBUG: Qdrant directory exists: {os.path.exists(QDRANT_DIR)}")
-    # DEBUG END
+    # Check for processed data directory
+    processed_data_dir_exists = os.path.exists(PROCESSED_DATA_DIR)
+    print(f"PROCESSED_DATA_DIR exists: {processed_data_dir_exists}")
+    print(f"Contents of current directory: {os.listdir('.')}")
+    
+    if processed_data_dir_exists:
+        print(f"Contents of PROCESSED_DATA_DIR: {os.listdir(PROCESSED_DATA_DIR)}")
+    
+    qdrant_dir_exists = os.path.exists(QDRANT_DIR)
+    print(f"QDRANT_DIR exists: {qdrant_dir_exists}")
     
-    # Check if directory exists
-    if not os.path.exists(QDRANT_DIR):
-        print(f"WARNING: Qdrant directory does not exist: {str(QDRANT_DIR)}")
-        raise ValueError(f"Qdrant directory not found at {str(QDRANT_DIR)}")
+    if qdrant_dir_exists:
+        print(f"Contents of QDRANT_DIR: {os.listdir(QDRANT_DIR)}")
     
-    # Try creating the client with minimal parameters
+    # Try creating the client with a simple path parameter
     try:
-        return QdrantClient(path=str(QDRANT_DIR))
+        client = QdrantClient(path=str(QDRANT_DIR))
+        print("Successfully created QdrantClient with path parameter")
+        return client
     except Exception as e:
-        # DEBUG START
-        print(f"DEBUG: Error initializing QdrantClient with path: {str(e)}")
-        # DEBUG END
+        print(f"Error creating QdrantClient with path: {str(e)}")
         
         # Try with location parameter
         try:
-            return QdrantClient(location=str(QDRANT_DIR))
+            client = QdrantClient(location=str(QDRANT_DIR))
+            print("Successfully created QdrantClient with location parameter")
+            return client
         except Exception as e2:
-            # DEBUG START
-            print(f"DEBUG: Error initializing with location: {str(e2)}")
-            # DEBUG END
+            print(f"Error creating QdrantClient with location: {str(e2)}")
             
-            # Last attempt with in-memory client
-            print("Attempting to create in-memory client")
+            # Last resort - try in-memory
+            print("Creating in-memory QdrantClient as fallback")
             return QdrantClient(":memory:")
 
 def retrieve_documents(query, k=5):
     """Retrieve relevant documents for a query."""
     # Define collection name
     collection_name = "kohavi_ab_testing_pdf_collection"
-    
-    # DEBUG START - HF Compatibility fix
-    print(f"DEBUG: Starting document retrieval for query: '{query[:30]}...'")
-    print(f"DEBUG: PROCESSED_DATA_DIR exists: {os.path.exists(PROCESSED_DATA_DIR)}")
-    print(f"DEBUG: CHUNKS_FILE exists: {os.path.exists(CHUNKS_FILE)}")
-    print(f"DEBUG: QDRANT_DIR exists: {os.path.exists(QDRANT_DIR)}")
-    # DEBUG END
+    print(f"======= QUERY: {query} =======")
     
     try:
+        # Check for processed data
+        print(f"CHUNKS_FILE exists: {os.path.exists(CHUNKS_FILE)}")
+        
         # Get models and data
         try:
             embedding_model = get_embedding_model()
+            print("Successfully created embedding model")
         except Exception as e:
-            # DEBUG START
-            print(f"DEBUG: Error getting embedding model: {str(e)}")
-            # DEBUG END
+            print(f"Error getting embedding model: {str(e)}")
+            # Try to fallback to direct API call instead of using LangChain
             return [], []
             
         try:
-            chunks = load_document_chunks()
-            # DEBUG START
-            print(f"DEBUG: Loaded {len(chunks)} document chunks")
-            # DEBUG END
+            print("Loading document chunks...")
+            if not os.path.exists(CHUNKS_FILE):
+                print(f"ERROR: CHUNKS_FILE not found at {CHUNKS_FILE}")
+                return [], []
+                
+            with open(CHUNKS_FILE, 'rb') as f:
+                chunks = pickle.load(f)
+            print(f"Successfully loaded {len(chunks)} document chunks")
         except Exception as e:
-            # DEBUG START
-            print(f"DEBUG: Error loading document chunks: {str(e)}")
-            # DEBUG END
+            print(f"Error loading document chunks: {str(e)}")
             return [], []
             
         try:
             client = setup_qdrant_client()
-            # DEBUG START
-            print("DEBUG: Successfully created Qdrant client")
-            # DEBUG END
+            print("Successfully created Qdrant client")
         except Exception as e:
-            # DEBUG START
-            print(f"DEBUG: Error setting up Qdrant client: {str(e)}")
-            # DEBUG END
+            print(f"Error setting up Qdrant client: {str(e)}")
             return [], []
         
         # Check if collection exists
         try:
             collections = client.get_collections()
-            # DEBUG START
-            print(f"DEBUG: Available collections: {collections}")
+            print(f"Available collections: {collections}")
             
             collection_info = client.get_collection(collection_name)
-            print(f"DEBUG: Collection info: {collection_info}")
-            # DEBUG END
+            print(f"Collection info: {collection_info}")
         except Exception as e:
-            # DEBUG START
-            print(f"DEBUG: Error checking collection: {str(e)}")
-            # DEBUG END
-            return [], []
+            print(f"Error checking collection: {str(e)}")
+            try:
+                # Try to initialize collection
+                print("Attempting to create collection...")
+                sample_embedding = embedding_model.embed_query("sample")
+                client.create_collection(
+                    collection_name=collection_name,
+                    vectors_config={
+                        "size": len(sample_embedding),
+                        "distance": "Cosine"
+                    }
+                )
+                print(f"Created new collection {collection_name}")
+            except Exception as e2:
+                print(f"Failed to create collection: {str(e2)}")
+                return [], []
         
         # Create a mapping of IDs to documents
         docs_by_id = {i: doc for i, doc in enumerate(chunks)}
         
         # Get query embedding
         try:
+            print(f"Generating embedding for query: {query}")
             query_embedding = embedding_model.embed_query(query)
-            # DEBUG START
-            print(f"DEBUG: Generated embedding of length {len(query_embedding)}")
-            # DEBUG END
+            print(f"Successfully generated embedding of length {len(query_embedding)}")
         except Exception as e:
-            # DEBUG START
-            print(f"DEBUG: Error creating query embedding: {str(e)}")
-            # DEBUG END
+            print(f"Error creating query embedding: {str(e)}")
             return [], []
         
         # Search for relevant documents
@@ -244,50 +348,27 @@ def retrieve_documents(query, k=5):
         
         # Try different querying approaches
         try:
-            # Simple query_points call
-            results = client.query_points(
+            print(f"Querying collection {collection_name}")
+            results = client.search(
                 collection_name=collection_name,
                 query_vector=query_embedding,
                 limit=k
             )
-            # DEBUG START
-            print(f"DEBUG: Retrieved {len(results)} results with query_points")
-            # DEBUG END
+            print(f"Retrieved {len(results)} results with search method")
         except Exception as e1:
-            # DEBUG START
-            print(f"DEBUG: First query approach failed: {str(e1)}")
-            # DEBUG END
-            
-            try:
-                # Try with minimum parameters
-                results = client.search(
-                    collection_name=collection_name,
-                    query_vector=query_embedding,
-                    limit=k
-                )
-                # DEBUG START
-                print(f"DEBUG: Retrieved {len(results)} results with search method")
-                # DEBUG END
-            except Exception as e2:
-                # DEBUG START
-                print(f"DEBUG: Second query approach failed: {str(e2)}")
-                # DEBUG END
-                return [], []
+            print(f"Search failed: {str(e1)}")
+            return [], []
         
         # Handle empty results
         if not results:
-            # DEBUG START
-            print("DEBUG: No results found in vector store")
-            # DEBUG END
+            print("No results found in vector store")
             return [], []
             
         # Process results
         documents = []
         sources_dict = {}
         
-        # DEBUG START
-        print(f"DEBUG: Processing {len(results)} search results")
-        # DEBUG END
+        print(f"Processing {len(results)} search results")
         
         for result in results:
             try:
@@ -325,27 +406,21 @@ def retrieve_documents(query, k=5):
                             "score": float(result.score),
                             "type": "pdf"
                         }
-                        # DEBUG START
-                        print(f"DEBUG: Added source: {title}, Page: {page}")
-                        # DEBUG END
+                        print(f"Added source: {title}, Page: {page}")
             except Exception as e:
-                # DEBUG START
-                print(f"DEBUG: Error processing result: {str(e)}")
-                # DEBUG END
+                print(f"Error processing result: {str(e)}")
                 continue
         
         # Convert sources dictionary to list
         sources = list(sources_dict.values())
         
-        # DEBUG START
-        print(f"DEBUG: Returning {len(documents)} documents with {len(sources)} unique sources")
-        # DEBUG END
+        print(f"Returning {len(documents)} documents with {len(sources)} unique sources")
         return documents, sources
         
     except Exception as e:
-        # DEBUG START
-        print(f"DEBUG: Unexpected error in retrieve_documents: {str(e)}")
-        # DEBUG END
+        print(f"Unexpected error in retrieve_documents: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return [], []
 
 def rephrase_query(query):