class React

Author: devvratbhardwaj

aimon/extensions/__init__.py

@@ -1 +1 @@
-from .react import ReactConfig, react
+from .react import ReactConfig, React

aimon/extensions/react.py

@@ -1,60 +1,94 @@
-from aimon import Detect
+from aimon import Client, Detect
 from dataclasses import dataclass
 
 @dataclass
 class ReactConfig:
     publish: bool
+    async_mode:bool
     model_name: str
     max_attempts: int
     aimon_api_key: str
     application_name: str
-    values_returned: list[str]
     hallucination_threshold: float
-    aimon_config: dict[str, dict[str, str]]
+    framework:str = None
 
-## ReAct -> Reason and Act
+class React:
+    
+    ## Initailze the AIMon Client here
+    def __init__(self, llm_app, react_configuration, context_extractor):
+                                                                                   
+        self.llm_app = llm_app
+        self.context_extractor = context_extractor
+        self.react_configuration = react_configuration
+        self.client = Client(auth_header="Bearer {}".format(self.react_configuration.aimon_api_key))
+
+    def create_payload(self, context, user_query, user_instructions, generated_text):
+        
+        aimon_payload = {
+            'context':context,
+            'user_query':user_query,
+            'generated_text':generated_text,
+            'instructions':user_instructions,
+        }
+
+        aimon_payload['publish'] = self.react_configuration.publish
+        aimon_payload['async_mode'] = self.react_configuration.async_mode
+        aimon_payload['config'] = { 'hallucination': {'detector_name': 'default'},
+                                    'instruction_adherence': {'detector_name': 'default'},}
+
+        if self.react_configuration.publish:
+            aimon_payload['application_name'] = self.react_configuration.application_name
+            aimon_payload['model_name'] = self.react_configuration.model_name
+
+        return aimon_payload
 
-def react(  llm_app,
-            user_query,
-            user_instructions,
-            context_extractor,
-            react_configuration,
-        ):
+
+    ## ReAct -> Reason and Act
+    def react(self, user_query, user_instructions,):
+
+        llm_response = self.llm_app(user_query, user_instructions, reprompted_flag=False)
+        
+        context = self.context_extractor(user_query, user_instructions, llm_response)    
+
+        ## Generated text for LLM Response, if the user employs the LlamaIndex framework
+        if llm_response.response or self.react_configuration.framework=="llamaindex":
+            generated_text = llm_response.response
+        else:
+            generated_text = llm_response
 
-    detect = Detect(values_returned = react_configuration.values_returned,
-                    api_key = react_configuration.aimon_api_key,
-                    config = react_configuration.aimon_config,
-                    publish = react_configuration.publish,
-                    application_name = react_configuration.application_name,
-                    model_name = react_configuration.model_name,
-                )
-
-    llm_response = llm_app(user_query, user_instructions, reprompted_flag=False)
-
-    ## Decorating the context_extractor function with AIMon's "detect"
-    context_extractor = detect(context_extractor)
+        aimon_payload = self.create_payload(context, user_query, user_instructions, generated_text)
 
-    _, _, _, query_result, aimon_response = context_extractor(user_query, user_instructions, llm_response)
+        detect_response = self.client.inference.detect(body=[aimon_payload])
+
+        for _ in range(self.react_configuration.max_attempts):
 
-    for _ in range(react_configuration.max_attempts):
+            failed_instructions = []
+            ## Loop to check for failed instructions
+            for x in detect_response.instruction_adherence['results']:
+                if x['adherence'] == False:
+                    failed_instructions.append(x['instruction'])
 
-        failed_instructions = []
-        ## Loop to check for failed instructions
-        for x in aimon_response.detect_response.instruction_adherence['results']:
-            if x['adherence'] == False:
-                failed_instructions.append(x['instruction'])
+            hallucination_score = detect_response.hallucination['score'] 
 
-        hallucination_score = aimon_response.detect_response.hallucination['score'] 
+            ## Check whether the hallucination score is greater than the required threshold OR if any of the supplied instructions are not complied with
+            if  self.react_configuration.hallucination_threshold > 0 and \
+                (hallucination_score > self.react_configuration.hallucination_threshold or len(failed_instructions)>0): 
+                
+                llm_response = self.llm_app(user_query, user_instructions, reprompted_flag=True, hallucination_score=hallucination_score)
+                
+                context = self.context_extractor(user_query, user_instructions, llm_response)
+                
+                ## Generated text for LLM Response, if the user employs the LlamaIndex framework
+                if llm_response.response or self.react_configuration.framework=="llamaindex":
+                    generated_text = llm_response.response
+                else:
+                    generated_text = llm_response
 
-        ## Check whether the hallucination score is greater than the required threshold OR if any of the supplied instructions are not complied with
-        if  react_configuration.hallucination_threshold > 0 and \
-            (hallucination_score > react_configuration.hallucination_threshold or len(failed_instructions)>0): 
-            
-            llm_response = llm_app(user_query, user_instructions, reprompted_flag=True, hallucination_score=hallucination_score)
+                new_aimon_payload = self.create_payload(context, user_query, user_instructions, generated_text)
 
-            _, _, _, query_result, aimon_response = context_extractor(user_query, user_instructions, llm_response)
+                detect_response = self.client.inference.detect(body=[new_aimon_payload])
 
-    if hallucination_score > react_configuration.hallucination_threshold:
-        return f"The generated LLM response, even after {react_configuration.max_attempts} attempts of ReAct is still hallucinated. The response: {query_result}"
+        if hallucination_score > self.react_configuration.hallucination_threshold:
+            return f"The generated LLM response, even after {self.react_configuration.max_attempts} attempts of ReAct is still hallucinated. The response: {generated_text}"
 
-    return query_result
+        return generated_text