from langchain_core.retrievers import BaseRetriever
from typing import List
from langchain_core.documents import Document
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain import globals
class CustomRetriever(BaseRetriever):
"""사용자 쿼리를 포함하는 상위 k 문서를 포함하는 장난감 검색기이다.
이 검색기는 동기화 메서드 _get_relevant_documents만 구현한다.
검색기가 파일 액세스 또는 네트워크 액세스를 포함하는 경우 `_aget_relevant_documents`의 네이티브 비동기 구현에서 이점을 얻을 수 있다.
평소와 같이 Runnables에는 다른 스레드에서 실행되는 동기화 구현에 위임하는 기본 비동기 구현이 제공된다."""
documents : List[Document]
"""List of documents to retrieve from."""
k : int
"""Number of top results to return"""
def _get_relevant_documents(self, query : str, *, run_manager : CallbackManagerForRetrieverRun) -> List[Document]:
"""검색기에 대한 동기화를 구현한다."""
matchingDocumentList = []
for document in self.documents:
if len(matchingDocumentList) > self.k:
return matchingDocumentList
if query.lower() in document.page_content.lower():
matchingDocumentList.append(document)
return matchingDocumentList
documentList = [
Document(
page_content = "Dogs are great companions, known for their loyalty and friendliness.",
metadata = {"type" : "dog", "trait" : "loyalty"}
),
Document(
page_content = "Cats are independent pets that often enjoy their own space.",
metadata = {"type" : "cat", "trait" : "independence"}
),
Document(
page_content = "Goldfish are popular pets for beginners, requiring relatively simple care.",
metadata = {"type" : "fish", "trait" : "low maintenance"}
),
Document(
page_content = "Parrots are intelligent birds capable of mimicking human speech.",
metadata = {"type" : "bird", "trait" : "intelligence"}
),
Document(
page_content = "Rabbits are social animals that need plenty of space to hop around.",
metadata = {"type" : "rabbit", "trait" : "social"}
)
]
globals.set_debug(False)
customRetriever = CustomRetriever(documents = documentList, k = 3)
resultDocumentList = customRetriever.batch(["dog", "cat"])
print(resultDocumentList)