Spaces:
Running
on
T4
Running
on
T4
refactoring for collection_name input
Browse files- utils/retriever.py +1 -1
- utils/vectorstore_interface.py +5 -5
utils/retriever.py
CHANGED
@@ -209,7 +209,7 @@ def get_context(
|
|
209 |
search_kwargs["filter"] = filter_obj
|
210 |
|
211 |
# Perform initial retrieval
|
212 |
-
retrieved_docs = vectorstore.search(query, top_k,
|
213 |
|
214 |
logging.info(f"Retrieved {len(retrieved_docs)} documents for query: {query[:50]}...")
|
215 |
|
|
|
209 |
search_kwargs["filter"] = filter_obj
|
210 |
|
211 |
# Perform initial retrieval
|
212 |
+
retrieved_docs = vectorstore.search(query, collection_name, top_k, **search_kwargs)
|
213 |
|
214 |
logging.info(f"Retrieved {len(retrieved_docs)} documents for query: {query[:50]}...")
|
215 |
|
utils/vectorstore_interface.py
CHANGED
@@ -55,7 +55,7 @@ class HuggingFaceSpacesVectorStore(VectorStoreInterface):
|
|
55 |
class QdrantVectorStore(VectorStoreInterface):
|
56 |
"""Vector store implementation for direct Qdrant connection."""
|
57 |
|
58 |
-
def __init__(self, url: str,
|
59 |
from qdrant_client import QdrantClient
|
60 |
from sentence_transformers import SentenceTransformer
|
61 |
|
@@ -101,7 +101,7 @@ class QdrantVectorStore(VectorStoreInterface):
|
|
101 |
|
102 |
return self._embedding_model
|
103 |
|
104 |
-
def search(self, query: str, top_k: int, **kwargs) -> List[Dict[str, Any]]:
|
105 |
"""Search using direct Qdrant connection."""
|
106 |
try:
|
107 |
# Get embedding model
|
@@ -118,7 +118,7 @@ class QdrantVectorStore(VectorStoreInterface):
|
|
118 |
# Perform vector search
|
119 |
logging.info(f"Searching Qdrant collection '{self.collection_name}' for top {top_k} results")
|
120 |
search_result = self.client.search(
|
121 |
-
collection_name=
|
122 |
query_vector=query_embedding,
|
123 |
query_filter=filter_obj, # Add filter support
|
124 |
limit=top_k,
|
@@ -162,10 +162,10 @@ def create_vectorstore(config: Any) -> VectorStoreInterface:
|
|
162 |
|
163 |
elif vectorstore_type.lower() == "qdrant":
|
164 |
url = config.get("vectorstore", "URL") # Use the full URL
|
165 |
-
collection_name = config.get("vectorstore", "COLLECTION_NAME")
|
166 |
api_key = auth_config["api_key"]
|
167 |
# Remove port parameter since it's included in the URL
|
168 |
-
return QdrantVectorStore(url,
|
169 |
|
170 |
else:
|
171 |
raise ValueError(f"Unsupported vector store type: {vectorstore_type}")
|
|
|
55 |
class QdrantVectorStore(VectorStoreInterface):
|
56 |
"""Vector store implementation for direct Qdrant connection."""
|
57 |
|
58 |
+
def __init__(self, url: str, api_key: Optional[str] = None):
|
59 |
from qdrant_client import QdrantClient
|
60 |
from sentence_transformers import SentenceTransformer
|
61 |
|
|
|
101 |
|
102 |
return self._embedding_model
|
103 |
|
104 |
+
def search(self, query: str, collection_name:str, top_k: int, **kwargs) -> List[Dict[str, Any]]:
|
105 |
"""Search using direct Qdrant connection."""
|
106 |
try:
|
107 |
# Get embedding model
|
|
|
118 |
# Perform vector search
|
119 |
logging.info(f"Searching Qdrant collection '{self.collection_name}' for top {top_k} results")
|
120 |
search_result = self.client.search(
|
121 |
+
collection_name=collection_name,
|
122 |
query_vector=query_embedding,
|
123 |
query_filter=filter_obj, # Add filter support
|
124 |
limit=top_k,
|
|
|
162 |
|
163 |
elif vectorstore_type.lower() == "qdrant":
|
164 |
url = config.get("vectorstore", "URL") # Use the full URL
|
165 |
+
#collection_name = config.get("vectorstore", "COLLECTION_NAME")
|
166 |
api_key = auth_config["api_key"]
|
167 |
# Remove port parameter since it's included in the URL
|
168 |
+
return QdrantVectorStore(url, api_key)
|
169 |
|
170 |
else:
|
171 |
raise ValueError(f"Unsupported vector store type: {vectorstore_type}")
|