new embeeding model + new create_quer_engine with keyword matching

index_retriever.py

@@ -27,38 +27,89 @@ def create_vector_index(documents):
     index = VectorStoreIndex.from_documents(documents)
     log_message("✓ Index created")
     return index
+from llama_index.core.vector_stores import MetadataFilters, ExactMatchFilter
+import re
+
+def extract_document_id(query):
+    """Extract GOST document ID from query"""
+    patterns = [
+        r'ГОСТ\s*Р?\s*([\d\.]+(?:-\d{4})?)',
+        r'НП-[\d\-]+',
+        r'ПН\s+АЭ\s+Г-[\d\-]+'
+    ]
+    
+    for pattern in patterns:
+        match = re.search(pattern, query, re.IGNORECASE)
+        if match:
+            doc_id = match.group(0)
+            # Normalize
+            doc_id = re.sub(r'ГОСТ\s*Р', 'ГОСТ Р', doc_id, flags=re.IGNORECASE)
+            if 'ГОСТ' in doc_id and '-' not in doc_id:
+                doc_id += '-2020'
+            return doc_id
+    return None
+
 
 def create_query_engine(vector_index):
-    """Create hybrid retrieval engine"""
+    """Create hybrid retrieval engine with document filtering"""
     log_message("Creating query engine...")
     
-    # Vector retriever
-    vector_retriever = VectorIndexRetriever(
-        index=vector_index,
-        similarity_top_k=50
-    )
-    
-    # BM25 retriever
-    bm25_retriever = BM25Retriever.from_defaults(
-        docstore=vector_index.docstore,
-        similarity_top_k=50
-    )
+    def retrieve_with_filter(query_str):
+        """Custom retrieval with optional document filtering"""
+        doc_id = extract_document_id(query_str)
+        
+        if doc_id:
+            log_message(f"Detected document filter: {doc_id}")
+            
+            # Try filtered retrieval first
+            filters = MetadataFilters(
+                filters=[ExactMatchFilter(key="document_id", value=doc_id)]
+            )
+            
+            filtered_retriever = VectorIndexRetriever(
+                index=vector_index,
+                similarity_top_k=30,
+                filters=filters
+            )
+            
+            filtered_results = filtered_retriever.retrieve(query_str)
+            log_message(f"Filtered retrieval: {len(filtered_results)} results from {doc_id}")
+            
+            if len(filtered_results) >= 10:
+                # Good enough, use filtered results
+                return filtered_results
+            else:
+                log_message("Not enough filtered results, falling back to hybrid")
+        
+        # Fallback to hybrid retrieval
+        vector_retriever = VectorIndexRetriever(
+            index=vector_index,
+            similarity_top_k=50
+        )
+        
+        bm25_retriever = BM25Retriever.from_defaults(
+            docstore=vector_index.docstore,
+            similarity_top_k=50
+        )
+        
+        hybrid_retriever = QueryFusionRetriever(
+            [vector_retriever, bm25_retriever],
+            similarity_top_k=60,
+            num_queries=1
+        )
+        
+        return hybrid_retriever.retrieve(query_str)
     
-    # Hybrid fusion
-    hybrid_retriever = QueryFusionRetriever(
-        [vector_retriever, bm25_retriever],
-        similarity_top_k=60,
-        num_queries=1
-    )
+    # Create custom query engine
+    class CustomRetriever:
+        def retrieve(self, query_str):
+            return retrieve_with_filter(query_str)
     
-    # Response synthesizer
     response_synthesizer = get_response_synthesizer()
-    
-    # Query engine
     query_engine = RetrieverQueryEngine(
-        retriever=hybrid_retriever,
+        retriever=CustomRetriever(),
         response_synthesizer=response_synthesizer
     )
     
-    log_message("✓ Query engine created")
+    log_message("✓ Query engine created with document filtering")
     return query_engine

utils.py

@@ -7,9 +7,12 @@ def get_llm_model(api_key, model_name="gemini-2.0-flash"):
     """Get LLM model"""
     return GoogleGenAI(model=model_name, api_key=api_key)
 
-def get_embedding_model(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"):
-    """Get embedding model"""
-    return HuggingFaceEmbedding(model_name=model_name)
+def get_embedding_model(model_name="intfloat/multilingual-e5-large"):
+    """Use better multilingual embedding model"""
+    return HuggingFaceEmbedding(
+        model_name=model_name,
+        trust_remote_code=True
+    )
 
 def get_reranker_model(model_name='cross-encoder/ms-marco-MiniLM-L-12-v2'):
     """Get reranker model"""