arxiv-rag-mvp / qdrant_manager.py
donb-hf's picture
update services
84deff7
# qdrant_manager.py
from langchain_qdrant import Qdrant
from langchain_openai import OpenAIEmbeddings
from typing import List, Dict, Any
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
class QdrantManager:
def __init__(self, api_url: str, api_key: str, collection_name: str):
self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
self.qdrant = Qdrant.from_documents(
[],
self.embeddings,
url=api_url,
api_key=api_key,
collection_name=collection_name,
)
def add_texts(self, chunks: List[str], metadata: Dict[str, Any]):
try:
self.qdrant.add_texts(chunks, metadatas=[metadata] * len(chunks))
logging.info(f"Added {len(chunks)} chunks to Qdrant for document: {metadata['title']}")
except Exception as e:
logging.error(f"Error adding to Qdrant: {str(e)}")
def get_embeddings(self, text: str) -> List[float]:
return self.embeddings.embed_query(text)