Update generate_answer_node.py

Author: VinciGit00

scrapegraphai/nodes/generate_answer_node.py

@@ -1,10 +1,8 @@
-"""
-Generate answer_node
-"""
 import re
 import json
 from typing import List, Optional
 import requests
+import asyncio
 from tqdm import tqdm
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
@@ -25,16 +23,6 @@ class GenerateAnswerNode(BaseNode):
     and the content extracted from a webpage. It constructs a prompt from the user's input
     and the scraped content, feeds it to the LLM, and parses the LLM's response to produce
     an answer.
-
-    Attributes:
-        llm_model: An instance of a language model client, configured for generating answers.
-        verbose (bool): A flag indicating whether to show print statements during execution.
-
-    Args:
-        input (str): Boolean expression defining the input keys needed from the state.
-        output (List[str]): List of output keys to be updated in the state.
-        node_config (dict): Additional configuration for the node.
-        node_name (str): The unique identifier name for the node, defaulting to "GenerateAnswer".
     """
 
     def __init__(
@@ -57,7 +45,33 @@ def __init__(
         self.additional_info = node_config.get("additional_info", "")
         self.api_key = node_config.get("config", {}).get("llm", {}).get("api_key", "")
 
+    async def _process_chunks_async(self, chunks, templates, user_prompt, format_instructions):
+        async def send_request(prompt):
+            url = "https://api.openai.com/v1/chat/completions"
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self.api_key}"
+            }
+            response = await requests.post(url, headers=headers, json={
+                "model": self.llm_model.model_name,
+                "messages": [{"role": "user", "content": prompt}],
+                "temperature": 0
+            }, timeout=10)
+            response_text = response.json()['choices'][0]['message']['content']
+            return parse_response_to_dict(response_text)
 
+        tasks = []
+        for i, chunk in enumerate(chunks):
+            prompt = templates['chunks'].format(
+                question=user_prompt,
+                context=chunk,
+                chunk_id=i + 1,
+                format_instructions=format_instructions
+            )
+            tasks.append(send_request(prompt))
+        
+        results = await asyncio.gather(*tasks)
+        return results
 
     def execute(self, state: dict) -> dict:
         self.logger.info(f"--- Executing {self.node_name} Node ---")
@@ -76,63 +90,60 @@ def execute(self, state: dict) -> dict:
                 'merge': TEMPLATE_MERGE_MD
             }
 
-            url = "https://api.openai.com/v1/chat/completions"
-            headers = {
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {self.api_key}"
-            }
-
             if len(doc) == 1:
                 prompt = templates['no_chunks'].format(
                     question=user_prompt,
                     context=doc[0],
                     format_instructions=format_instructions
                 )
-                response = requests.post(url, headers=headers, json={
-                    "model": self.llm_model.model_name,
-                    "messages": [{"role": "user", "content": prompt}],
-                    "temperature": 0
-                }, timeout=10)
+                response = requests.post(
+                    url="https://api.openai.com/v1/chat/completions",
+                    headers={
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {self.api_key}"
+                    },
+                    json={
+                        "model": self.llm_model.model_name,
+                        "messages": [{"role": "user", "content": prompt}],
+                        "temperature": 0
+                    },
+                    timeout=10
+                )
 
                 response_text = response.json()['choices'][0]['message']['content']
                 cleaned_response = parse_response_to_dict(response_text)
                 state.update({self.output[0]: cleaned_response})
                 return state
 
-            chunks_responses = []
-            for i, chunk in enumerate(
-                tqdm(doc, desc="Processing chunks",
-                     disable=not self.verbose)):
-                prompt = templates['chunks'].format(
+            else:
+                chunks_responses = asyncio.run(
+                    self._process_chunks_async(doc, templates, user_prompt, format_instructions)
+                )
+
+                merge_context = " ".join([json.dumps(chunk) for chunk in chunks_responses])
+                merge_prompt = templates['merge'].format(
                     question=user_prompt,
-                    context=chunk,
-                    chunk_id=i + 1,
+                    context=merge_context,
                     format_instructions=format_instructions
                 )
-                response = requests.post(url, headers=headers, json={
-                    "model": self.llm_model.model_name,
-                    "messages": [{"role": "user", "content": prompt}],
-                    "temperature": 0
-                }, timeout=10)
-                chunk_response = response.json()['choices'][0]['message']['content']
-                cleaned_chunk_response = parse_response_to_dict(chunk_response)
-                chunks_responses.append(cleaned_chunk_response)
-
-            merge_context = " ".join([json.dumps(chunk) for chunk in chunks_responses])
-            merge_prompt = templates['merge'].format(
-                question=user_prompt,
-                context=merge_context,
-                format_instructions=format_instructions
-            )
-            response = requests.post(url, headers=headers, json={
-                "model": self.llm_model.model_name,
-                "messages": [{"role": "user", "content": merge_prompt}],
-                "temperature": 0
-            }, timeout=10)
-            response_text = response.json()['choices'][0]['message']['content']
-            cleaned_response = parse_response_to_dict(response_text)
-            state.update({self.output[0]: cleaned_response})
-            return state
+                response = requests.post(
+                    url="https://api.openai.com/v1/chat/completions",
+                    headers={
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {self.api_key}"
+                    },
+                    json={
+                        "model": self.llm_model.model_name,
+                        "messages": [{"role": "user", "content": merge_prompt}],
+                        "temperature": 0
+                    },
+                    timeout=10
+                )
+
+                response_text = response.json()['choices'][0]['message']['content']
+                cleaned_response = parse_response_to_dict(response_text)
+                state.update({self.output[0]: cleaned_response})
+                return state
 
         else:
             templates = {
@@ -142,13 +153,15 @@ def execute(self, state: dict) -> dict:
             }
 
             if self.additional_info:
-                templates = {key: self.additional_info + template for key, template in templates.items()}
+                templates = {key: self.additional_info + 
+                             template for key, template in templates.items()}
 
             if len(doc) == 1:
                 prompt = PromptTemplate(
                     template=templates['no_chunks'],
                     input_variables=["question"],
-                    partial_variables={"context": doc, "format_instructions": format_instructions}
+                    partial_variables={"context": doc[0], 
+                                       "format_instructions": format_instructions}
                 )
                 chain = prompt | self.llm_model | output_parser
                 answer = chain.invoke({"question": user_prompt})