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])