File size: 882 Bytes
dc22afe 29cf6b3 dc22afe | 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 |
# hybrid_retrieval.py
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
import os
print(os.listdir())
class Retriever:
def __init__(self,
chroma_dir: str = 'src/chroma_db',
collection_name: str = 'my_collection'
):
self.vectorstore = Chroma(
collection_name=collection_name,
persist_directory=chroma_dir
)
def retrieve(self, query: str):
"""
Récupère les documents via le retriever Embedding.
"""
return self.vectorstore.similarity_search_with_score(query,k=3)
def add_document(self, doc: str, metadata: dict, id: str):
self.vectorstore.collection.add(documents=[doc], metadatas=[metadata], ids=[id])
|