diff --git a/simplemind/conversation.py b/simplemind/conversation.py new file mode 100644 index 0000000..cb07fbf --- /dev/null +++ b/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 \ No newline at end of file diff --git a/simplemind/integrations/__init__.py b/simplemind/integrations/__init__.py index 6942e5c..b69523f 100644 --- a/simplemind/integrations/__init__.py +++ b/simplemind/integrations/__init__.py @@ -1,2 +1,3 @@ from .anthropic import Anthropic from .openai import OpenAI +from .ollama import Ollama \ No newline at end of file diff --git a/simplemind/integrations/ollama.py b/simplemind/integrations/ollama.py new file mode 100644 index 0000000..6bb1d7e --- /dev/null +++ b/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 OpenAI 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 OpenAI 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)