Tell-Me-Mo
diff --git a/‎backend/alembic/versions/6c7942ee0af2_add_context_id_to_conversations.py‎
Lines changed: 41 additions & 0 deletions b/‎backend/alembic/versions/6c7942ee0af2_add_context_id_to_conversations.py‎
Lines changed: 41 additions & 0 deletions
diff --git a/‎backend/models/conversation.py‎
Lines changed: 1 addition & 0 deletions b/‎backend/models/conversation.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎backend/routers/conversations.py‎
Lines changed: 24 additions & 23 deletions b/‎backend/routers/conversations.py‎
Lines changed: 24 additions & 23 deletions
diff --git a/‎backend/services/rag/hybrid_search.py‎
Lines changed: 37 additions & 52 deletions b/‎backend/services/rag/hybrid_search.py‎
Lines changed: 37 additions & 52 deletions
diff --git a/‎backend/tests/integration/test_rag_pipeline.py‎
Lines changed: 97 additions & 0 deletions b/‎backend/tests/integration/test_rag_pipeline.py‎
Lines changed: 97 additions & 0 deletions
@@ -0,0 +1,41 @@
+"""add_context_id_to_conversations
+
+Revision ID: 6c7942ee0af2
+Revises: convert_enum_to_varchar
+Create Date: 2025-10-16 08:36:51.719536
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '6c7942ee0af2'
+down_revision: Union[str, Sequence[str], None] = 'convert_enum_to_varchar'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # Add context_id column to conversations table
+    op.add_column('conversations', sa.Column('context_id', sa.String(length=255), nullable=True))
+
+    # Add index for faster filtering
+    op.create_index(
+        'ix_conversations_context_id',
+        'conversations',
+        ['context_id'],
+        unique=False
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # Drop index first
+    op.drop_index('ix_conversations_context_id', table_name='conversations')
+
+    # Drop column
+    op.drop_column('conversations', 'context_id')
@@ -17,6 +17,7 @@ class Conversation(Base):
     id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     project_id = Column(UUID(as_uuid=True), nullable=True)  # Can store project_id, program_id, portfolio_id, or None for org-level. No FK constraint to allow flexibility.
     organization_id = Column(UUID(as_uuid=True), ForeignKey("organizations.id", ondelete="CASCADE"), nullable=False)
+    context_id = Column(String(255), nullable=True)  # Optional context ID for filtering (e.g., 'risk_<uuid>', 'task_<uuid>')
 
     # Conversation metadata
     title = Column(String(255), nullable=False)
 
@@ -32,11 +32,13 @@ class MessageModel(BaseModel):
 class ConversationCreateRequest(BaseModel):
     title: str
     messages: List[MessageModel] = []
+    context_id: Optional[str] = None  # Optional context ID (e.g., 'risk_<uuid>')
 
 
 class ConversationUpdateRequest(BaseModel):
     title: Optional[str] = None
     messages: Optional[List[MessageModel]] = None
+    context_id: Optional[str] = None
 
 
 class ConversationResponse(BaseModel):
@@ -51,6 +53,7 @@ class ConversationResponse(BaseModel):
 @router.get("/{project_id}/conversations", response_model=List[ConversationResponse])
 async def get_conversations(
     project_id: str,
+    context_id: Optional[str] = None,
     session: AsyncSession = Depends(get_db),
     current_org: Organization = Depends(get_current_organization),
     current_user: User = Depends(get_current_user),
@@ -61,39 +64,33 @@ async def get_conversations(
 
     Args:
         project_id: UUID of the entity or 'organization' for org-level conversations
+        context_id: Optional context ID to filter conversations (e.g., 'risk_<uuid>')
         session: Database session
 
     Returns:
         List of conversations for the entity
     """
-    logger.info(f"Getting conversations for entity {sanitize_for_log(project_id)}")
+    logger.info(f"Getting conversations for entity {sanitize_for_log(project_id)}, context_id: {sanitize_for_log(context_id)}")
 
     try:
+        # Build WHERE clause conditions
+        conditions = [Conversation.organization_id == current_org.id]
+
         # Handle organization-level conversations
         if project_id == 'organization':
-            result = await session.execute(
-                select(Conversation)
-                .where(
-                    and_(
-                        Conversation.project_id.is_(None),  # Organization-level conversations have NULL project_id
-                        Conversation.organization_id == current_org.id
-                    )
-                )
-                .order_by(desc(Conversation.last_accessed_at))
-            )
+            conditions.append(Conversation.project_id.is_(None))
         else:
