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