Update retriever.py

retriever.py

@@ -7,8 +7,19 @@ def get_relevant_passages(query, df, top_k=20):
     # Create a copy to avoid modifying the original dataframe
     df_copy = df.copy()
     
+    # Print shape for debugging
+    print(f"DataFrame shape: {df_copy.shape}")
+    print(f"DataFrame columns: {df_copy.columns.tolist()}")
+    
+    # Handle missing columns gracefully
+    for col in ['description', 'test_type', 'adaptive_support', 'remote_support', 'duration']:
+        if col not in df_copy.columns:
+            df_copy[col] = 'N/A'
+    
     # Ensure URL field is properly formatted
-    if 'url' in df_copy.columns:
+    if 'url' not in df_copy.columns:
+        df_copy['url'] = 'https://www.shl.com/missing-url'
+    else:
         # Clean up URLs if needed
         df_copy['url'] = df_copy['url'].astype(str)
         # Ensure URLs start with http or https
@@ -18,16 +29,16 @@ def get_relevant_passages(query, df, top_k=20):
     # Format test_type for better representation
     def format_test_type(test_types):
         if isinstance(test_types, list):
-            return ', '.join(test_types)
+            return ', '.join([str(t) for t in test_types if t])
         return str(test_types)
     
     # Concatenate all fields into a single string per row
     corpus = df_copy.apply(
-        lambda row: f"{row['description']} "
+        lambda row: f"{row.get('assessment_name', '')} {row.get('description', '')} "
                    f"Test types: {format_test_type(row['test_type'])}. "
                    f"Adaptive support: {row['adaptive_support']}. "
                    f"Remote support: {row['remote_support']}. "
-                   f"Duration: {row['duration'] if pd.notna(row['duration']) else 'N/A'} minutes.",
+                   f"Duration: {row['duration']} minutes.",
         axis=1
     ).tolist()
     
@@ -35,4 +46,4 @@ def get_relevant_passages(query, df, top_k=20):
     query_embedding = model.encode(query, convert_to_tensor=True)
     hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=top_k)[0]
     
-    return df_copy.iloc[[hit['corpus_id'] for hit in hits]]
+    return df_copy.iloc[[hit['corpus_id'] for hit in hits]]