Fix Haystack 2.1.0 compatibility

pipelines.py

@@ -7,17 +7,8 @@ from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.components.embedders import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder
 from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
 
-# Robust import handling for Haystack 2.1.0
-try:
-    # Try the new import structure (Haystack >= 2.1.1)
-    from haystack.components.rankers.sentence_transformers import SentenceTransformersSimilarityRanker
-except ImportError:
-    try:
-        # Try the direct import (Haystack 2.1.0)
-        from haystack.components.rankers import SentenceTransformersSimilarityRanker
-    except ImportError:
-        # Fallback to legacy import
-        from haystack.nodes.ranker import SentenceTransformersRanker as SentenceTransformersSimilarityRanker
+# CORRECT IMPORT FOR HAYSTACK 2.1.0
+from haystack.nodes.ranker import SentenceTransformersRanker
 
 from haystack_integrations.components.generators.google_ai import GoogleAIGeminiGenerator
 from haystack.components.preprocessors import DocumentSplitter
@@ -32,29 +23,19 @@ document_store = InMemoryDocumentStore()
 # Optimized for CPU
 doc_embedder = SentenceTransformersDocumentEmbedder(
     model="BAAI/bge-base-en-v1.5",
-    use_gpu=False,
-    onnx_execution_provider="CPUExecutionProvider"
+    use_gpu=False
 )
 text_embedder = SentenceTransformersTextEmbedder(
     model="BAAI/bge-base-en-v1.5",
-    use_gpu=False,
-    onnx_execution_provider="CPUExecutionProvider"
+    use_gpu=False
 )
 retriever = InMemoryEmbeddingRetriever(document_store=document_store, top_k=3)
 
-# Initialize ranker based on what was imported
-if "SentenceTransformersRanker" in globals():
-    # Using the legacy ranker
-    reranker = SentenceTransformersRanker(
-        model_name_or_path="cross-encoder/ms-marco-TinyBERT-L-2-v2",
-        use_gpu=False
-    )
-else:
-    # Using the new ranker
-    reranker = SentenceTransformersSimilarityRanker(
-        model="cross-encoder/ms-marco-TinyBERT-L-2-v2",
-        use_gpu=False
-    )
+# Initialize ranker - DIFFERENT INITIALIZATION FOR 2.1.0
+reranker = SentenceTransformersRanker(
+    model_name_or_path="cross-encoder/ms-marco-TinyBERT-L-2-v2",
+    use_gpu=False
+)
 
 # Initialize generator
 generator = GoogleAIGeminiGenerator(
@@ -79,20 +60,13 @@ try:
     logger.info("Warming up components...")
     doc_embedder.warm_up()
     text_embedder.warm_up()
-    
-    # Handle different warm_up methods
-    if hasattr(reranker, 'warm_up'):
-        reranker.warm_up()
-    elif hasattr(reranker, 'prepared'):
-        reranker.prepared = True  # Legacy versions didn't require warm_up
-        
+    reranker.prepared = True  # Different warmup for 2.1.0
     logger.info("Components warmed up")
 except Exception as e:
     logger.error(f"Warmup failed: {e}")
 
 def add_documents(texts: list[str], meta_list: list[dict]) -> int:
     """Process and store documents with chunking"""
-    # Create base documents
     docs = [
         Document(content=text, meta=meta)
         for text, meta in zip(texts, meta_list)
@@ -102,14 +76,12 @@ def add_documents(texts: list[str], meta_list: list[dict]) -> int:
     if not docs:
         return 0
         
-    # Split into chunks
     split_result = splitter.run(docs)
     split_docs = split_result.get("documents", [])
     
     if not split_docs:
         return 0
         
-    # Batch embedding with reduced batch size
     embedded_docs = []
     batch_size = 8
     
@@ -128,60 +100,34 @@ def add_documents(texts: list[str], meta_list: list[dict]) -> int:
 def query_rag(question: str, session_id: str) -> dict:
     """Query the RAG system with session filtering"""
     try:
-        # Validate input
         if not question.strip():
-            return {
-                "answer": "Please provide a non-empty question.",
-                "sources": []
-            }
+            return {"answer": "Please provide a non-empty question.", "sources": []}
             
-        # Embed question
         embedding_result = text_embedder.run(question)
         query_emb = embedding_result.get("embedding")
         
         if not query_emb:
-            return {
-                "answer": "Failed to process your question.",
-                "sources": []
-            }
+            return {"answer": "Failed to process your question.", "sources": []}
         
-        # Retrieve documents with session filter
         filters = {"field": "meta.session_id", "operator": "==", "value": session_id}
-        retrieved_docs = retriever.run(
-            query_embedding=query_emb, 
-            filters=filters
-        ).get("documents", [])
+        retrieved_docs = retriever.run(query_embedding=query_emb, filters=filters).get("documents", [])
         
         if not retrieved_docs:
-            return {
-                "answer": "No documents found for this session. Please upload a file first.",
-                "sources": []
-            }
+            return {"answer": "No documents found. Upload a file first.", "sources": []}
         
-        # Handle different ranker interfaces
-        if hasattr(reranker, 'run'):
-            # New interface
-            reranked_docs = reranker.run(
-                query=question, 
-                documents=retrieved_docs[:5]
-            ).get("documents", [])[:3]
-        else:
-            # Legacy interface
-            reranked_docs = reranker.predict(
-                query=question,
-                documents=retrieved_docs[:5],
-                top_k=3
-            )
+        # DIFFERENT USAGE FOR 2.1.0 RANKER
+        reranked_docs = reranker.predict(
+            query=question,
+            documents=retrieved_docs[:5],
+            top_k=3
+        )
         
-        # Generate answer with context
         context = "\n\n".join([doc.content for doc in reranked_docs])
         prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
         
-        # Handle generator response
         response = generator.run(parts=[prompt])
-        answer = response.get("replies", [""])[0] if response and response.get("replies") else "No response generated"
+        answer = response.get("replies", [""])[0] if response and response.get("replies") else "No response"
         
-        # Format sources
         sources = [
             {
                 "filename": d.meta.get("filename", "Unknown"),
@@ -195,7 +141,4 @@ def query_rag(question: str, session_id: str) -> dict:
     
     except Exception as e:
         logger.exception(f"Query failed: {e}")
-        return {
-            "answer": "Sorry, I encountered an error processing your request.",
-            "sources": []
-        }
+        return {"answer": "Sorry, I encountered an error.", "sources": []}