Repository: CIME-Software/perplexipy (language: python)

Issue:
Implement API call with streaming



Discussion / hints:

Fix patch:
diff --git a/LICENSE b/LICENSE.txt
similarity index 99%
rename from LICENSE
rename to LICENSE.txt
index 0c8d1e2..9fea731 100644
--- a/LICENSE
+++ b/LICENSE.txt
@@ -26,3 +26,4 @@ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
diff --git a/perplexipy/__init__.py b/perplexipy/__init__.py
index 9446a28..60869d2 100644
--- a/perplexipy/__init__.py
+++ b/perplexipy/__init__.py
@@ -1,7 +1,7 @@
 # See: https://github.com/CIME-Software/perplexipy/blob/master/LICENSE.txt
 
 
-__VERSION__ = '0.0.4'
+__VERSION__ = '0.0.5'
 
 
 from dotenv import load_dotenv
@@ -13,14 +13,19 @@
 
 
 
-PERPLEXITY_DEFAULT_MODEL = 'mistral-7b-instruct'
-PERPLEXITY_DEFAULT_ROLE = 'user'
 """
 `PERPLEXITY_API_KEY` is set to the environment variable of the same name if
 present, otherwise it's set to the empty string `''`.
 """
 PERPLEXITY_API_KEY = os.environ.get('PERPLEXITY_API_KEY', default = '')
 PERPLEXITY_API_URL = 'https://api.perplexity.ai'
+"""
+The default model is **mistral-7b-instruct** because of it's efficiency and
+performance qualities.  Ref:  https://arxiv.org/abs/2310.06825
+"""
+PERPLEXITY_DEFAULT_MODEL = 'mistral-7b-instruct'
+PERPLEXITY_DEFAULT_ROLE = 'user'
+PERPLEXITY_TIMEOUT = 30.0 # seconds
 
 
 class PerplexityClientError(Exception):
@@ -31,6 +36,29 @@ def __init__(self, errorMessage):
         super().__init__(errorMessage)
 
 
+class Responses:
+    """
+    Encapsulates all the streaming responses from a query and enables access to
+    it using a purpose-built iterable.  Strips all the OpenAI response metadata
+    and returns only the textual response.  It's a streaming iterable,
+    open-ended.
+    """
+    def __init__(self, responsesStream):
+        self._responsesStream = responsesStream
+
+
+    def __iter__(self):
+        return self
+
+
+    def __next__(self):
+        try:
+            result = self._responsesStream.__next__().choices[0].delta.content
+        except Exception as e:
+            raise e
+        return result
+
+
 class PerplexityClient:
     """
     PerplexityClient objects encapsulate all the API functionality.  They can be
@@ -75,10 +103,14 @@ def __init__(self, key = PERPLEXITY_API_KEY, endpoint = PERPLEXITY_API_URL):
         self._key = key
         self._role = PERPLEXITY_DEFAULT_ROLE
         self.model = PERPLEXITY_DEFAULT_MODEL
-        self._client = OpenAI(api_key = self._key, base_url = self._endpoint)
+        self._client = OpenAI(
+            api_key = self._key,
+            base_url = self._endpoint,
+            timeout = PERPLEXITY_TIMEOUT,
+        )
 
 
-    def query(self, query):
+    def query(self, query: str) -> str:
         """
         Send a single message query to the service, receive a single response.
 
@@ -99,7 +131,13 @@ def query(self, query):
         whatever the API raised.
 
         In most cases, the API returns `openai.BadRequestError`.
