Merge pull request #3 from Kheiden/add-ollama

.env.template

@@ -0,0 +1,3 @@
+OPENAI_API_KEY=
+OLLAMA_HOST_URL=http://host.docker.internal:11434
+OLLAMA_MODEL=llama3.2

Dockerfile

@@ -0,0 +1,12 @@
+FROM python:3.12.0
+
+RUN apt-get update -y && apt-get upgrade -y
+RUN pip install --upgrade pip
+
+COPY requirements.txt /src/requirements.txt
+
+WORKDIR /src
+
+RUN pip install -r requirements.txt
+
+ENTRYPOINT ["python", "build.py"]

README.md

@@ -89,6 +89,11 @@ To get started:
 3. Make your changes.
 4. Submit a pull request.
 
+## Building
+1. Clone the repository.
+2. `cd` to the root directory.
+3. Run `docker-compose up --build`
+
 ## License
 SimpleMind is licensed under the MIT License.
 

build.py

@@ -0,0 +1,2 @@
+import simplemind
+print('Build successful.')

docker-compose.yaml

@@ -0,0 +1,10 @@
+services:
+  simplemind:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    volumes:
+      - ./simplemind:/src/simplemind
+      - ./build.py:/src/build.py
+    env_file:
+      - .env

requirements.txt

@@ -0,0 +1,5 @@
+instructor
+anthropic
+requests
+openai
+ollama

simplemind/conversation.py

@@ -0,0 +1,12 @@
+class Conversation:
+    def __init__(self, ai_client):
+        self.messages = []
+        self.ai_client = ai_client
+    
+    def say(self, message):
+        self.messages.append({'role': 'user', 'content': message})
+
+    def get_reply(self):
+        reply = self.ai_client.message(messages=self.messages)
+        self.messages.append({'role': 'system', 'content': reply.text})
+        return reply

simplemind/integrations/__init__.py

@@ -1,2 +1,3 @@
 from .anthropic import Anthropic
 from .openai import OpenAI
+from .ollama import Ollama

simplemind/integrations/ollama.py

@@ -0,0 +1,107 @@
+import os
+
+from ollama import Client as BaseOllama
+
+from .base import BaseClientProvider
+from ..models import AIResponse
+from ..conversation import Conversation
+
+TIMEOUT = 60
+
+class Ollama(BaseClientProvider):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.login()
+        self.conversation = []
+
+    def login(self):
+        """Initialize Ollama client, with Instructor enabled."""
+        if not os.environ.get('OLLAMA_HOST_URL'):
+           raise ValueError("Please set the OLLAMA_HOST_URL environment variable")
+        
+        if not os.environ.get('OLLAMA_MODEL'):
+            raise ValueError("Please set the OLLAMA_MODEL environment variable")
+        else:
+            self.model = os.environ.get('OLLAMA_MODEL')
+
+        self.client = BaseOllama(
+            timeout=TIMEOUT,
+            host=os.environ.get('OLLAMA_HOST_URL'))
+        assert self.test_connection()
+
+    @property
+    def available_models(self):
+        """Returns the available models from the Ollama client."""
+
+        def gen():
+            for model in self.client.list().get('models'):
+                yield model
+
+        return [g for g in gen()]
+
+    def test_connection(self):
+        """Test the connection to Ollama. Returns True if successful."""
+
+        return bool(len(self.available_models))
+
+    def generate_text(self, prompt, *, response_model=False, **kwargs):
+        use_instructor = bool(response_model)
+
+        client = self.instructor_client if use_instructor else self.client
+
+        # Parameters for the Ollama client.
+        params = {
+            "prompt": prompt,
+            "model": self.model,
+        }
+        params.update(kwargs)
+
+        if use_instructor:
+            params["response_model"] = response_model
+
+        # Make the request to Ollama.
+        completion = client.generate(**params)
+        if use_instructor:
+            return completion.model_dump()
+
+        else:
+            return AIResponse(
+                response=completion,
+                text=completion.get('response'),
+            )
+
+    def message(self, message=None, message_history=None, response_model=False, **kwargs):
+        """Generates a response from the Ollama client."""
+        use_instructor = bool(response_model)
+
+        client = self.instructor_client if use_instructor else self.client
+
+        # Parameters for the Ollama client.
+        all_messages = []
+        if message_history:
+            all_messages.extend(message_history)
+        if message:
+            all_messages.append({'role': 'user', 'content': message})
+        params = {
+            "messages": all_messages,
+            "model": self.model,
+        }
+        params.update(kwargs)
+
+        if use_instructor:
+            params["response_model"] = response_model
+
+        # Make the request to Ollama.
+        completion = client.chat(**params)
+        if use_instructor:
+            return completion.model_dump()
+
+        else:
+            return AIResponse(
+                response=completion,
+                text=completion.get('message').get('content'),
+            )
+
+    def start_conversation(self):
+        return Conversation(self)

t3.py

@@ -0,0 +1,28 @@
+import simplemind
+
+aiclient = simplemind.Ollama()
+
+print('Messaging client')
+message_response = aiclient.message(message="Once upon a time in a land far away...")
+print(message_response)
+
+print('Generating Text')
+generated_text = aiclient.generate_text(prompt="Once upon a time in a land far away...")
+print(generated_text)
+
+print('Initiating Conversation')
+conversation = aiclient.start_conversation()
+
+# Add a message to the conversation
+conversation.say("Please remember the number 42")
+
+# Get the AI's response
+reply = conversation.get_reply()
+print(reply)
+
+# Add another message to the conversation
+conversation.say("What number did I ask you to remember?")
+
+# Get the AI's response
+reply = conversation.get_reply()
+print(reply)