seekrai 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

seekrai/__init__.py

@@ -23,30 +23,15 @@ log: str | None = None  # Set to either 'debug' or 'info', controls console logg
 
 if TYPE_CHECKING:
     import requests
-    from aiohttp import ClientSession
 
 requestssession: "requests.Session" | Callable[[], "requests.Session"] | None = None
 
-aiosession: ContextVar["ClientSession" | None] = ContextVar(
-    "aiohttp-session", default=None
-)
-
 from seekrai.client import AsyncClient, AsyncSeekrFlow, Client, SeekrFlow
 
 
 api_key: str | None = None  # To be deprecated in the next major release
 
-# Legacy functions
-from seekrai.legacy.complete import AsyncComplete, Complete, Completion
-from seekrai.legacy.embeddings import Embeddings
-from seekrai.legacy.files import Files
-from seekrai.legacy.finetune import Finetune
-from seekrai.legacy.images import Image
-from seekrai.legacy.models import Models
-
-
 __all__ = [
-    "aiosession",
     "constants",
     "version",
     "SeekrFlow",

seekrai/abstract/api_requestor.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import email.utils
 import json
 import sys
@@ -10,34 +9,30 @@ from json import JSONDecodeError
 from random import random
 from typing import (
     Any,
-    AsyncContextManager,
     AsyncGenerator,
+    AsyncIterator,
     Dict,
     Iterator,
+    Mapping,
     Tuple,
     overload,
 )
 from urllib.parse import urlencode, urlsplit, urlunsplit
 
-import aiohttp
-import requests
-from tqdm.utils import CallbackIOWrapper
-
 
 if sys.version_info >= (3, 8):
     from typing import Literal
 else:
     from typing_extensions import Literal
 
-import seekrai
+import httpx
+
 from seekrai import error, utils
 from seekrai.constants import (
     BASE_URL,
     INITIAL_RETRY_DELAY,
-    MAX_CONNECTION_RETRIES,
     MAX_RETRIES,
     MAX_RETRY_DELAY,
-    MAX_SESSION_LIFETIME_SECS,
     TIMEOUT_SECS,
 )
 from seekrai.seekrflow_response import SeekrFlowResponse
@@ -58,43 +53,32 @@ def _build_api_url(url: str, query: str) -> str:
     return str(urlunsplit((scheme, netloc, path, query, fragment)))
 
 
-def _make_session(max_retries: int | None = None) -> requests.Session:
-    if seekrai.requestssession:
-        if isinstance(seekrai.requestssession, requests.Session):
-            return seekrai.requestssession
-        return seekrai.requestssession()
-    s = requests.Session()
-    s.mount(
-        "https://",
-        requests.adapters.HTTPAdapter(max_retries=max_retries),
-    )
-    return s
-
-
-def parse_stream_helper(line: bytes) -> str | None:
-    if line and line.startswith(b"data:"):
-        if line.startswith(b"data: "):
+def parse_stream_helper(line: str) -> str | None:
+    if line and line.startswith("data:"):
+        if line.startswith("data: "):
             # SSE event may be valid when it contains whitespace
-            line = line[len(b"data: ") :]
+            line = line[len("data: ") :]
         else:
-            line = line[len(b"data:") :]
-        if line.strip() == b"[DONE]":
+            line = line[len("data:") :]
+        if line.strip() == "[DONE]":
             # return here will cause GeneratorExit exception in urllib3
             # and it will close http connection with TCP Reset
             return None
         else:
-            return line.decode("utf-8")
+            return line
     return None
 
 
-def parse_stream(rbody: Iterator[bytes]) -> Iterator[str]:
+def parse_stream(rbody: Iterator[str]) -> Iterator[str]:
     for line in rbody:
         _line = parse_stream_helper(line)
         if _line is not None:
             yield _line
 
 
-async def parse_stream_async(rbody: aiohttp.StreamReader) -> AsyncGenerator[str, Any]:
+async def parse_stream_async(
+    rbody: AsyncIterator[str],
+) -> AsyncGenerator[str, Any]:
     async for line in rbody:
         _line = parse_stream_helper(line)
         if _line is not None:
@@ -169,42 +153,12 @@ class APIRequestor:
         timeout = sleep_seconds * jitter
         return timeout if timeout >= 0 else 0
 
-    def _retry_request(
-        self,
-        options: SeekrFlowRequest,
-        remaining_retries: int,
-        response_headers: Dict[str, Any] | None,
-        *,
-        stream: bool,
-        request_timeout: float | Tuple[float, float] | None = None,
-    ) -> requests.Response:
-        remaining = remaining_retries - 1
-        if remaining == 1:
-            utils.log_debug("1 retry left")
-        else:
-            utils.log_debug(f"{remaining} retries left")
-
-        timeout = self._calculate_retry_timeout(remaining, response_headers)
-        ("Retrying request to %s in %f seconds", options.url, timeout)
-
-        # In a synchronous context we are blocking the entire thread. Up to the library user to run the client in a
-        # different thread if necessary.
-        time.sleep(timeout)
-
-        return self.request_raw(
-            options=options,
-            stream=stream,
-            request_timeout=request_timeout,
-            remaining_retries=remaining,
-        )
-
     @overload
     def request(
         self,
         options: SeekrFlowRequest,
         stream: Literal[True],
-        remaining_retries: int | None = ...,
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[Iterator[SeekrFlowResponse], bool, str]:
         pass
 
@@ -213,8 +167,7 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         stream: Literal[False] = ...,
-        remaining_retries: int | None = ...,
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[SeekrFlowResponse, bool, str]:
         pass
 
@@ -223,8 +176,7 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         stream: bool = ...,
-        remaining_retries: int | None = ...,
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[SeekrFlowResponse | Iterator[SeekrFlowResponse], bool, str]:
         pass
 
@@ -232,29 +184,22 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         stream: bool = False,
-        remaining_retries: int | None = None,
-        request_timeout: float | Tuple[float, float] | None = None,
+        request_timeout: float | None = None,
     ) -> Tuple[
         SeekrFlowResponse | Iterator[SeekrFlowResponse],
         bool,
         str | None,
     ]:
-        result = self.request_raw(
-            options=options,
-            remaining_retries=remaining_retries or self.retries,
-            stream=stream,
-            request_timeout=request_timeout,
-        )
-
-        resp, got_stream = self._interpret_response(result, stream)
-        return resp, got_stream, self.api_key
+        result = self.request_raw(options, stream, request_timeout)
+        resp = self._interpret_response(result, stream)
+        return resp, stream, self.api_key
 
     @overload
     async def arequest(
         self,
         options: SeekrFlowRequest,
         stream: Literal[True],
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[AsyncGenerator[SeekrFlowResponse, None], bool, str]:
         pass
 
@@ -264,7 +209,7 @@ class APIRequestor:
         options: SeekrFlowRequest,
         *,
         stream: Literal[True],
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[AsyncGenerator[SeekrFlowResponse, None], bool, str]:
         pass
 
@@ -273,7 +218,7 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         stream: Literal[False] = ...,
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[SeekrFlowResponse, bool, str]:
         pass
 
@@ -282,7 +227,7 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         stream: bool = ...,
-        request_timeout: float | Tuple[float, float] | None = ...,
+        request_timeout: float | None = ...,
     ) -> Tuple[SeekrFlowResponse | AsyncGenerator[SeekrFlowResponse, None], bool, str]:
         pass
 
@@ -290,43 +235,49 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         stream: bool = False,
-        request_timeout: float | Tuple[float, float] | None = None,
-    ) -> Tuple[SeekrFlowResponse | AsyncGenerator[SeekrFlowResponse, None], bool, str]:
-        ctx = AioHTTPSession()
-        session = await ctx.__aenter__()
-        result = None
+        request_timeout: float | None = None,
+    ) -> Tuple[
+        SeekrFlowResponse | AsyncGenerator[SeekrFlowResponse, None], bool, str | None
+    ]:
+        abs_url, headers, data = self._prepare_request_raw(options, False)
+        client = httpx.AsyncClient(http2=True)
+        req = client.build_request(
+            options.method,
+            abs_url,
+            headers=headers,
+            data=data,  # type: ignore
+            files=options.files,
+            timeout=request_timeout or self.timeout,
+        )
         try:
