File size: 7,617 Bytes
5a21d6e
a5aba38
 
 
5a21d6e
a5aba38
 
 
 
 
 
 
 
79d4fd5
a5aba38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79d4fd5
 
a5aba38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5a21d6e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a5aba38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79d4fd5
a5aba38
 
 
 
 
 
 
 
 
 
 
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
from typing import List, Dict, Any, Optional, Set
from langchain_qdrant import QdrantVectorStore
from langchain_core.documents import Document
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams, Filter, FieldCondition, MatchValue

from config import (
    QDRANT_PATH,
    COLLECTION_NAME,
    EMBEDDING_DIMENSION,
    USE_MEMORY_MODE
)
from embeddings import get_embedder
from qdrant_client_manager import get_qdrant_client


class VectorStore:
    """Qdrant vector store wrapper."""
    
    def __init__(
        self,
        path: str = QDRANT_PATH,
        collection_name: str = COLLECTION_NAME,
        use_memory: bool = USE_MEMORY_MODE,
        embedder=None
    ):
        self.collection_name = collection_name
        self.use_memory = use_memory
        self.path = path
        
        # Use shared Qdrant client to prevent multiple instance conflicts
        self._client = get_qdrant_client()
        
        self._ensure_collection_exists()
        self._embedder = embedder or get_embedder()
        
        self._vector_store = QdrantVectorStore(
            client=self._client,
            collection_name=self.collection_name,
            embedding=self._embedder
        )
    
    def _ensure_collection_exists(self):
        """Create collection if it doesn't exist."""
        collections = self._client.get_collections().collections
        names = [c.name for c in collections]
        
        if self.collection_name not in names:
            self._client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=EMBEDDING_DIMENSION,
                    distance=Distance.COSINE
                )
            )
    
    def add_documents(
        self,
        texts: List[str],
        metadatas: Optional[List[Dict[str, Any]]] = None
    ) -> List[str]:
        """Add documents to the vector store."""
        if not texts:
            return []
        
        if metadatas is None:
            metadatas = [{} for _ in texts]
        
        documents = [
            Document(page_content=text, metadata=meta)
            for text, meta in zip(texts, metadatas)
        ]
        
        ids = self._vector_store.add_documents(documents)
        return ids
    
    def search(self, query: str, top_k: int = 20) -> List[Dict[str, Any]]:
        """Search for similar documents."""
        results = self._vector_store.similarity_search_with_score(
            query=query,
            k=top_k
        )
        
        formatted = []
        for doc, score in results:
            formatted.append({
                "id": doc.metadata.get("_id", ""),
                "score": score,
                "text": doc.page_content,
                "source": doc.metadata.get("source", "Unknown"),
                "chunk_index": doc.metadata.get("chunk_index", -1),
                "page_number": doc.metadata.get("page_number", -1),
                "metadata": doc.metadata
            })
        
        return formatted
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """Get collection statistics."""
        try:
            info = self._client.get_collection(self.collection_name)
            count = info.points_count or 0
            return {
                "name": self.collection_name,
                "vectors_count": count,
                "points_count": count,
                "status": str(info.status)
            }
        except Exception:
            return {
                "name": self.collection_name,
                "vectors_count": 0,
                "points_count": 0,
                "status": "error"
            }
    
    def clear_collection(self):
        """Delete and recreate the collection."""
        self._client.delete_collection(self.collection_name)
        self._ensure_collection_exists()
        
        self._vector_store = QdrantVectorStore(
            client=self._client,
            collection_name=self.collection_name,
            embedding=self._embedder
        )
    
    def collection_exists(self) -> bool:
        """Check if collection has documents."""
        stats = self.get_collection_stats()
        return stats["points_count"] > 0

    def document_exists(self, source: str) -> bool:
        """
        Check if a document with the given source name exists in the collection.
        
        Args:
            source: The source filename to check for
            
        Returns:
            True if document exists, False otherwise
        """
        try:
            result = self._client.scroll(
                collection_name=self.collection_name,
                scroll_filter=Filter(
                    must=[
                        FieldCondition(
                            key="metadata.source",
                            match=MatchValue(value=source)
                        )
                    ]
                ),
                limit=1,
                with_payload=False,
                with_vectors=False
            )
            points, _ = result
            return len(points) > 0
        except Exception:
            return False

    def get_loaded_sources(self) -> Set[str]:
        """
        Get set of all unique source names in the collection.
        
        Returns:
            Set of source filenames
        """
        sources = set()
        try:
            offset = None
            while True:
                result = self._client.scroll(
                    collection_name=self.collection_name,
                    limit=100,
                    offset=offset,
                    with_payload=True,
                    with_vectors=False
                )
                points, offset = result
                
                for point in points:
                    if point.payload:
                        # Check both possible metadata structures
                        source = None
                        if "metadata" in point.payload and isinstance(point.payload["metadata"], dict):
                            source = point.payload["metadata"].get("source")
                        if not source:
                            source = point.payload.get("source")
                        if source:
                            sources.add(source)
                
                if offset is None:
                    break
            
            return sources
        except Exception:
            return sources


_vector_store_instance = None


def get_vector_store() -> VectorStore:
    """Return singleton vector store instance."""
    global _vector_store_instance
    if _vector_store_instance is None:
        _vector_store_instance = VectorStore()
    return _vector_store_instance


def reset_vector_store():
    """Reset singleton for testing."""
    global _vector_store_instance
    if _vector_store_instance is not None:
        try:
            _vector_store_instance.clear_collection()
        except Exception:
            pass
    _vector_store_instance = None


if __name__ == "__main__":
    store = VectorStore(use_memory=True)
    
    texts = ["Atlas ERP sistemi.", "Finans modülü özellikleri."]
    metadatas = [
        {"source": "test.pdf", "chunk_index": 0},
        {"source": "test.pdf", "chunk_index": 1}
    ]
    
    store.add_documents(texts, metadatas)
    print(f"Stats: {store.get_collection_stats()}")
    
    results = store.search("ERP nedir?", top_k=2)
    for r in results:
        print(f"Score: {r['score']:.4f} - {r['text'][:50]}")