File size: 2,837 Bytes
a33458e |
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 |
import os
import sys
from langchain.vectorstores import Qdrant
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
# Add project root to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from app.config import VECTOR_DB_PATH, COLLECTION_NAME
from app.core.llm import get_llm, get_embeddings, get_chat_model
class MemoryManager:
"""Manages the RAG memory system using a vector database."""
def __init__(self):
self.embeddings = get_embeddings()
self.llm = get_llm()
self.chat_model = get_chat_model()
self.client = self._init_qdrant_client()
self.vectorstore = self._init_vector_store()
self.memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True
)
def _init_qdrant_client(self):
"""Initialize the Qdrant client."""
os.makedirs(VECTOR_DB_PATH, exist_ok=True)
return QdrantClient(path=VECTOR_DB_PATH)
def _init_vector_store(self):
"""Initialize the vector store."""
collections = self.client.get_collections().collections
collection_names = [collection.name for collection in collections]
# Get vector dimension from the embedding model
vector_size = len(self.embeddings.embed_query("test"))
if COLLECTION_NAME not in collection_names:
self.client.create_collection(
collection_name=COLLECTION_NAME,
vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
)
return Qdrant(
client=self.client,
collection_name=COLLECTION_NAME,
embeddings=self.embeddings
)
def get_retriever(self):
"""Get the retriever for RAG."""
return self.vectorstore.as_retriever(
search_type="similarity",
search_kwargs={"k": 5}
)
def create_rag_chain(self):
"""Create a RAG chain for question answering."""
# Using the chat model created with the regular LLM
return ConversationalRetrievalChain.from_llm(
llm=self.llm,
retriever=self.get_retriever(),
memory=self.memory,
return_source_documents=True
)
def add_texts(self, texts, metadatas=None):
"""Add texts to the vector store."""
return self.vectorstore.add_texts(texts=texts, metadatas=metadatas)
def similarity_search(self, query, k=5):
"""Perform a similarity search."""
return self.vectorstore.similarity_search(query, k=k) |