File size: 11,897 Bytes
b30f068 d94ff21 b30f068 d94ff21 b30f068 d94ff21 b30f068 d94ff21 b30f068 d94ff21 b30f068 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 | """
Qdrant Vector Store
Manages vector embeddings in Qdrant database
"""
import os
import sys
import threading
from pathlib import Path
# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from typing import List, Dict
from src.config.paths import QDRANT_DIR
class QdrantVectorStore:
"""
Qdrant Vector Store Manager
Handles storage and retrieval of document embeddings
Usage:
store = QdrantVectorStore()
store.add_document_chunk("chunk_id", embedding, metadata)
results = store.search_documents(query_embedding)
"""
def __init__(self, host: str = "localhost", port: int = 6333):
"""
Initialize Qdrant client
Args:
host: Qdrant host (default: localhost)
port: Qdrant port (default: 6333)
"""
self.host = host
self.port = port
self.embedding_dim = 1536 # OpenAI text-embedding-3-small
self.using_docker = False
# Serializes access to the client. The embedded (on-disk) Qdrant client is
# not safe under concurrent threads, so search/upsert take this lock. Paired
# with the process-wide singleton (get_shared_vector_store) this makes the
# multi-user Space safe: one client, one-at-a-time access.
self._lock = threading.RLock()
# QDRANT_LOCAL_PATH env still overrides; default is the project-root
# anchored QDRANT_DIR (build scratch). Empty string forces in-memory.
local_path_env = os.environ.get("QDRANT_LOCAL_PATH")
local_path = local_path_env if local_path_env is not None else str(QDRANT_DIR)
# Build mode: write directly to the on-disk store and skip the Docker
# probe (used by scripts/build_index.py). Avoids accidentally connecting
# to a stray empty Docker on :6333 during a build.
if os.environ.get("QDRANT_FORCE_LOCAL", "0") == "1":
self.client = QdrantClient(path=local_path)
print(f"📦 QDRANT_FORCE_LOCAL: using on-disk store at '{local_path}' (build mode)")
self.initialize_collections()
return
# Served mode: try Docker Qdrant first (the deploy restores the committed
# snapshot into it and both processes read from it).
try:
self.client = QdrantClient(host=host, port=port, timeout=5)
# Test connection by getting collections
_ = self.client.get_collections()
self.using_docker = True
print(f"✅ Connected to Qdrant Docker at {host}:{port}")
self.initialize_collections()
return
except Exception as e:
docker_error = e
# Fail fast in the hosted deploy: a served instance MUST read the Docker
# Qdrant restored from the committed snapshot. Silently falling back to an
# empty in-memory (or fresh on-disk) store would answer nothing or trigger
# a full re-embed. Opt in with QDRANT_REQUIRE_DOCKER=1 in the deploy env.
if os.environ.get("QDRANT_REQUIRE_DOCKER", "0") == "1":
raise RuntimeError(
f"QDRANT_REQUIRE_DOCKER=1 but no Qdrant reachable at {host}:{port}. "
f"Start Docker Qdrant and restore the committed snapshot "
f"(data/qdrant_snapshots/, see RUNBOOK). Original error: {docker_error}"
)
# Legacy / dev fallback. Persistent on-disk store so embeddings survive
# restarts (fixes the cold-start re-embed ISA reported); in-memory only
# if the path is unavailable or explicitly disabled (QDRANT_LOCAL_PATH="").
print(f"⚠️ Qdrant Docker not available at {host}:{port}")
print(f" Error: {str(docker_error)}")
print(" 💡 For production, start Docker: docker run -d -p 6333:6333 qdrant/qdrant")
if local_path:
try:
self.client = QdrantClient(path=local_path)
print(f" 📦 Using persistent on-disk store at '{local_path}' (survives restarts)")
except Exception as e2:
print(f" ⚠️ On-disk store unavailable ({e2}); using in-memory (data lost on restart)")
self.client = QdrantClient(":memory:")
else:
print(" 📦 Using in-memory mode (data lost on restart)")
self.client = QdrantClient(":memory:")
self.initialize_collections()
def initialize_collections(self):
"""Create collections for APIs and PDF documents"""
# Collection for API descriptions (for future use)
if not self.client.collection_exists("api_catalog"):
try:
self.client.create_collection(
collection_name="api_catalog",
vectors_config=VectorParams(
size=self.embedding_dim,
distance=Distance.COSINE
)
)
except Exception:
pass # Collection might already exist
# Collection for PDF document chunks (CDMS)
if not self.client.collection_exists("cdms_documents"):
try:
self.client.create_collection(
collection_name="cdms_documents",
vectors_config=VectorParams(
size=self.embedding_dim,
distance=Distance.COSINE
)
)
except Exception:
pass # Collection might already exist
def add_document_chunk(
self,
chunk_id: str,
embedding: List[float],
metadata: Dict
):
"""
Add PDF chunk to vector store
Args:
chunk_id: Unique identifier for the chunk (string, will be converted to int)
embedding: Vector embedding (1536 dimensions)
metadata: Dict with chunk metadata
"""
try:
# Convert string ID to integer hash for Qdrant (Qdrant requires int or UUID)
import hashlib
int_id = int(hashlib.md5(chunk_id.encode()).hexdigest()[:15], 16) # Use first 15 hex chars as int
with self._lock:
self.client.upsert(
collection_name="cdms_documents",
points=[PointStruct(
id=int_id, # Qdrant requires integer or UUID
vector=embedding,
payload=metadata
)]
)
except Exception as e:
print(f"⚠️ Warning: Could not add chunk to Qdrant: {e}")
def search_documents(
self,
query_embedding: List[float],
limit: int = 10,
score_threshold: float = 0.4, # Balanced: 0.3 let weak/wrong-product
# matches through (they were being answered from); 0.5 was too strict.
# Tune with real queries. Better to abstain than cite the wrong label.
query_filter=None,
) -> List[Dict]:
"""
Search for similar document chunks
Args:
query_embedding: Query vector embedding
limit: Maximum number of results
score_threshold: Minimum similarity score (0-1)
query_filter: Optional Qdrant Filter to scope the search (e.g. by
product) so retrieval is constrained at the vector level rather
than post-filtered from a global, dominant-product-heavy top-k.
Returns:
List of search results with metadata and scores
"""
try:
with self._lock:
results = self.client.search(
collection_name="cdms_documents",
query_vector=query_embedding,
query_filter=query_filter,
limit=limit,
score_threshold=score_threshold,
with_payload=True
)
# Format results
formatted_results = []
for result in results:
payload = result.payload
formatted_results.append({
"id": result.id,
"score": result.score,
"content": payload.get("content", ""),
"source_file": payload.get("source_file", "Unknown"),
"page_number": payload.get("page_number", 0),
"document_id": payload.get("document_id", ""),
"pdf_url": payload.get("pdf_url", ""), # PHASE 1 FIX: Extract PDF URL from payload
"url_hash": payload.get("url_hash", ""), # PHASE 1 FIX: Extract URL hash from payload
"metadata": payload
})
return formatted_results
except Exception as e:
print(f"⚠️ Warning: Could not search Qdrant: {e}")
return []
def get_collection_info(self) -> Dict:
"""Get information about collections"""
try:
info = {
"using_docker": self.using_docker,
"host": self.host if self.using_docker else "in-memory",
"port": self.port if self.using_docker else None
}
if self.client.collection_exists("cdms_documents"):
collection_info = self.client.get_collection("cdms_documents")
info["cdms_documents"] = {
"points_count": collection_info.points_count,
# vectors_count was removed in qdrant-client >=1.16; fall back
# to points_count so this stays version-tolerant.
"vectors_count": getattr(collection_info, "vectors_count", None)
or collection_info.points_count
}
return info
except Exception as e:
return {"error": str(e)}
# --- Process-wide singleton --------------------------------------------------
# The embedded (on-disk) Qdrant permits only ONE client per path per process: a
# second QdrantClient(path=...) raises "already accessed by another instance".
# Every request previously built its own store, so two concurrent users on the
# Space would collide -> the loser fell back to an empty in-memory store and
# abstained. Share ONE store across all requests (double-checked lock for a safe
# first init); per-op locking inside the store serialises concurrent access.
_SHARED_VS = None
_SHARED_VS_LOCK = threading.Lock()
def get_shared_vector_store() -> "QdrantVectorStore":
"""Return the process-wide QdrantVectorStore, creating it once on first use."""
global _SHARED_VS
if _SHARED_VS is None:
with _SHARED_VS_LOCK:
if _SHARED_VS is None:
_SHARED_VS = QdrantVectorStore()
return _SHARED_VS
# Test function
if __name__ == "__main__":
print("Testing Qdrant Vector Store...")
print("-" * 70)
try:
store = QdrantVectorStore()
print("✅ Qdrant connection successful!")
info = store.get_collection_info()
if info:
print("\n📊 Collection Info:")
for collection, data in info.items():
print(f" {collection}: {data.get('points_count', 0)} points")
else:
print("\n📊 Collections initialized (empty)")
print("\n💡 To use:")
print(" 1. Make sure Qdrant is running:")
print(" docker run -d -p 6333:6333 qdrant/qdrant")
print(" 2. Or it will use in-memory mode automatically")
except Exception as e:
print(f"❌ Error: {e}")
|