Update AI.py

AI.py

@@ -1,27 +1,22 @@
 import requests
-from requests import get
 from uuid import uuid4
-import io
-import re
 import json
-from typing import Any
-import logging
+from typing import Any, Dict, Optional
+
 class OPENGPT:
     def __init__(
         self,
         max_tokens: int = 600,
         timeout: int = 30,
-        intro: str = None,
-        filepath: str = None,
+        intro: Optional[str] = None,
+        filepath: Optional[str] = None,
         update_file: bool = True,
-        proxies: dict = {},
+        proxies: Dict[str, str] = {},
         history_offset: int = 10250,
-        act: str = None,
+        act: Optional[str] = None,
     ):
         """Instantiates OPENGPT
-
         Args:
-            is_conversation (bool, optional): Flag for chatting conversationally. Defaults to True
             max_tokens (int, optional): Maximum number of tokens to be generated upon completion. Defaults to 600.
             timeout (int, optional): Http request timeout. Defaults to 30.
             intro (str, optional): Conversation introductory prompt. Defaults to None.
@@ -33,9 +28,7 @@ class OPENGPT:
         """
         self.session = requests.Session()
         self.max_tokens_to_sample = max_tokens
-        self.chat_endpoint = (
-            "https://opengpts-example-vz4y4ooboq-uc.a.run.app/runs/stream"
-        )
+        self.chat_endpoint = "https://opengpts-example-vz4y4ooboq-uc.a.run.app/runs/stream"
         self.stream_chunk_size = 64
         self.timeout = timeout
         self.last_response = {}
@@ -56,51 +49,22 @@ class OPENGPT:
             "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
         }
 
-
     def ask(
         self,
         prompt: str,
         stream: bool = False,
         raw: bool = False,
-        optimizer: str = None,
-    ) -> dict:
+        optimizer: Optional[str] = None,
+    ) -> Dict[str, Any]:
         """Chat with AI
-
         Args:
             prompt (str): Prompt to be send.
             stream (bool, optional): Flag for streaming response. Defaults to False.
             raw (bool, optional): Stream back raw response as received. Defaults to False.
             optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`. Defaults to None.
-            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
         Returns:
-           dict : {}
-        ```json
-        {
-            "messages": [
-                {
-                    "content": "Hello there",
-                    "additional_kwargs": {},
-                    "type": "human",
-                    "example": false
-                },
-                {
-                    "content": "Hello! How can I assist you today?",
-                    "additional_kwargs": {
-                    "agent": {
-                        "return_values": {
-                            "output": "Hello! How can I assist you today?"
-                            },
-                        "log": "Hello! How can I assist you today?",
-                        "type": "AgentFinish"
-                    }
-                },
-                "type": "ai",
-                "example": false
-                }]
-        }
-        ```
+           dict: Response from the AI.
         """
-
         self.session.headers.update(self.headers)
         self.session.headers.update(
             dict(
@@ -145,7 +109,6 @@ class OPENGPT:
                     yield value if raw else resp[1]
                 except json.decoder.JSONDecodeError:
                     pass
-            )
 
         def for_non_stream():
             for _ in for_stream():
@@ -158,43 +121,31 @@ class OPENGPT:
         self,
         prompt: str,
         stream: bool = False,
-        optimizer: str = None,
+        optimizer: Optional[str] = None,
     ) -> str:
         """Generate response `str`
         Args:
             prompt (str): Prompt to be send.
             stream (bool, optional): Flag for streaming response. Defaults to False.
             optimizer (str, optional): Prompt optimizer name - `[code, shell_command]`. Defaults to None.
-            conversationally (bool, optional): Chat conversationally when using optimizer. Defaults to False.
         Returns:
             str: Response generated
         """
-
         def for_stream():
-            for response in self.ask(
-                prompt, True
-            ):
+            for response in self.ask(prompt, True):
                 yield self.get_message(response)
 
         def for_non_stream():
-            return self.get_message(
-                self.ask(
-                    prompt,
-                    False,
-                    optimizer=optimizer,
-                )
-            )
+            return self.get_message(self.ask(prompt, False, optimizer=optimizer))
 
         return for_stream() if stream else for_non_stream()
 
-    def get_message(self, response: dict) -> str:
+    def get_message(self, response: Dict[str, Any]) -> str:
         """Retrieves message only from response
-
         Args:
             response (dict): Response generated by `self.ask`
-
         Returns:
             str: Message extracted
         """
         assert isinstance(response, dict), "Response should be of dict data-type only"
-        return response["content"]
+        return response["content"]