Feature Request: Add Async Support for acompress_documents Method in JinaRerank Class

[fbossiere]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Add asynchronous support to JinaRerank and implement acompress_documents

Motivation

The JinaRerank class currently uses synchronous HTTP requests via the requests library to interact with the Jina Rerank API. This design limits the potential for high-throughput, parallel processing in asynchronous applications. With growing demand for scalable, real-time document reranking solutions, supporting asynchronous operations is essential for handling multiple rerank requests efficiently, reducing latency, and improving overall performance in distributed or high-load environments.

Additionally, the JinaRerank class does not include an acompress_documents method, which means that users cannot perform asynchronous document compression. Implementing this method would extend the utility of JinaRerank in modern applications that rely on asynchronous workflows.

Proposal (If applicable)

1. Replace requests with aiohttp:

   • Transition from requests to aiohttp, a popular asynchronous HTTP client library, to support non-blocking API calls to the Jina Rerank API.
   • Update all HTTP request methods in the JinaRerank class, such as rerank and compress_documents, to use asynchronous functions (async def) with await statements for aiohttp requests.
   • Add exception handling to manage potential errors (e.g., timeout, network issues) specific to asynchronous operations.

2. Implement acompress_documents Method:

   • Create an asynchronous version of compress_documents called acompress_documents.
   • This method should mirror the logic of compress_documents, using the asynchronous rerank function and processing the response asynchronously to construct a list of compressed documents.
   • Implement support for asynchronous callbacks, allowing users to trigger async callback functions during the document compression process, if specified.

3. Update Environment Validation for Async Compatibility:

   • Modify the validate_environment method to ensure that aiohttp.ClientSession is initialized for asynchronous API requests instead of a synchronous requests.Session.
   • Add logic to check for asynchronous requirements in the environment (e.g., proper handling of async API keys) to ensure compatibility with the new async design.

4. Backward Compatibility:

   • For users who rely on synchronous workflows, consider maintaining the current requests-based version as an option or offering a synchronous compress_documents method.
   • Document the differences between synchronous and asynchronous modes, including usage guidelines for each mode, to ensure that users can adopt the new async functionality without disruptions.

5. Example Usage:

   • Provide examples in the documentation showcasing how to initialize and use JinaRerank with the async methods acompress_documents and rerank.
   • Include sample code illustrating common async patterns, such as awaiting multiple rerank calls with asyncio.gather() for concurrent processing.

Benefits of the Proposed Solution:

• Scalability: Asynchronous API calls would enable JinaRerank to handle multiple rerank requests in parallel, ideal for applications that process a high volume of documents or need to maintain low latency.
• Performance: Switching to async I/O operations can significantly reduce blocking time, allowing the application to utilize resources more efficiently and improve response times.
• Enhanced Flexibility: With both compress_documents and acompress_documents methods, users can choose between synchronous and asynchronous workflows based on their application's requirements, making JinaRerank more versatile.

Code Suggestions:

Here's a high-level example of what the changes might look like:

import aiohttp
from langchain_core.utils import get_from_dict_or_env

JINA_API_URL: str = "https://api.jina.ai/v1/rerank"

class JinaRerank(BaseDocumentCompressor):
    session: Optional[aiohttp.ClientSession] = None

    @classmethod
    async def async_validate_environment(cls, values: Dict) -> Any:
        jina_api_key = get_from_dict_or_env(values, "jina_api_key", "JINA_API_KEY")
        user_agent = values.get("user_agent", "langchain")
        session = aiohttp.ClientSession()
        session.headers.update({
            "Authorization": f"Bearer {jina_api_key}",
            "Accept-Encoding": "identity",
            "Content-type": "application/json",
            "user-agent": user_agent,
        })
        values["session"] = session
        return values

    async def async_rerank(self, documents: Sequence[Union[str, Document, dict]], query: str, **kwargs) -> List[Dict[str, Any]]:
        if len(documents) == 0:
            return []
        docs = [doc.page_content if isinstance(doc, Document) else doc for doc in documents]
        data = {
            "query": query,
            "documents": docs,
            "model": kwargs.get("model", self.model),
            "top_n": kwargs.get("top_n", self.top_n),
        }
        
        async with self.session.post(JINA_API_URL, json=data) as resp:
            resp_json = await resp.json()
            if "results" not in resp_json:
                raise RuntimeError(resp_json["detail"])
            
            results = resp_json["results"]
            return [{"index": res["index"], "relevance_score": res["relevance_score"]} for res in results]

    async def acompress_documents(self, documents: Sequence[Document], query: str, callbacks: Optional[Callbacks] = None) -> Sequence[Document]:
        reranked = await self.async_rerank(documents, query)
        compressed = [
            Document(documents[res["index"]].page_content, metadata={**documents[res["index"]].metadata, "relevance_score": res["relevance_score"]})
            for res in reranked
        ]
        return compressed

This proposal provides an actionable roadmap to add asynchronous support to JinaRerank, enhancing its performance and making it more suitable for modern, high-demand applications.

[fbossiere]

The implementation of acompress_documents is not necessary in the JinaRerank class as this functionality is already available through the parent class BaseDocumentCompressor.

Indeed, BaseDocumentCompressor already implements the acompress_documents method which automatically executes the synchronous compress_documents method in a separate thread using asyncio's run_in_executor. This approach allows any synchronous method to be executed asynchronously without having to rewrite it.

Therefore, rather than reimplementing acompress_documents, we can rely on the existing implementation from the parent class which already handles asynchronous execution correctly.

This solution is simpler and avoids code duplication while ensuring proper asynchronous behavior.

If performance optimizations are needed, it would be more relevant to focus on using aiohttp for API calls rather than rewriting acompress_documents.