+
+            PerplexityClientError
+        If the query is `None` or empty.
         """
+        if not query:
+            raise PerplexityClientError('query cannot be None or empty')
+
         messages = [ { 'role': self._role, 'content': query, }, ]
 
         response = self._client.chat.completions.create(
@@ -111,3 +149,84 @@ def query(self, query):
 
         return result
 
+
+    def queryBatch(self, query: str) -> tuple:
+        """
+        Send a single message query to the service, receive a single response.
+
+        Arguments
+        ---------
+            query
+        A string with the query in one of the model's supported languages.
+
+        Returns
+        -------
+        A tuple with a batch of 1 or more response strings.
+
+        Raises
+        ------
+            Exception
+        An `openai.BadRequestError` or similar if the query is malformed or has
+        something other than text data.  The error raised passes through
+        whatever the API raised.
+
+        In most cases, the API returns `openai.BadRequestError`.
+
+            PerplexityClientError
+        If the query is `None` or empty.
+        """
+        if not query:
+            raise PerplexityClientError('query cannot be None or empty')
+
+        messages = [ { 'role': self._role, 'content': query, }, ]
+
+        response = self._client.chat.completions.create(
+            model = self.model,
+            messages = messages,
+        )
+
+        result = tuple(choice.message.content for choice in response.choices)
+
+        return result
+
+
+    def queryStreamable(self, query: str) -> Responses:
+        """
+        Send a query and return a long, streamable response.
+
+        Arguments
+        ---------
+            query
+        A string with the query in one of the model's supported languages.
+
+        Returns
+        -------
+        Returns a `perplexipy.Responses` object for streaming the textual
+        responses from the client.  The response strips all metadata and returns
+        only content strings, for simpler and easier integration.
+
+        Raises
+        ------
+            Exception
+        An `openai.BadRequestError` or similar if the query is malformed or has
+        something other than text data.  The error raised passes through
+        whatever the API raised.
+
+        In most cases, the API returns `openai.BadRequestError`.
+
+            PerplexityClientError
+        If the query is `None` or empty.
+        """
+        if not query:
+            raise PerplexityClientError('query cannot be None or empty')
+
+        messages = [ { 'role': self._role, 'content': query, }, ]
+
+        response = self._client.chat.completions.create(
+            model = self.model,
+            messages = messages,
+            stream = True,
+        )
+
+        return Responses(response)
+
diff --git a/requirements.txt b/requirements.txt
index 9002c7f..cd7407f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,2 @@
-openai
+openai>=1.12.0
 python-dotenv


Test patch:
diff --git a/tests/perplexipy-test.py b/tests/perplexipy-test.py
index a79b254..3035b73 100644
--- a/tests/perplexipy-test.py
+++ b/tests/perplexipy-test.py
@@ -2,6 +2,7 @@
 
 from perplexipy import PerplexityClient
 from perplexipy import PerplexityClientError
+from perplexipy import Responses
 
 import pytest
 
@@ -10,6 +11,7 @@
 
 TEST_BOGUS_MODEL = 'bogus-llm-1b'
 TEST_QUERY = 'Brief answer to the ultimate question about life, the Universe, and everything?'
+TEST_QUERY_LONG = 'Give me a comprehensive list of US presidents.'
 
 
 # +++ globals +++
@@ -57,6 +59,28 @@ def test_PerplexityClient_bogusModel(testClient):
     result = testClient.query(TEST_QUERY)
     assert result
 
+
+def test_PerplexityClient_queryBatch(testClient):
+    result = testClient.queryBatch(TEST_QUERY)
+    assert isinstance(result, tuple)
+    assert isinstance(result[0], str)
+
+
+def test_PerplexityClient_queryStreamable(testClient):
+    t = ''
+    results = testClient.queryStreamable(TEST_QUERY_LONG)
+    assert isinstance(results, Responses)
+    for result in results:
+        assert isinstance(result, str)
+        t += result
+    assert len(t)
+
+    with pytest.raises(PerplexityClientError):
+        testClient.queryStreamable('')
+    with pytest.raises(PerplexityClientError):
+        testClient.queryStreamable(None)
+
+
 # _testClient = PerplexityClient()
-# test_PerplexityClient_bogusModel(_testClient)
+# test_PerplexityClient_queryStreamable(_testClient)