Spaces:
Running
Running
File size: 7,097 Bytes
09801ca | 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 | import os
import uuid
from typing import List, Dict, Any, Optional
from datetime import datetime
import logging
os.environ["USE_TF"] = "0"
os.environ["USE_TORCH"] = "1"
os.environ["TF_USE_LEGACY_KERAS"] = "1"
try:
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
from sentence_transformers import SentenceTransformer
QDRANT_AVAILABLE = True
except (ImportError, Exception):
QDRANT_AVAILABLE = False
logger = logging.getLogger(__name__)
class VectorStoreService:
_instance = None
def __new__(cls, *args, **kwargs):
if cls._instance is None:
cls._instance = super(VectorStoreService, cls).__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if getattr(self, '_initialized', False):
return
self.is_ready = QDRANT_AVAILABLE
if not self.is_ready:
logger.warning("Qdrant or SentenceTransformers not installed. Vector store disabled.")
self._initialized = True
return
# Initialize local Qdrant
# This creates a 'qdrant_data' directory locally (no Docker required!)
qdrant_path = os.path.join(os.getcwd(), "qdrant_data")
os.makedirs(qdrant_path, exist_ok=True)
import time
for _ in range(5):
try:
self.client = QdrantClient(path=qdrant_path)
break
except Exception as e:
logger.warning(f"Qdrant locked, waiting 1s... ({e})")
time.sleep(1)
else:
logger.error("Failed to acquire Qdrant lock after 5 seconds. Falling back to memory mode.")
self.client = QdrantClient(":memory:")
# Initialize lightweight local embedding model
# all-MiniLM-L6-v2 is extremely fast and great for general semantic search
logger.info("Loading local embedding model (SentenceTransformers)...")
self.model = SentenceTransformer("all-MiniLM-L6-v2")
self.vector_size = self.model.get_sentence_embedding_dimension()
# Setup collections
self.chat_collection = "chat_memory"
self.doc_collection = "document_chunks"
self._ensure_collections()
self._initialized = True
def _ensure_collections(self):
"""Ensure required collections exist in Qdrant"""
collections = [c.name for c in self.client.get_collections().collections]
for collection_name in [self.chat_collection, self.doc_collection]:
if collection_name not in collections:
logger.info(f"Creating Qdrant collection: {collection_name}")
self.client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=self.vector_size, distance=Distance.COSINE),
)
def add_chat_message(self, user_id: str, role: str, content: str, conversation_id: str) -> bool:
if not self.is_ready: return False
try:
# Embed the message content
vector = self.model.encode(content).tolist()
# Store in Qdrant
point_id = str(uuid.uuid4())
self.client.upsert(
collection_name=self.chat_collection,
points=[
PointStruct(
id=point_id,
vector=vector,
payload={
"user_id": user_id,
"conversation_id": conversation_id,
"role": role,
"content": content,
"timestamp": datetime.utcnow().isoformat()
}
)
]
)
return True
except Exception as e:
logger.error(f"Failed to add chat to vector store: {e}")
return False
def search_chat_history(self, user_id: str, query: str, limit: int = 5) -> List[Dict]:
"""Semantically search the user's past chat history (RAG over memory)"""
if not self.is_ready: return []
try:
query_vector = self.model.encode(query).tolist()
# Filter by user_id
user_filter = Filter(
must=[
FieldCondition(
key="user_id",
match=MatchValue(value=user_id)
)
]
)
search_result = self.client.search(
collection_name=self.chat_collection,
query_vector=query_vector,
query_filter=user_filter,
limit=limit
)
return [hit.payload for hit in search_result]
except Exception as e:
logger.error(f"Failed to search chat history: {e}")
return []
def connect_custom_qdrant(self, url: str, api_key: Optional[str] = None, collection_name: str = "dataset_metadata") -> Dict[str, Any]:
"""Connect to Qdrant Cloud or custom Qdrant instance with robust URL normalization and failover."""
if not QDRANT_AVAILABLE:
raise Exception("qdrant_client package is not installed.")
url_clean = url.strip().rstrip('/')
api_key_clean = api_key.strip() if api_key else None
candidate_urls = []
if not url_clean.startswith('http://') and not url_clean.startswith('https://'):
candidate_urls.append(f"https://{url_clean}")
candidate_urls.append(f"http://{url_clean}")
else:
candidate_urls.append(url_clean)
if ':6333' in url_clean:
candidate_urls.append(url_clean.replace(':6333', ''))
else:
candidate_urls.append(f"{url_clean}:6333")
last_error = None
new_client = None
successful_url = ""
for cand in candidate_urls:
try:
test_c = QdrantClient(url=cand, api_key=api_key_clean, timeout=8)
cols = test_c.get_collections().collections
new_client = test_c
successful_url = cand
break
except Exception as ex:
last_error = ex
if not new_client:
raise Exception(f"Failed to connect to Qdrant at {url}: {str(last_error)}")
self.client = new_client
self.is_ready = True
if collection_name:
self.doc_collection = collection_name
self._ensure_collections()
cols = self.client.get_collections().collections
return {
"status": "connected",
"active_url": successful_url,
"collections_count": len(cols),
"collections": [c.name for c in cols]
}
# Singleton instance
vector_store = VectorStoreService()
|