-            result = await self.arequest_raw(
-                options,
-                session,
-                request_timeout=request_timeout,
+            result = await client.send(req, stream=stream)
+        except httpx.TimeoutException as e:
+            utils.log_debug("Encountered httpx.TimeoutException")
+
+            raise error.Timeout("Request timed out: {}".format(e)) from e
+        except httpx.RequestError as e:
+            utils.log_debug("Encountered httpx.RequestError")
+
+            raise error.APIConnectionError(
+                "Error communicating with API: {}".format(e)
+            ) from e
+
+        # retry on 5XX error or rate-limit
+        if 500 <= result.status_code < 600 or result.status_code == 429:
+            utils.log_debug(
+                f"Encountered httpx.HTTPError. Error code: {result.status_code}"
             )
-            resp, got_stream = await self._interpret_async_response(result, stream)
-        except Exception:
-            # Close the request before exiting session context.
-            if result is not None:
-                result.release()
-            await ctx.__aexit__(None, None, None)
-            raise
-        if got_stream:
-
-            async def wrap_resp() -> AsyncGenerator[SeekrFlowResponse, None]:
-                assert isinstance(resp, AsyncGenerator)
-                try:
-                    async for r in resp:
-                        yield r
-                finally:
-                    # Close the request before exiting session context. Important to do it here
-                    # as if stream is not fully exhausted, we need to close the request nevertheless.
-                    result.release()
-                    await ctx.__aexit__(None, None, None)
-
-            return wrap_resp(), got_stream, self.api_key  # type: ignore
-        else:
-            # Close the request before exiting session context.
-            result.release()
-            await ctx.__aexit__(None, None, None)
-            return resp, got_stream, self.api_key  # type: ignore
+
+        utils.log_debug(
+            "SeekrFlow API response",
+            path=abs_url,
+            response_code=result.status_code,
+            processing_ms=result.headers.get("x-total-time"),
+            request_id=result.headers.get("CF-RAY"),
+        )
+        resp = await self._interpret_async_response(result, stream)
+
+        return resp, stream, self.api_key
 
     @classmethod
     def handle_error_response(
@@ -425,12 +376,11 @@ class APIRequestor:
         self,
         options: SeekrFlowRequest,
         absolute: bool = False,
-    ) -> Tuple[str, Dict[str, str], Dict[str, str] | CallbackIOWrapper | bytes | None]:
+    ) -> Tuple[str, Dict[str, str], Mapping[str, Any] | None | str]:
         abs_url = options.url if absolute else "%s%s" % (self.api_base, options.url)
         headers = self._validate_headers(options.headers or self.supplied_headers)
 
