Add GenAI chat feature

Author: om-khade-algobulls

pyalgotrading/algobulls/api.py

@@ -30,6 +30,8 @@ def __init__(self, connection):
         self.__key_papertrading = {}  # strategy-cstc_id mapping
         self.__key_realtrading = {}  # strategy-cstc_id mapping
         self.pattern = re.compile(r'(?<!^)(?=[A-Z])')
+        self.genai_api_key = None
+        self.genai_session_id = None
 
     def __convert(self, _dict):
         # Helps convert _dict keys from camelcase to snakecase
@@ -473,3 +475,47 @@ def get_reports(self, strategy_code: str, trading_type: TradingType, report_type
         response = self._send_request(endpoint=endpoint, params=params)
 
         return response
+
+    def get_genai(self, user_prompt: str, session_id: int, chat_gpt_model: str = ''):
+        """
+        Fetch GenAI response.
+
+        Args:
+           user_prompt: User question
+           session_id: Session id of the GenAI session
+           chat_gpt_model: Chat gpt model name
+        Returns:
+           GenAI response
+
+        Info: ENDPOINT
+           `GET` v1/build/python/genai              Get GenAI response
+        """
+        endpoint = 'v1/build/python/genai'
+        params = {"userPrompt": user_prompt, 'sessionId': self.genai_session_id, 'openaiApiKey': self.genai_api_key, 'chat_gpt_model': chat_gpt_model}
+        response = self._send_request(endpoint=endpoint, params=params)
+
+        return response
+
+    def get_genai_response(self):
+        """
+        Fetch GenAI response.
+
+        Args:
+        Returns:
+           GenAI response for current session. Last active session is used when session_id is None.
+
+        Info: ENDPOINT
+           `GET` v1/build/python/genai/response     Pooling API to get response in case of timeout
+        """
+        endpoint = 'v1/build/python/genai/response'
+        params = {'sessionId': self.genai_session_id}
+        response = self._send_request(endpoint=endpoint, params=params)
+
+        return response
+
+    def get_genai_sessions(self, page_no):
+        endpoint = 'v1/build/python/genai/sessions'
+        params = {'sessionId': self.genai_session_id}
+        response = self._send_request(endpoint=endpoint, params=params)
+
+        return response

pyalgotrading/algobulls/connection.py

@@ -90,15 +90,54 @@ def set_generative_ai_keys(self, api_key, secret_key, token_key):
             # todo: learn about different generative AIs and how many and what keys are required ?
                 #  also confirm with backend the format of API keys to be received
         """
-
+        assert isinstance(api_key, str), f'Argument "api_key" should be a string'
+        self.api.openai_key = api_key
         return 'SUCCESS' or 'FAILURE'
 
-    def generate_strategy(self):
-        input_prompt = str(input())
+    def start_chat(self, chat_gpt_model):
+        while True:
+            user_prompt = str(input())
+            if user_prompt.lower() == 'exit':
+                print("Thanks for the chat")
+                return
+
+            response = self.api.get_genai(user_prompt, chat_gpt_model)
+            while response['status_code'] == 504:
+                response = self.api.get_genai_response()
+
+            print(f"GenAI: {response['message']}")
+
+    def continue_from_previous_session(self, page_no):
+        """
+        display previous sessions
+        Returns:
 
-        # call the api
+        """
+        customer_genai_sessions = self.api.get_genai_sessions(page_no)
+        for i, session in enumerate(customer_genai_sessions):
+            print(f"Session {i}: ID: {session['id']}, Started: {session['last_user_prompt']}")
 
-        return  # strategy in strings
+        if len(customer_genai_sessions) < 20:
+            print("End")
+        else:
+            print(f"Type 'next' to view the next 20 sessions.")
+
+        user_input = input("Enter session number or 'next': ")
+        if user_input.lower() == "next" and len(customer_genai_sessions) > 20:
+            self.continue_from_previous_session(page_no=page_no + 1)
+        elif user_input.isdigit() and 1 <= int(user_input) <= len(customer_genai_sessions):
+            selected_session_index = page_no + int(user_input) - 1
+            selected_session_id = customer_genai_sessions[selected_session_index]["id"]
+            self.api.genai_api_key = selected_session_id
+
+    def initiate_chat(self, start_fresh=None, chat_gpt_model=None):
+        if start_fresh:
+            # reset session
+            self.api.genai_api_key = None
+        elif start_fresh is not None:
+            self.continue_from_previous_session(page_no=1)
+
+        self.start_chat(chat_gpt_model)
 
     def save_latest_generated_strategy(self):
         pass