ollama provider

pyproject.toml

@@ -33,6 +33,7 @@ azure = "exchange.providers.azure:AzureProvider"
 databricks = "exchange.providers.databricks:DatabricksProvider"
 anthropic = "exchange.providers.anthropic:AnthropicProvider"
 bedrock = "exchange.providers.bedrock:BedrockProvider"
+ollama = "exchange.providers.ollama:OllamaProvider"
 
 [project.entry-points."exchange.moderator"]
 passive = "exchange.moderators.passive:PassiveModerator"

src/exchange/providers/__init__.py

@@ -5,7 +5,9 @@
 from exchange.providers.base import Provider, Usage  # noqa
 from exchange.providers.databricks import DatabricksProvider  # noqa
 from exchange.providers.openai import OpenAiProvider  # noqa
+from exchange.providers.ollama import OllamaProvider  # noqa
 from exchange.providers.azure import AzureProvider  # noqa
+
 from exchange.utils import load_plugins
 
 

src/exchange/providers/ollama.py

@@ -0,0 +1,93 @@
+import os
+from typing import Any, Dict, List, Tuple, Type
+
+import httpx
+
+from exchange.message import Message
+from exchange.providers.base import Provider, Usage
+from exchange.providers.utils import (
+    messages_to_openai_spec,
+    openai_response_to_message,
+    openai_single_message_context_length_exceeded,
+    raise_for_status,
+    tools_to_openai_spec,
+)
+from exchange.tool import Tool
+
+OLLAMA_HOST = "http://localhost:11434/"
+
+#
+# NOTE: this is experimental, best used with 70B model or larger if you can. 
+# Example profile config to try:
+class OllamaProvider(Provider):
+    """Provides chat completions for models hosted by Ollama"""
+
+    """
+
+        ollama:
+          provider: ollama
+          processor: llama3.1
+          accelerator: llama3.1
+          moderator: passive
+          toolkits:
+          - name: developer
+            requires: {}
+    """    
+
+    def __init__(self, client: httpx.Client) -> None:
+        print('PLEASE NOTE: the ollama provider is experimental, use with care')
+        super().__init__()
+        self.client = client
+
+    @classmethod
+    def from_env(cls: Type["OllamaProvider"]) -> "OllamaProvider":
+        url = os.environ.get("OLLAMA_HOST", OLLAMA_HOST)
+        client = httpx.Client(
+            base_url=url,
+            headers={"Content-Type": "application/json"},
+            timeout=httpx.Timeout(60 * 10),
+        )
+        return cls(client)
+
+    @staticmethod
+    def get_usage(data: dict) -> Usage:
+        usage = data.get("usage", {})
+        input_tokens = usage.get("prompt_tokens", 0)
+        output_tokens = usage.get("completion_tokens", 0)
+        total_tokens = usage.get("total_tokens", input_tokens + output_tokens)
+
+        return Usage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+        )
+
+    def complete(
+        self,
+        model: str,
+        system: str,
+        messages: List[Message],
+        tools: Tuple[Tool],
+        **kwargs: Dict[str, Any],
+    ) -> Tuple[Message, Usage]:
+        payload = dict(
+            messages=[
+                {"role": "system", "content": system},
+                *messages_to_openai_spec(messages),
+            ],
+            model=model,
+            tools=tools_to_openai_spec(tools) if tools else [],
+            **kwargs,
+        )
+        payload = {k: v for k, v in payload.items() if v}
+        response = self.client.post("v1/chat/completions", json=payload)
+
+        # Check for context_length_exceeded error for single, long input message
+        if "error" in response.json() and len(messages) == 1:
+            openai_single_message_context_length_exceeded(response.json()["error"])
+
+        data = raise_for_status(response).json()
+
+        message = openai_response_to_message(data)
+        usage = self.get_usage(data)
+        return message, usage

tests/providers/test_ollama.py

@@ -0,0 +1,56 @@
+import os
+from unittest.mock import patch
+
+import pytest
+from exchange import Message, Text
+from exchange.providers.ollama import OllamaProvider
+
+
+@pytest.fixture
+@patch.dict(os.environ, {})
+def ollama_provider():
+    return OllamaProvider.from_env()
+
+
+@patch("httpx.Client.post")
+@patch("time.sleep", return_value=None)
+@patch("logging.warning")
+@patch("logging.error")
+def test_ollama_completion(mock_error, mock_warning, mock_sleep, mock_post, ollama_provider):
+    mock_response = {
+        "choices": [{"message": {"role": "assistant", "content": "Hello!"}}],
+    }
+
+    mock_post.return_value.json.return_value = mock_response
+
+    model = "llama2"
+    system = "You are a helpful assistant."
+    messages = [Message.user("Hello")]
+    tools = ()
+
+    reply_message, _ = ollama_provider.complete(model=model, system=system, messages=messages, tools=tools)
+
+    assert reply_message.content == [Text(text="Hello!")]
+    mock_post.assert_called_once_with(
+        "v1/chat/completions",
+        json={
+            "messages": [
+                {"role": "system", "content": system},
+                {"role": "user", "content": "Hello"},
+            ],
+            "model": model,
+        },
+    )
+
+
+@pytest.mark.integration
+def test_ollama_integration():
+    provider = OllamaProvider.from_env()
+    model = "llama2"  # specify a valid model
+    system = "You are a helpful assistant."
+    messages = [Message.user("Hello")]
+
+    reply = provider.complete(model=model, system=system, messages=messages, tools=None)
+
+    assert reply[0].content is not None
+    print("Completion content from Ollama:", reply[0].content)