-        data = None
-        data_bytes = None
+        data: Mapping[str, Any] | None | str = None
         if options.method.lower() == "get" or options.method.lower() == "delete":
             if options.params:
                 encoded_params = urlencode(
@@ -438,11 +388,9 @@ class APIRequestor:
                 )
                 abs_url = _build_api_url(abs_url, encoded_params)
         elif options.method.lower() in {"post", "put"}:
-            if options.params and (options.files or options.override_headers):
-                data = options.params
-            elif options.params and not options.files:
-                data_bytes = json.dumps(options.params).encode()
-                headers["Content-Type"] = "application/json"
+            data = options.params
+            if options.params and not options.files:
+                data = json.dumps(data)
 
         else:
             raise error.APIConnectionError(
@@ -458,87 +406,44 @@ class APIRequestor:
             "Request to SeekrFlow API",
             method=options.method,
             path=abs_url,
-            post_data=(data or data_bytes),
+            post_data=data,
             headers=json.dumps(headers),
         )
 
-        return abs_url, headers, (data or data_bytes)
+        return abs_url, headers, data
 
     def request_raw(
         self,
         options: SeekrFlowRequest,
-        remaining_retries: int,
-        *,
         stream: bool = False,
-        request_timeout: float | Tuple[float, float] | None = None,
-        absolute: bool = False,
-    ) -> requests.Response:
-        abs_url, headers, data = self._prepare_request_raw(options, absolute)
-        if not hasattr(_thread_context, "session"):
-            _thread_context.session = _make_session(MAX_CONNECTION_RETRIES)
-            _thread_context.session_create_time = time.time()
-        elif (
-            time.time() - getattr(_thread_context, "session_create_time", 0)
-            >= MAX_SESSION_LIFETIME_SECS
-        ):
-            _thread_context.session.close()
-            _thread_context.session = _make_session(MAX_CONNECTION_RETRIES)
-            _thread_context.session_create_time = time.time()
+        request_timeout: float | None = None,
+    ) -> httpx.Response:
+        abs_url, headers, data = self._prepare_request_raw(options, False)
+        client = httpx.Client(http2=True)
+        req = client.build_request(
+            options.method,
+            abs_url,
+            headers=headers,
+            data=data,  # type: ignore
+            files=options.files,
+            timeout=request_timeout or self.timeout,
+        )
         try:
