openassistant/app/providers/r2r.py

from typing import Optional, Any
from r2r import R2RClient

from app.libs.util import verify_jwt_expiration
from config.llm import tool_settings
import nest_asyncio
import asyncio

# Apply nest_asyncio to allow nested event loops
nest_asyncio.apply()


class R2R:
    client: R2RClient

    def __init__(self):
        self.auth_enabled = tool_settings.R2R_USERNAME and tool_settings.R2R_PASSWORD
        self.client = None

    def init(self):
        if not self.auth_enabled:
            return
        if not self.client:
            self.client = R2RClient(tool_settings.R2R_BASE_URL)
            asyncio.run(
                self.client.users.login(
                    tool_settings.R2R_USERNAME, tool_settings.R2R_PASSWORD
                )
            )

    def ingest_file(self, file_path: str, metadata: Optional[dict]):
        self._check_login()
        ingest_response = self.client.documents.create(
            file_path=file_path, metadata=metadata if metadata else None, id=None
        )
        return ingest_response.get("results")

    def search(self, query: str, filters: dict[str, Any]):
        self._check_login()
        search_response = self.client.retrieval.search(
            query=query,
            search_settings={
                "filters": filters,
                "limit": tool_settings.R2R_SEARCH_LIMIT,
            },
        )
        return search_response.get("results").get("chunk_search_results")

    def _check_login(self):
        if not self.auth_enabled:
            return
        if verify_jwt_expiration(self.client.access_token):
            return
        else:
            self.init()


# 创建 R2R 实例
r2r = R2R()