Add back Ollama integration

.envrc.template

@@ -1,3 +1,4 @@
 export OPENAI_API_KEY=""
 export ANTHROPIC_API_KEY=""
 export XAI_API_KEY=""
+export OLLAMA_HOST_URL=""

.gitignore

@@ -166,3 +166,4 @@ cython_debug/
 .env
 
 src/**
+requirements.txt

Dockerfile

@@ -9,4 +9,4 @@ WORKDIR /src
 
 RUN pip install -r requirements.txt
 
-ENTRYPOINT ["python", "build.py"]
+ENTRYPOINT ["python", "main.py"]

README.md

@@ -159,10 +159,16 @@ To get started:
 4. Submit a pull request.
 
 ## Building
+
 1. Clone the repository.
 2. `cd` to the root directory.
+3. Generate the requirements.txt file `python -m piptools compile pyproject.toml`
 3. Run `docker-compose up --build`
 
+Two containers will run in sequence:
+1) `simplemind` - Builds and runs the tests
+2) `simplemind-test` - Runs the tests again, if the 1st container suceeds.
+
 ## License
 SimpleMind is licensed under the Apache 2.0 License.
 

docker-compose.yaml

@@ -1,10 +1,22 @@
 services:
   simplemind:
+    depends_on:
+      simplemind-test:
+        condition: service_completed_successfully
     build:
       context: .
       dockerfile: Dockerfile
     volumes:
       - ./simplemind:/src/simplemind
-      - ./build.py:/src/build.py
+      - ./test_ollama.py:/src/main.py
+    env_file:
+      - .env
+  simplemind-test:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    volumes:
+      - ./simplemind:/src/simplemind
+      - ./test_ollama.py:/src/main.py
     env_file:
       - .env

pyproject.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 description = "An experimental client for AI providers that intends to replace LangChain and LangGraph for most common use cases."
 readme = "README.md"
 requires-python = ">=3.11"
-dependencies = ["pydantic", "pydantic-settings", "instructor", "openai", "anthropic"]
+dependencies = ["pydantic", "pydantic-settings", "instructor", "openai", "anthropic", "ollama"]
 
 [build-system]
 requires = ["hatchling"]

simplemind/providers/__init__.py

@@ -1,5 +1,6 @@
 from .openai import OpenAI
 from .anthropic import Anthropic
 from .xai import XAI
+from .ollama import Ollama
 
-providers = [OpenAI, Anthropic, XAI]
+providers = [OpenAI, Anthropic, XAI, Ollama]

simplemind/providers/ollama.py

@@ -0,0 +1,63 @@
+import ollama as ol
+
+from ..settings import settings
+
+PROVIDER_NAME = "ollama"
+DEFAULT_MODEL = "llama3.2"
+TIMEOUT = 60
+NOT_IMPLEMENTED_REASON = """
+# TODO: instructor does not natively support ollama.
+# Alternate python dependency may be required
+"""
+class Ollama:
+    __name__ = PROVIDER_NAME
+    DEFAULT_MODEL = DEFAULT_MODEL
+
+    def __init__(self, host_url: str = None):
+        self.host_url = host_url or settings.OLLAMA_HOST_URL
+
+    @property
+    def client(self):
+        """The raw Ollama client."""
+        return ol.Client(
+            timeout=TIMEOUT,
+            host=self.host_url)
+
+    @property
+    def structured_client(self):
+        """A client patched with Instructor."""
+        raise NotImplementedError(NOT_IMPLEMENTED_REASON)
+
+    def send_conversation(self, conversation: "Conversation"):
+        """Send a conversation to the Ollama API."""
+        from ..models import Message
+        messages = [
+            {"role": msg.role, "content": msg.text} for msg in conversation.messages
+        ]
+        response = self.client.chat(
+            model=conversation.llm_model or DEFAULT_MODEL, messages=messages
+        )
+        assistant_message = response.get("message")
+
+        # Create and return a properly formatted Message instance
+        return Message(
+            role="assistant",
+            text=assistant_message.get("content"),
+            raw=response,
+            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_provider=PROVIDER_NAME,
+        )
+
+    def structured_response(self, *args, **kwargs):
+        raise NotImplementedError(NOT_IMPLEMENTED_REASON)
+
+    def generate_text(self, prompt, *, llm_model):
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.chat(
+            messages=messages, model=llm_model
+        )
+
+        return response.get("message").get("content")

simplemind/settings.py

@@ -6,6 +6,7 @@ class Settings(BaseSettings):
     OPENAI_API_KEY: str = Field(..., env="OPENAI_API_KEY")
     ANTHROPIC_API_KEY: str = Field(..., env="ANTHROPIC_API_KEY")
     XAI_API_KEY: str = Field(..., env="XAI_API_KEY")
+    OLLAMA_HOST_URL: str = Field(..., env="OLLAMA_HOST_URL")
 
 
 settings = Settings()

test_ollama.py

@@ -0,0 +1,58 @@
+import unittest
+import simplemind as sm
+from pydantic import BaseModel
+
+class TestOllama(unittest.TestCase):
+
+  def test_generate_text(self):
+    result = sm.generate_text(prompt="What is the meaning of life?", llm_provider="ollama", llm_model="llama3.2")
+    self.assertIsNotNone(result)
+
+  def test_create_conversation(self):
+    conversation = sm.create_conversation(llm_provider="ollama", llm_model="llama3.2")
+    conversation.add_message("user", "Remember the number 42.")
+    result = conversation.send()
+    self.assertIsNotNone(result)
+    self.assertIsInstance(result, sm.models.Message)
+
+  def test_memory(self):
+    conversation = sm.create_conversation(llm_provider="ollama", llm_model="llama3.2")
+    class SimpleMemoryPlugin:
+      def __init__(self):
+          self.memories = [
+              "the earth has fictionally beeen destroyed.",
+              "the moon is made of cheese.",
+          ]
+
+      def yield_memories(self):
+          return (m for m in self.memories)
+
+      def send_hook(self, conversation: sm.Conversation):
+          for m in self.yield_memories():
+              conversation.add_message(role="system", text=m)
+    
+    conversation.add_plugin(SimpleMemoryPlugin())
+    conversation.add_message(
+        role="user",
+        text="Write a poem about the moon",
+    )
+    result = conversation.send()
+    self.assertIsNotNone(result)
+    self.assertIsInstance(result, sm.models.Message)
+
+  def test_structure_response(self):
+    class Poem(BaseModel):
+        title: str
+        content: str
+    with self.assertRaises(NotImplementedError):
+      data_obj = sm.generate_data(
+        prompt="Write a poem about love",
+        llm_provider="ollama",
+        llm_model="llama3.2",
+        response_model=Poem)
+      self.assertIsNotNone(data_obj)
+      self.assertIsInstance(data_obj, Poem)
+
+
+if __name__ == '__main__':
+  unittest.main()