-            # Get conversations directly by entity ID (no entity type validation needed)
-            # The entity ID could be a project, program, or portfolio
-            result = await session.execute(
-                select(Conversation)
-                .where(
-                    and_(
-                        Conversation.project_id == project_id,
-                        Conversation.organization_id == current_org.id
-                    )
-                )
-                .order_by(desc(Conversation.last_accessed_at))
-            )
+            conditions.append(Conversation.project_id == project_id)
+
+        # Filter by context_id if provided
+        if context_id:
+            conditions.append(Conversation.context_id == context_id)
+
+        result = await session.execute(
+            select(Conversation)
+            .where(and_(*conditions))
+            .order_by(desc(Conversation.last_accessed_at))
+        )
         conversations = result.scalars().all()
 
         return [
@@ -142,6 +139,7 @@ async def create_conversation(
         conversation = Conversation(
             project_id=None if project_id == 'organization' else uuid.UUID(project_id),
             organization_id=current_org.id,
+            context_id=request.context_id,  # Store context_id
             title=request.title,
             messages=[msg.model_dump() for msg in request.messages],
             created_by=current_user.email or "unknown",
@@ -231,6 +229,9 @@ async def update_conversation(
         if request.messages is not None:
             update_data["messages"] = [msg.model_dump() for msg in request.messages]
 
+        if request.context_id is not None:
+            update_data["context_id"] = request.context_id
+
         # Update conversation
         await session.execute(
             update(Conversation)
 
@@ -78,6 +78,9 @@ class HybridSearchConfig:
     final_result_count: int = 15
     diversity_boost: float = 0.1
 
+    # Diversity optimization
+    diversity_similarity_threshold: float = 0.85  # Results with similarity > this are considered duplicates
+
     # Quality filters
     min_confidence_score: float = 0.3
     filter_low_quality: bool = True
@@ -225,6 +228,7 @@ async def hybrid_search(
 
             # Calculate pipeline metrics
             processing_time = int((time.time() - start_time) * 1000)
+            diversity_score = await self._calculate_diversity_score(final_results)
             pipeline = SearchPipeline(
                 query=query,
                 config=self.config,
@@ -236,7 +240,7 @@ async def hybrid_search(
                 semantic_result_count=len(semantic_results),
                 keyword_result_count=len(keyword_results),
                 overlap_count=self._calculate_overlap(semantic_results, keyword_results),
-                diversity_score=self._calculate_diversity_score(final_results),
+                diversity_score=diversity_score,
                 confidence_distribution=self._calculate_confidence_distribution(final_results),
                 processing_time_ms=processing_time
             )
@@ -786,62 +790,42 @@ async def _optimize_final_results(
         return final_results
 
     async def _diversify_results(
-        self, 
-        results: List[SearchResult], 
+        self,
+        results: List[SearchResult],
         query: str
     ) -> List[SearchResult]:
         """Apply diversity optimization to avoid redundant results."""
         if not results or len(results) <= 1:
             return results
-        
+
         # Simple diversity based on content similarity
         if self.sentence_transformer:
             try:
                 # Get embeddings for result texts
                 texts = [r.text[:200] for r in results]  # Limit length
-                embeddings = self.sentence_transformer.encode(texts)
-                
-                # Select diverse results using MMR-like approach
-                diverse_results = []
-                remaining_indices = set(range(len(results)))
-                
-                # Start with highest scoring result
-                best_idx = 0
-                diverse_results.append(results[best_idx])
-                remaining_indices.remove(best_idx)
-                
-                # Select remaining results balancing relevance and diversity
-                while remaining_indices and len(diverse_results) < len(results):
-                    best_score = -1
-                    best_idx = -1
-                    
-                    for idx in remaining_indices:
-                        # Relevance score
-                        relevance = results[idx].final_score
-                        
-                        # Diversity score (minimum similarity to selected results)
-                        similarities = []
-                        for selected_result in diverse_results:
-                            selected_idx = results.index(selected_result)
-                            sim = self._cosine_similarity(
-                                embeddings[idx], 
-                                embeddings[selected_idx]
-                            )
-                            similarities.append(sim)
-                        
-                        diversity = 1.0 - max(similarities) if similarities else 1.0
-                        
-                        # Combined score (balance relevance and diversity)
-                        combined_score = 0.7 * relevance + 0.3 * diversity
-                        
-                        if combined_score > best_score:
-                            best_score = combined_score
-                            best_idx = idx
-                    
-                    if best_idx != -1:
-                        diverse_results.append(results[best_idx])
-                        remaining_indices.remove(best_idx)
-                
+                # Run blocking sentence transformer in thread pool to avoid blocking event loop
+                embeddings = await asyncio.to_thread(self.sentence_transformer.encode, texts)
+
+                # SIMPLIFIED: Just filter out highly similar consecutive results
+                # This is much faster than full MMR and good enough for our use case
+                diverse_results = [results[0]]  # Always keep the best result
+                selected_indices = [0]  # Track indices of selected results
+
+                for i in range(1, len(results)):
+                    # Check similarity to all selected results
+                    is_diverse = True
+                    for selected_idx in selected_indices:
+                        sim = self._cosine_similarity(embeddings[i], embeddings[selected_idx])
+
+                        # If too similar to any selected result, skip it
+                        if sim > self.config.diversity_similarity_threshold:
+                            is_diverse = False
+                            break
+
+                    if is_diverse:
+                        diverse_results.append(results[i])
+                        selected_indices.append(i)
+
                 return diverse_results
 
             except Exception as e:
@@ -862,24 +846,25 @@ def _calculate_overlap(
         keyword_ids = set(r.chunk_id for r in keyword_results)
         return len(semantic_ids.intersection(keyword_ids))
 
-    def _calculate_diversity_score(self, results: List[SearchResult]) -> float:
+    async def _calculate_diversity_score(self, results: List[SearchResult]) -> float:
         """Calculate diversity score for result set."""
         if len(results) <= 1:
             return 0.0
-        
+
         # Diversity based on search method coverage
         search_types = set()
         for result in results:
             search_types.update(result.search_types)
-        
+
         type_diversity = len(search_types) / 4.0  # Max 4 search types
-        
+
         # Content diversity (if embeddings available)
         content_diversity = 0.5  # Default
         if self.sentence_transformer and len(results) > 1:
             try:
                 texts = [r.text[:100] for r in results[:10]]  # Sample for performance
-                embeddings = self.sentence_transformer.encode(texts)
+                # Run blocking sentence transformer in thread pool
+                embeddings = await asyncio.to_thread(self.sentence_transformer.encode, texts)
 
                 similarities = []
                 for i in range(len(embeddings)):
 
@@ -1056,3 +1056,100 @@ async def test_delete_vectors_by_project_id(
 
         assert len(results_proj1) == 0  # Project 1 deleted
         assert len(results_proj2) >= 1  # Project 2 still exists
+
+
+# ============================================================================
+# Section 6.2: RAG Pipeline - Diversity Optimization (Async Threading)
+# ============================================================================
+
+class TestDiversityOptimization:
+    """Test that diversity optimization uses async threading to avoid blocking."""
+
+    @pytest.mark.asyncio
+    async def test_diversify_results_with_async_threading(self, mock_embedding_model):
+        """
+        Test that _diversify_results uses asyncio.to_thread() for sentence transformer
+        encoding to avoid blocking the async event loop.
+
+        This test covers the fix for the hanging issue where synchronous
+        sentence_transformer.encode() calls would block the entire async event loop.
+        """
+        from services.rag.hybrid_search import HybridSearchService, SearchResult, SearchType
+        import asyncio
+
+        # Arrange - Create service with mock model
+        service = HybridSearchService()
+        service.sentence_transformer = mock_embedding_model
+
+        # Create 26 mock search results (same as real scenario that was hanging)
+        results = [
+            SearchResult(
+                chunk_id=f"chunk_{i}",
+                text=f"Meeting notes about security and risk assessment session {i}",
+                metadata={"title": f"Meeting {i}"},
+                semantic_score=0.8,
+                keyword_score=0.6,
+                hybrid_score=0.7,
+                final_score=0.75 - (i * 0.01),  # Decreasing scores
+                search_types=[SearchType.SEMANTIC, SearchType.KEYWORD],
+                confidence_score=0.7
+            )
+            for i in range(26)
+        ]
+
+        # Act - Run diversity optimization with timeout to detect hanging
+        try:
+            diverse_results = await asyncio.wait_for(
+                service._diversify_results(results, "test query"),
+                timeout=10.0  # Should complete in <1 second, 10s is generous
+            )
+
+            # Assert - Should complete without timing out
+            assert diverse_results is not None
+            assert len(diverse_results) > 0
+            assert len(diverse_results) <= len(results)  # May filter some
+
+        except asyncio.TimeoutError:
+            pytest.fail(
+                "Diversity optimization timed out! The asyncio.to_thread() fix "
+                "is not working - sentence_transformer.encode() is blocking the event loop."
+            )
+
+    @pytest.mark.asyncio
+    async def test_calculate_diversity_score_with_async_threading(self, mock_embedding_model):
+        """
+        Test that _calculate_diversity_score uses async threading for embedding generation.
+        """
+        from services.rag.hybrid_search import HybridSearchService, SearchResult, SearchType
+        import asyncio
+
+        # Arrange
+        service = HybridSearchService()
+        service.sentence_transformer = mock_embedding_model
+
+        results = [
+            SearchResult(
+                chunk_id=f"chunk_{i}",
+                text=f"Content {i}",
+                metadata={},
+                semantic_score=0.8,
+                final_score=0.8,
+                search_types=[SearchType.SEMANTIC],
+                confidence_score=0.8
+            )
+            for i in range(10)
+        ]
+
+        # Act - Should complete without hanging
+        try:
+            diversity_score = await asyncio.wait_for(
+                service._calculate_diversity_score(results),
+                timeout=5.0
+            )
+
+            # Assert
+            assert isinstance(diversity_score, float)
+            assert 0 <= diversity_score <= 1.0
+
+        except asyncio.TimeoutError:
+            pytest.fail("Diversity score calculation timed out - async threading issue!")