Source code for langchain_community.retrievers.breebs

from typing import List

import requests
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.documents.base import Document
from langchain_core.retrievers import BaseRetriever


class BreebsRetriever(BaseRetriever):
    """为 `Breebs` 创建的检索器类。

    请访问 https://www.breebs.com/ 了解更多信息。
    参数:
        breeb_key: 触发 breeb（特定主题上的专业知识药丸）的键。

    要检索所有可用的 Breebs 列表：可以调用 https://breebs.promptbreeders.com/web/listbreebs"""

    breeb_key: str
    url = "https://breebs.promptbreeders.com/knowledge"

    def __init__(self, breeb_key: str):
        super().__init__(breeb_key=breeb_key)  # type: ignore[call-arg]
        self.breeb_key = breeb_key

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """获取给定查询的上下文。
请注意，目前还没有得分。
"""
        r = requests.post(
            self.url,
            json={
                "breeb_key": self.breeb_key,
                "query": query,
            },
        )
        if r.status_code != 200:
            return []
        else:
            chunks = r.json()
            return [
                Document(
                    page_content=chunk["content"],
                    metadata={"source": chunk["source_url"], "score": 1},
                )
                for chunk in chunks
            ]