ollama-client (#48)

Author: gauron99

python/ollama-client/README.md

@@ -0,0 +1,53 @@
+# Python Ollama-Client Function (HTTP)
+
+Welcome to your Python Ollama Client Function. It uses the [ollama](https://github.com/ollama/ollama)
+library.
+
+## The Function
+
+Your Function can be found in `function/func.py`. It handles HTTP requests in
+the `handle(self,scope,receive,send)` which is also the ASGI's handle signature
+(It's ASGI compatible). The only requests handled elsewhere are readiness and
+liveness checks -- `ready` and `alive` functions respectivelly.
+
+### What it does
+
+During initialization, we set a the Ollama's client with the correct server
+adress. That's it. Everything else happens in the `handle` function itself.
+
+`handle` function includes some error handling and simple http body extraction
+and subsequently it makes an API request to the ollama server using Ollama's
+`client.chat()` function.
+
+### Expected data
+
+Any `GET` request will simply echo the standard 'OK' string.
+
+`POST` request should be in json format and include `prompt` key. This is your
+prompt for the LLM. Additionally you can include `model` key which is the name
+of the model you want to use.
+
+Example of a curl command:
+
+```bash
+# use the default model
+curl localhost:11434 -d '{"prompt":"How to cook eggs properly?"}'
+
+# use different model
+curl localhost:11434 -d '{"prompt":"How to cook eggs properly?","model":"llama3.2:3b"}'
+```
+
+These values are simply extracted from the request and if provided it feeds them
+to the request for the LLM in a ollama complient way (see the construction of
+`self.client.chat()` function call).
+
+## Extra
+
+As per usual, the Function also contains a readiness and liveness checks
+implemented at the bottom of the Function class in their matching function names.
+The `start` and `stop` function are also available. See the function comments
+for more descriptive information.
+
+For more info about the Ollama library, please visit [ollama github page](https://github.com/ollama/ollama)
+
+For more info about Functions, see [the complete documentation]('https://github.com/knative/func/tree/main/docs')

python/ollama-client/function/__init__.py

@@ -0,0 +1 @@
+from .func import new

python/ollama-client/function/func.py

@@ -0,0 +1,124 @@
+# Function
+import logging
+from ollama import Client
+import json
+import os
+
+def new():
+    """ New is the only method that must be implemented by a Function.
+    The instance returned can be of any name.
+    """
+    return Function()
+
+# helper function for sending responses
+async def send_it(send,msg:str|None):
+    if msg == None:
+        msg = ""
+
+    await send({
+        'type': 'http.response.start',
+        'status': 200,
+        'headers': [
+            [b'content-type', b'text/plain'],
+        ],
+    })
+    await send({
+        'type': 'http.response.body',
+        'body': msg.encode(),
+    })
+
+class Function:
+    def __init__(self):
+        """ The init method is an optional method where initialization can be
+        performed. See the start method for a startup hook which includes
+        configuration.
+       """
+        self.client = Client(
+                # where your OLLAMA server is running
+                host=os.environ.get("OLLAMA_HOST","127.0.0.1:11434")
+                )
+
+    async def handle(self, scope, receive, send):
+        """ Handle all HTTP requests to this Function other than readiness
+        and liveness probes.
+
+        To communicate with the LLM following curl data is expected:
+            {
+                "prompt":"Your prompt for LLM",
+                "model": "Your preffered ollama-compatible model",
+            }
+
+            Note: Both of these have defaults, therefore you dont need to
+            provide them.
+
+            example: curl <host:port> -d '{"prompt":"What is philosophy exactly"}'
+             """
+        logging.info("OK: Request Received")
+
+        if scope["method"] == "GET":
+            await send_it(send,'OK')
+            return
+
+        # 1) extract the whole body from request
+        body = b''
+        more_body = True
+        while more_body:
+            message = await receive()
+            body += message.get('body', b'')
+            more_body = message.get('more_body', False)
+
+        # 2) decode the request and fetch info
+        data = json.loads(body.decode('utf-8'))
+        prompt = data.get('prompt','Who are you?')
+        model = data.get('model',"llama3.2:1b")
+
+        print(f"using model {model}")
+        # 3) make /api/chat request to the ollama server
+        response = self.client.chat(
+                # assign your model here
+                model=model,
+                messages=[
+                    {
+                    'role':'user',
+                    'content':prompt,
+                    },
+                ])
+
+        # 4) return the response to the calling client
+        await send_it(send,response.message.content)
+
+    def start(self, cfg):
+        """ start is an optional method which is called when a new Function
+        instance is started, such as when scaling up or during an update.
+        Provided is a dictionary containing all environmental configuration.
+        Args:
+            cfg (Dict[str, str]): A dictionary containing environmental config.
+                In most cases this will be a copy of os.environ, but it is
+                best practice to use this cfg dict instead of os.environ.
+        """
+        logging.info("Function starting")
+
+    def stop(self):
+        """ stop is an optional method which is called when a function is
+        stopped, such as when scaled down, updated, or manually canceled.  Stop
+        can block while performing function shutdown/cleanup operations.  The
+        process will eventually be killed if this method blocks beyond the
+        platform's configured maximum studown timeout.
+        """
+        logging.info("Function stopping")
+
+    def alive(self):
+        """ alive is an optional method for performing a deep check on your
+        Function's liveness.  If removed, the system will assume the function
+        is ready if the process is running. This is exposed by default at the
+        path /health/liveness.  The optional string return is a message.
+        """
+        return True, "Alive"
+
+    def ready(self):
+        """ ready is an optional method for performing a deep check on your
+        Function's readiness.  If removed, the system will assume the function
+        is ready if the process is running.  This is exposed by default at the
+        path /health/rediness.
+        """
+        return True, "Ready"

python/ollama-client/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "function"
+description = ""
+version = "0.1.0"
+requires-python = ">=3.9"
+readme = "README.md"
+license = "MIT"
+dependencies = [
+  "httpx",
+  "pytest",
+  "pytest-asyncio",
+  "ollama"
+]
+authors = [
+  { name="Your Name", email="you@example.com"},
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"
+

python/ollama-client/tests/test_func.py

@@ -0,0 +1,38 @@
+"""
+An example set of unit tests which confirm that the main handler (the
+callable function) returns 200 OK for a simple HTTP GET.
+"""
+import pytest
+from function import new
+
+
+@pytest.mark.asyncio
+async def test_function_handle():
+    f = new()  # Instantiate Function to Test
+
+    sent_ok = False
+    sent_headers = False
+    sent_body = False
+
+    # Mock Send
+    async def send(message):
+        nonlocal sent_ok
+        nonlocal sent_headers
+        nonlocal sent_body
+
+        if message.get('status') == 200:
+            sent_ok = True
+
+        if message.get('type') == 'http.response.start':
+            sent_headers = True
+
+        if message.get('type') == 'http.response.body':
+            sent_body = True
+
+    # Invoke the Function
+    await f.handle({}, {}, send)
+
+    # Assert send was called
+    assert sent_ok, "Function did not send a 200 OK"
+    assert sent_headers, "Function did not send headers"
+    assert sent_body, "Function did not send a body"