From 23389c3a6217b70d1902f2be4ed27e525096ed98 Mon Sep 17 00:00:00 2001 From: Kurt Heiden Date: Mon, 28 Oct 2024 21:22:15 -0600 Subject: [PATCH] Add back Ollama integration --- .envrc.template | 1 + .gitignore | 1 + Dockerfile | 2 +- README.md | 6 +++ docker-compose.yaml | 14 ++++++- pyproject.toml | 2 +- simplemind/providers/__init__.py | 3 +- simplemind/providers/ollama.py | 63 ++++++++++++++++++++++++++++++++ simplemind/settings.py | 1 + test_ollama.py | 58 +++++++++++++++++++++++++++++ 10 files changed, 147 insertions(+), 4 deletions(-) create mode 100644 simplemind/providers/ollama.py create mode 100644 test_ollama.py diff --git a/.envrc.template b/.envrc.template index 8438939..07f7d8f 100644 --- a/.envrc.template +++ b/.envrc.template @@ -1,3 +1,4 @@ export OPENAI_API_KEY="" export ANTHROPIC_API_KEY="" export XAI_API_KEY="" +export OLLAMA_HOST_URL="" diff --git a/.gitignore b/.gitignore index b66e3ef..e43c919 100644 --- a/.gitignore +++ b/.gitignore @@ -166,3 +166,4 @@ cython_debug/ .env src/** +requirements.txt diff --git a/Dockerfile b/Dockerfile index 5d0954e..7599519 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,4 +9,4 @@ WORKDIR /src RUN pip install -r requirements.txt -ENTRYPOINT ["python", "build.py"] \ No newline at end of file +ENTRYPOINT ["python", "main.py"] \ No newline at end of file diff --git a/README.md b/README.md index 95d75d9..adebc02 100644 --- a/README.md +++ b/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. diff --git a/docker-compose.yaml b/docker-compose.yaml index d756f99..85d363c 100644 --- a/docker-compose.yaml +++ b/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 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1c4cebc..e71e648 100644 --- a/pyproject.toml +++ b/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"] diff --git a/simplemind/providers/__init__.py b/simplemind/providers/__init__.py index 87787ea..c3112c6 100644 --- a/simplemind/providers/__init__.py +++ b/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] diff --git a/simplemind/providers/ollama.py b/simplemind/providers/ollama.py new file mode 100644 index 0000000..de1bfdc --- /dev/null +++ b/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") diff --git a/simplemind/settings.py b/simplemind/settings.py index 7828508..8b8e6f3 100644 --- a/simplemind/settings.py +++ b/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() diff --git a/test_ollama.py b/test_ollama.py new file mode 100644 index 0000000..4d40d38 --- /dev/null +++ b/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() \ No newline at end of file