Merge pull request #88 from seungwon2/master

[Fix] : fix HyDE RAG fusion error

Author: gonsoomoon-ml

genai/aws-gen-ai-kr/utils/rag_summit.py

@@ -35,6 +35,7 @@
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain.embeddings.sagemaker_endpoint import EmbeddingsContentHandler
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate
 
 import threading
@@ -313,7 +314,19 @@ class retriever_utils():
         length_function=len,
     )
     token_limit = 300
+    
+    @classmethod
+    def control_streaming_mode(cls, llm, stream=True):
+
+        if stream:
+            llm.streaming = True
+            llm.callbacks = [StreamingStdOutCallbackHandler()]
+        else:
+            llm.streaming = False
+            llm.callbacks = None
 
+        return llm
+    
     @classmethod
     # semantic search based
     def get_semantic_similar_docs_by_langchain(cls, **kwargs):
@@ -476,9 +489,10 @@ def get_rag_fusion_similar_docs(cls, **kwargs):
         llm_text = kwargs["llm_text"]
         query_augmentation_size = kwargs["query_augmentation_size"]
         query_transformation_prompt = kwargs["query_transformation_prompt"]
-               
+
+        llm_text = cls.control_streaming_mode(llm_text, stream=False) ## trun off llm streaming
         generate_queries = query_transformation_prompt | llm_text | StrOutputParser() | (lambda x: x.split("\n"))
-        
+
         rag_fusion_query = generate_queries.invoke(
             {
                 "query": kwargs["query"],
@@ -495,6 +509,9 @@ def get_rag_fusion_similar_docs(cls, **kwargs):
             print("\n")
             print("===== RAG-Fusion Queries =====")
             print(rag_fusion_query)
+        
+        llm_text = cls.control_streaming_mode(llm_text, stream=True)## trun on llm streaming
+
 
         tasks = []
         for query in rag_fusion_query:
@@ -547,6 +564,7 @@ def _get_hyde_response(query, prompt, llm_text):
         hyde_query = kwargs["hyde_query"]
 
         tasks = []
+        llm_text = cls.control_streaming_mode(llm_text, stream=False) ## trun off llm streaming
         for template_type in hyde_query:
             hyde_response = partial(
                 _get_hyde_response,
@@ -557,10 +575,9 @@ def _get_hyde_response(query, prompt, llm_text):
             tasks.append(cls.hyde_pool.apply_async(hyde_response,))
         hyde_answers = [task.get() for task in tasks]
         hyde_answers.insert(0, query)
-        augmentation = hyde_answers
-        
 
         tasks = []
+        llm_text = cls.control_streaming_mode(llm_text, stream=True) ## trun on llm streaming
         for hyde_answer in hyde_answers:
             semantic_search = partial(
                 cls.get_semantic_similar_docs,
@@ -583,10 +600,13 @@ def _get_hyde_response(query, prompt, llm_text):
             c=60,
             k=kwargs["k"],
         )
+        
         if kwargs["verbose"]:
             print("\n")
             print("===== HyDE Answers =====")
             print(hyde_answers)
+            
+        augmentation = hyde_answers[1]
 
         return similar_docs