feat(serve): migrate server implementation to use asyncio and grpc.aio for asynchronous handling (#6)

* feat(serve): migrate server implementation to use asyncio and grpc.aio for asynchronous handling

* fix(summarize): await an async function

Author: david20571015

README.md

@@ -20,7 +20,7 @@ The following environment variables are required (`export` them or place them in
 - `QDRANT_COLLECTION`: The Qdrant collection name.
 
 ```shell
-uv run serve --config configs/config.toml
+python3 scripts/serve.py --config configs/config.toml
 ```
 
 ## Features

llm_backend/search/service.py

@@ -36,16 +36,16 @@ def __init__(
         self.query_template = config.query.prompt_template
         self.similarity_top_k = config.query.similarity_top_k
 
-    def Search(
+    async def Search(
         self,
         request: search_pb2.SearchRequest,
-        context: grpc.ServicerContext,
+        context: grpc.aio.ServicerContext,
     ):
         prompt = self.query_template.format(keywords=", ".join(request.keywords))
         similarity_top_k = request.similarity_top_k or self.similarity_top_k
 
         retriever = self.index.as_retriever(similarity_top_k=similarity_top_k)
-        results: list[NodeWithScore] = retriever.retrieve(prompt)
+        results: list[NodeWithScore] = await retriever.aretrieve(prompt)
 
         return search_pb2.SearchResponse(
             results=[

llm_backend/summarize/service.py

@@ -55,11 +55,11 @@ def __init__(
         self.query_str = config.query.query_str
         self.content_formatter = CONTENT_FORMATTERS[config.query.content_format]
 
-    def Summarize(
+    async def Summarize(
         self,
         request: summarize_pb2.SummarizeRequest,
-        context: grpc.ServicerContext,
+        context: grpc.aio.ServicerContext,
     ):
         texts = self.content_formatter(request.contents)
-        summary = str(self.summarizer.get_response(self.query_str, texts))
+        summary = str(await self.summarizer.aget_response(self.query_str, texts))
         return summarize_pb2.SummarizeResponse(summary=summary)

pyproject.toml

@@ -18,7 +18,6 @@ dependencies = [
 ]
 
 [project.scripts]
-serve = "scripts.serve:main"
 gen-protos = "scripts.gen_protos:generate"
 
 [dependency-groups]

scripts/serve.py

@@ -1,7 +1,7 @@
 import argparse
+import asyncio
 import logging
 import os
-import sys
 import tomllib
 from concurrent import futures
 
@@ -10,33 +10,6 @@
 from llm_backend import Config, setup_search_service, setup_summarize_service
 
 
-def start_server(server: grpc.Server, config: Config):
-    try:
-        server_config = config.server
-        address = f"{server_config.host}:{server_config.port}"
-        server.add_insecure_port(address=address)
-        server.start()
-        logger.info("Server started on %s", address)
-        server.wait_for_termination()
-    except Exception as e:
-        logger.error("Error occurred while starting server: %s", e)
-        raise
-
-
-def serve(config: Config):
-    server = grpc.server(
-        futures.ThreadPoolExecutor(max_workers=config.server.max_workers)
-    )
-
-    setup_search_service(config, server)
-    logger.info("Added SearchService to server")
-
-    setup_summarize_service(config, server)
-    logger.info("Added SummarizeService to server")
-
-    start_server(server, config)
-
-
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -58,22 +31,44 @@ def load_config(config_path):
     return Config.model_validate(config)
 
 
-def main():
-    logging.basicConfig(
-        format="%(asctime)s\t%(levelname)s: %(message)s",
-        handlers=[
-            logging.StreamHandler(sys.stdout),
-            logging.FileHandler("server.log", "w"),
-        ],
+async def serve(config: Config, logger: logging.Logger):
+    server = grpc.aio.server(
+        futures.ThreadPoolExecutor(max_workers=config.server.max_workers)
     )
-    logger.setLevel(logging.INFO)
+    setup_search_service(config, server)
+    logger.info("Added SearchService to server")
 
-    args = parse_args()
-    config = load_config(args.config)
-    serve(config)
+    setup_summarize_service(config, server)
+    logger.info("Added SummarizeService to server")
+
+    server_config = config.server
+    address = f"{server_config.host}:{server_config.port}"
+    server.add_insecure_port(address=address)
+    logger.info("Server started on %s", address)
+
+    await server.start()
 
+    async def server_graceful_shutdown():
+        logging.info("Starting graceful shutdown...")
+        await server.stop(3)
+
+    _cleanup_coroutines.append(server_graceful_shutdown())
+
+    await server.wait_for_termination()
 
-logger = logging.getLogger("server")
 
 if __name__ == "__main__":
-    main()
+    logging.basicConfig(format="%(asctime)s\t%(levelname)s: %(message)s")
+    logger = logging.getLogger("server")
+    logger.setLevel(logging.INFO)
+
+    args = parse_args()
+    config = load_config(args.config)
+
+    loop = asyncio.new_event_loop()
+    _cleanup_coroutines = []
+    try:
+        loop.run_until_complete(serve(config, logger))
+    finally:
+        loop.run_until_complete(*_cleanup_coroutines)
+        loop.close()