-            result = _thread_context.session.request(
-                options.method,
-                abs_url,
-                headers=headers,
-                data=data,
-                files=options.files,
-                stream=stream,
-                timeout=request_timeout or self.timeout,
-                proxies=_thread_context.session.proxies,
-                allow_redirects=options.allow_redirects,
-            )
-        except requests.exceptions.Timeout as e:
-            utils.log_debug("Encountered requests.exceptions.Timeout")
-
-            if remaining_retries > 0:
-                return self._retry_request(
-                    options,
-                    remaining_retries=remaining_retries,
-                    response_headers=dict(result.headers),
-                    stream=stream,
-                    request_timeout=request_timeout,
-                )
+            result = client.send(req, stream=stream)
+        except httpx.TimeoutException as e:
+            utils.log_debug("Encountered httpx.TimeoutException")
 
             raise error.Timeout("Request timed out: {}".format(e)) from e
-        except requests.exceptions.RequestException as e:
-            utils.log_debug("Encountered requests.exceptions.RequestException")
-
-            if remaining_retries > 0:
-                return self._retry_request(
-                    options,
-                    remaining_retries=remaining_retries,
-                    response_headers=dict(result.headers),
-                    stream=stream,
-                    request_timeout=request_timeout,
-                )
+        except httpx.RequestError as e:
+            utils.log_debug("Encountered httpx.HTTPError")
 
             raise error.APIConnectionError(
                 "Error communicating with API: {}".format(e)
             ) from e
 
         # retry on 5XX error or rate-limit
-        if 500 <= result.status_code < 600 or result.status_code == 429:
-            utils.log_debug(
-                f"Encountered requests.exceptions.HTTPError. Error code: {result.status_code}"
-            )
-
-            if remaining_retries > 0:
-                return self._retry_request(
-                    options,
-                    remaining_retries=remaining_retries,
-                    response_headers=dict(result.headers),
-                    stream=stream,
-                    request_timeout=request_timeout,
-                )
+        if result.status_code > 300:
+            raise httpx.HTTPError("Error communicating with API: {}".format(result))
 
         utils.log_debug(
             "SeekrFlow API response",
@@ -547,113 +452,53 @@ class APIRequestor:
             processing_ms=result.headers.get("x-total-time"),
             request_id=result.headers.get("CF-RAY"),
         )
-
-        return result  # type: ignore
-
-    async def arequest_raw(
-        self,
-        options: SeekrFlowRequest,
-        session: aiohttp.ClientSession,
-        *,
-        request_timeout: float | Tuple[float, float] | None = None,
-        absolute: bool = False,
-    ) -> aiohttp.ClientResponse:
-        abs_url, headers, data = self._prepare_request_raw(options, absolute)
-
-        if isinstance(request_timeout, tuple):
-            timeout = aiohttp.ClientTimeout(
-                connect=request_timeout[0],
-                total=request_timeout[1],
-            )
-        else:
-            timeout = aiohttp.ClientTimeout(total=request_timeout or self.timeout)
-
-        if options.files:
-            data, content_type = requests.models.RequestEncodingMixin._encode_files(  # type: ignore
-                options.files, data
-            )
-            headers["Content-Type"] = content_type
-
-        request_kwargs = {
-            "headers": headers,
-            "data": data,
-            "timeout": timeout,
-            "allow_redirects": options.allow_redirects,
-        }
-
-        try:
-            result = await session.request(
-                method=options.method, url=abs_url, **request_kwargs
-            )
-            utils.log_debug(
-                "SeekrFlow API response",
-                path=abs_url,
-                response_code=result.status,
-                processing_ms=result.headers.get("x-total-time"),
-                request_id=result.headers.get("CF-RAY"),
-            )
-            # Don't read the whole stream for debug logging unless necessary.
-            if seekrai.log == "debug":
-                utils.log_debug(
-                    "API response body", body=result.content, headers=result.headers
-                )
-            return result
-        except (aiohttp.ServerTimeoutError, asyncio.TimeoutError) as e:
-            raise error.Timeout("Request timed out") from e
-        except aiohttp.ClientError as e:
-            raise error.APIConnectionError("Error communicating with SeekrFlow") from e
+        return result
 
     def _interpret_response(
-        self, result: requests.Response, stream: bool
-    ) -> Tuple[SeekrFlowResponse | Iterator[SeekrFlowResponse], bool]:
+        self, result: httpx.Response, stream: bool
+    ) -> SeekrFlowResponse | Iterator[SeekrFlowResponse]:
         """Returns the response(s) and a bool indicating whether it is a stream."""
