diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..bf26e2e --- /dev/null +++ b/Dockerfile @@ -0,0 +1,21 @@ +# # Use the official Python 3.10 image as the base image +FROM python:3.11 + +# Set the working directory +WORKDIR /ollama-python + +# Install Poetry +RUN pip install --no-cache-dir poetry + +# Copy the project files into the working directory +COPY . . + +# Install project dependencies using Poetry +RUN poetry config virtualenvs.create false \ + && poetry install --only main + +# Install pre-commit and set up the hooks +RUN poetry run pre-commit install +# Specify the command to run your application +CMD ["sh", "-c", "while :; do sleep 10; done"] + diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..34c254d --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,11 @@ +version: '3.8' + +services: + ollama-python: + build: + context: . + dockerfile: Dockerfile + volumes: + - .:/ollama-python + stdin_open: true # Keep container running + tty: true diff --git a/ollama_python/endpoints/async_base_api.py b/ollama_python/endpoints/async_base_api.py new file mode 100644 index 0000000..e4e113b --- /dev/null +++ b/ollama_python/endpoints/async_base_api.py @@ -0,0 +1,63 @@ +from ollama_python.endpoints.base import BaseAPI +import aiohttp +import json +from typing import Optional, Callable, AsyncGenerator + + +class AsyncBaseApi(BaseAPI): + def __init__(self, base_url: str = "http://localhost:11434/api"): + super().__init__(base_url) + self.session = aiohttp.ClientSession() + + async def _stream( + self, endpoint: str, parameters: dict, return_type: Optional[Callable] = None + ) -> AsyncGenerator: + """ + Stream the response from the given endpoint + :param endpoint: The endpoint to stream from + :param parameters: The parameters to send + :return: A generator that yields the response + """ + async with self.session.post( + f"{self.base_url}/{endpoint}", + json=parameters, + stream=True, + raise_for_status=True, + ) as session: + async for line in session.iter_lines(): + if line: + resp = json.loads(line) + yield return_type(**resp) if return_type else resp + + async def _post( + self, + endpoint: str, + parameters: Optional[dict] = None, + return_type: Optional[Callable] = None, + ): + """ + Send a POST request to the given endpoint + :param endpoint: + :param parameters: + :param return_type: + :return: + """ + async with self.session.post( + f"{self.base_url}/{endpoint}", json=parameters, raise_for_status=True + ) as session: + data = await session.json() + + return return_type(**data) if return_type else session.status + + async def _get(self, endpoint: str, return_type: Optional[Callable] = None): + """ + Send a GET request to the given endpoint + :param endpoint: + :param return_type: + :return: + """ + async with self.session.get( + f"{self.base_url}/{endpoint}", raise_for_status=True + ) as session: + data = await session.json() + return return_type(**data) if return_type else session.status