Merge pull request #38 from airsimonhan/streaming

update streaming output

Author: airsimonhan

README.md

@@ -181,6 +181,25 @@ docker-compose up
 
 For comprehensive Docker deployment instructions, configuration options, troubleshooting, and production deployment guidelines, please refer to our detailed [Docker Deployment Guide](./web-ui/DOCKER.md).
 
+## :bulb: Simple Test Demo
+
+1. Run by steps
+```bash
+conda activate rag
+cd Hyper_RAG/reproduce
+python reproduce/Step_0.py
+python reproduce/Step_1.py
+
+cd Hyper-RAG
+python -m uvicorn service_api:app --app-dir . --host 0.0.0.0 --port 8000
+```
+2. Open `testHTML_light.html` in your web browser.
+3. Selecting the model (`hyper`,`hyper-lite`,`naive`) and whether to output in streaming mode
+
+<div align="center">
+  <img src="./assets/hyperrag-streaming.gif" alt="Efficiency analysis" width="80%" />
+</div>
+
 ## :checkered_flag: Evaluation
 In this work, we propose two evaluation strategys: the **selection-based** and **scoring-based** evaluation. 
 

assets/hyperrag-streaming.gif

@@ -33,6 +33,7 @@
     convert_response_to_json,
     logger,
     set_logger,
+    limit_async_gen_call
 )
 from .base import (
     BaseKVStorage,
@@ -42,6 +43,8 @@
     BaseHypergraphStorage,
 )
 
+from .operate import hyper_query_stream, hyper_query_lite_stream, naive_query_stream, llm_query_stream
+
 
 def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
     try:
@@ -60,7 +63,7 @@ class HyperRAG:
     working_dir: str = field(
         default_factory=lambda: f"./HyperRAG_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
     )
-    print(working_dir)
+    # print(working_dir)
 
     current_log_level = logger.level
     log_level: str = field(default=current_log_level)
@@ -78,7 +81,7 @@ class HyperRAG:
     relation_keywords_to_max_tokens: int = 100
 
     embedding_func: EmbeddingFunc = field(default_factory=lambda: openai_embedding)
-    embedding_batch_num: int = 32
+    embedding_batch_num: int = 8
     embedding_func_max_async: int = 16
 
     # LLM
@@ -89,6 +92,8 @@ class HyperRAG:
     llm_model_max_async: int = 16
     llm_model_kwargs: dict = field(default_factory=dict)
 
+    llm_model_stream_func: callable = None
+
     # storage
     key_string_value_json_storage_cls: Type[BaseKVStorage] = JsonKVStorage
     vector_db_storage_cls: Type[BaseVectorStorage] = NanoVectorDBStorage
@@ -166,6 +171,16 @@ def __post_init__(self):
             )
         )
 
+        if getattr(self, "llm_model_stream_func", None) is not None:
+            # 先把 hashing_kv 注入到 stream func（供 openai_complete_stream_if_cache 使用）
+            self.llm_model_stream_func = limit_async_gen_call(self.llm_model_max_async)(
+                partial(
+                    self.llm_model_stream_func,
+                    hashing_kv=self.llm_response_cache,
+                    **self.llm_model_kwargs,
+                )
+            )
+
     def insert(self, string_or_strings):
         loop = always_get_an_event_loop()
         return loop.run_until_complete(self.ainsert(string_or_strings))
@@ -304,6 +319,61 @@ async def aquery(self, query: str, param: QueryParam = QueryParam()):
         await self._query_done()
         return response
 
+    async def astream_query(self, query: str, param: QueryParam = QueryParam()):
+        """
+        流式查询：返回 async generator（逐 token / 逐块）
+        依赖 self.llm_model_stream_func，不提供则抛错。
+        """
+        if self.llm_model_stream_func is None:
+            raise AttributeError("llm_model_stream_func is not set, streaming is unavailable.")
+
+        # 把 stream func 放进 global_config
+        cfg = asdict(self)
+        cfg["llm_model_stream_func"] = self.llm_model_stream_func
+
+        if param.mode == "hyper":
+            async for tok in hyper_query_stream(
+                    query,
+                    self.chunk_entity_relation_hypergraph,
+                    self.entities_vdb,
+                    self.relationships_vdb,
+                    self.text_chunks,
+                    param,
+                    cfg,
+            ):
+                yield tok
+
+        elif param.mode == "hyper-lite":
+            async for tok in hyper_query_lite_stream(
+                    query,
+                    self.chunk_entity_relation_hypergraph,
+                    self.entities_vdb,
+                    self.text_chunks,
+                    param,
+                    cfg,
+            ):
+                yield tok
+
+        elif param.mode == "naive":
+            async for tok in naive_query_stream(
+                    query,
+                    self.chunks_vdb,
+                    self.text_chunks,
+                    param,
+                    cfg,
+            ):
+                yield tok
+
+        elif param.mode == "llm":
+            async for tok in llm_query_stream(query, param, cfg):
+                yield tok
+
+        else:
+            raise ValueError(f"Unknown mode {param.mode}")
+
+        await self._query_done()
+
+
     async def _query_done(self):
         tasks = []
         for storage_inst in [self.llm_response_cache]:

hyperrag/hyperrag.py

@@ -73,6 +73,68 @@ async def openai_complete_if_cache(
         )
     return response.choices[0].message.content
 
+async def openai_complete_stream_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    base_url=None,
+    api_key=None,
+    chunk_size: int = 32,
+    **kwargs,
+):
+    """
+    OpenAI-compatible 流式输出（async generator）
+    - 命中缓存：按 chunk_size 分块 yield
+    - 不命中：stream=True 逐 token yield，并在结束后写缓存
+    """
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+
+    openai_async_client = (
+        AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
+    )
+
+    hashing_kv: BaseKVStorage = kwargs.pop("hashing_kv", None)
+
+    messages = []
+    if system_prompt is not None:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    # 1) cache 命中：直接回放
+    if hashing_kv is not None:
+        args_hash = compute_args_hash(model, messages)
+        if_cache_return = await hashing_kv.get_by_id(args_hash)
+        if if_cache_return is not None:
+            cached = if_cache_return["return"] or ""
+            # 按块 yield，避免一次性返回
+            for i in range(0, len(cached), chunk_size):
+                yield cached[i:i + chunk_size]
+            return
+
+    # 2) cache 未命中：真实 stream
+    full_text = []
+    stream = await openai_async_client.chat.completions.create(
+        model=model,
+        messages=messages,
+        stream=True,
+        **kwargs,
+    )
+
+    async for event in stream:
+        delta = None
+        if event.choices:
+            delta = getattr(event.choices[0].delta, "content", None)
+        if delta:
+            full_text.append(delta)
+            yield delta
+
+    # 3) 写入 cache
+    if hashing_kv is not None:
+        text = "".join(full_text)
+        await hashing_kv.upsert({args_hash: {"return": text, "model": model}})
 
 @retry(
     stop=stop_after_attempt(3),