+
         if stream and "text/event-stream" in result.headers.get("Content-Type", ""):
-            return (
+            iterator = (
                 self._interpret_response_line(
                     line, result.status_code, result.headers, stream=True
                 )
-                for line in parse_stream(result.iter_lines())
-            ), True
+                for line in parse_stream(result.iter_text())
+            )
+            return iterator
         else:
-            return (
-                self._interpret_response_line(
-                    result.content.decode("utf-8"),
-                    result.status_code,
-                    result.headers,
-                    stream=False,
-                ),
-                False,
+            return self._interpret_response_line(
+                result.content.decode("utf-8"),
+                result.status_code,
+                result.headers,
+                stream=False,
             )
 
     async def _interpret_async_response(
-        self, result: aiohttp.ClientResponse, stream: bool
-    ) -> (
-        tuple[AsyncGenerator[SeekrFlowResponse, None], bool]
-        | tuple[SeekrFlowResponse, bool]
-    ):
+        self, result: httpx.Response, stream: bool
+    ) -> AsyncGenerator[SeekrFlowResponse, None] | SeekrFlowResponse:
         """Returns the response(s) and a bool indicating whether it is a stream."""
         if stream and "text/event-stream" in result.headers.get("Content-Type", ""):
-            return (
+            iterator = (
                 self._interpret_response_line(
-                    line, result.status, result.headers, stream=True
+                    line, result.status_code, result.headers, stream=True
                 )
-                async for line in parse_stream_async(result.content)
-            ), True
+                async for line in parse_stream_async(result.aiter_text())
+            )
+            return iterator
         else:
             try:
-                await result.read()
-            except (aiohttp.ServerTimeoutError, asyncio.TimeoutError) as e:
+                result.read()
+            except httpx.TimeoutException as e:
                 raise error.Timeout("Request timed out") from e
-            except aiohttp.ClientError as e:
+            except httpx.HTTPError as e:
                 utils.log_warn(e, body=result.content)
-            return (
-                self._interpret_response_line(
-                    (await result.read()).decode("utf-8"),
-                    result.status,
-                    result.headers,
-                    stream=False,
-                ),
-                False,
+            return self._interpret_response_line(
+                (result.read()).decode("utf-8"),
+                result.status_code,
+                result.headers,
+                stream=False,
             )
 
     def _interpret_response_line(
@@ -687,24 +532,3 @@ class APIRequestor:
         if not 200 <= rcode < 300:
             raise self.handle_error_response(resp, rcode, stream_error=stream)
         return resp
-
-
-class AioHTTPSession(AsyncContextManager[aiohttp.ClientSession]):
-    def __init__(self) -> None:
-        self._session: aiohttp.ClientSession | None = None
-        self._should_close_session: bool = False
-
-    async def __aenter__(self) -> aiohttp.ClientSession:
-        self._session = seekrai.aiosession.get()
-        if self._session is None:
-            self._session = await aiohttp.ClientSession().__aenter__()
-            self._should_close_session = True
-
-        return self._session
-
-    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        if self._session is None:
-            raise RuntimeError("Session is not initialized")
-
-        if self._should_close_session:
-            await self._session.__aexit__(exc_type, exc_value, traceback)