Spaces:
Paused
Paused
Update rag_engine.py
Browse files- rag_engine.py +21 -6
rag_engine.py
CHANGED
|
@@ -52,12 +52,26 @@ class RAGEngine:
|
|
| 52 |
print(f"✅ Collection '{self.collection_name}' created!")
|
| 53 |
self._build_index()
|
| 54 |
else:
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 61 |
except Exception as e:
|
| 62 |
print(f"⚠️ Collection check/creation failed: {e}")
|
| 63 |
# Try to build anyway, maybe wrapper handles it
|
|
@@ -183,6 +197,7 @@ class RAGEngine:
|
|
| 183 |
try:
|
| 184 |
self.client.delete_collection(self.collection_name)
|
| 185 |
self._build_index()
|
|
|
|
| 186 |
return "✅ Knowledge Base Refreshed on Cloud!"
|
| 187 |
except Exception as e:
|
| 188 |
return f"❌ Error refreshing: {e}"
|
|
|
|
| 52 |
print(f"✅ Collection '{self.collection_name}' created!")
|
| 53 |
self._build_index()
|
| 54 |
else:
|
| 55 |
+
# Check if dataset is already indexed
|
| 56 |
+
dataset_filter = models.Filter(
|
| 57 |
+
must=[
|
| 58 |
+
models.FieldCondition(
|
| 59 |
+
key="metadata.source",
|
| 60 |
+
match=models.MatchValue(value="hf_dataset")
|
| 61 |
+
)
|
| 62 |
+
]
|
| 63 |
+
)
|
| 64 |
+
dataset_count = self.client.count(
|
| 65 |
+
collection_name=self.collection_name,
|
| 66 |
+
count_filter=dataset_filter
|
| 67 |
+
).count
|
| 68 |
+
|
| 69 |
+
print(f"✅ Qdrant Collection '{self.collection_name}' ready with {count} vectors.")
|
| 70 |
+
|
| 71 |
+
if dataset_count == 0:
|
| 72 |
+
print("⚠️ Phishing dataset not found in collection. Loading...")
|
| 73 |
+
self.load_from_huggingface()
|
| 74 |
+
|
| 75 |
except Exception as e:
|
| 76 |
print(f"⚠️ Collection check/creation failed: {e}")
|
| 77 |
# Try to build anyway, maybe wrapper handles it
|
|
|
|
| 197 |
try:
|
| 198 |
self.client.delete_collection(self.collection_name)
|
| 199 |
self._build_index()
|
| 200 |
+
self.load_from_huggingface()
|
| 201 |
return "✅ Knowledge Base Refreshed on Cloud!"
|
| 202 |
except Exception as e:
|
| 203 |
return f"❌ Error refreshing: {e}"
|