From 0e0c992f5940e9df683ab6fd83342bded2cf33b4 Mon Sep 17 00:00:00 2001
From: Moishe Lettvin <moishel@gmail.com>
Date: Sat, 10 Feb 2024 09:22:52 -0500
Subject: [PATCH] Ollama LLM service

---
 src/dailyai/services/ollama_ai_services.py | 42 ++++++++++++++++++++++
 src/dailyai/services/open_ai_services.py   |  5 +--
 2 files changed, 43 insertions(+), 4 deletions(-)
 create mode 100644 src/dailyai/services/ollama_ai_services.py

diff --git a/src/dailyai/services/ollama_ai_services.py b/src/dailyai/services/ollama_ai_services.py
new file mode 100644
index 000000000..c065e0b7b
--- /dev/null
+++ b/src/dailyai/services/ollama_ai_services.py
@@ -0,0 +1,42 @@
+from openai import AsyncOpenAI
+
+import json
+from collections.abc import AsyncGenerator
+
+from dailyai.services.ai_services import LLMService
+
+
+class OLLamaLLMService(LLMService):
+    def __init__(self, model="llama2", base_url='http://localhost:11434/v1'):
+        super().__init__()
+        self._model = model
+        self._client = AsyncOpenAI(api_key="ollama", base_url=base_url)
+
+    async def get_response(self, messages, stream):
+        return await self._client.chat.completions.create(
+            stream=stream,
+            messages=messages,
+            model=self._model
+        )
+
+    async def run_llm_async(self, messages) -> AsyncGenerator[str, None]:
+        messages_for_log = json.dumps(messages)
+        self.logger.debug(f"Generating chat via openai: {messages_for_log}")
+
+        chunks = await self._client.chat.completions.create(model=self._model, stream=True, messages=messages)
+        async for chunk in chunks:
+            if len(chunk.choices) == 0:
+                continue
+
+            if chunk.choices[0].delta.content:
+                yield chunk.choices[0].delta.content
+
+    async def run_llm(self, messages) -> str | None:
+        messages_for_log = json.dumps(messages)
+        self.logger.debug(f"Generating chat via openai: {messages_for_log}")
+
+        response = await self._client.chat.completions.create(model=self._model, stream=False, messages=messages)
+        if response and len(response.choices) > 0:
+            return response.choices[0].message.content
+        else:
+            return None
diff --git a/src/dailyai/services/open_ai_services.py b/src/dailyai/services/open_ai_services.py
index 7892af06b..2aea60fa9 100644
--- a/src/dailyai/services/open_ai_services.py
+++ b/src/dailyai/services/open_ai_services.py
@@ -1,15 +1,12 @@
-import requests
 import aiohttp
-import asyncio
 from PIL import Image
 import io
 from openai import AsyncOpenAI
 
-import os
 import json
 from collections.abc import AsyncGenerator
 
-from dailyai.services.ai_services import AIService, TTSService, LLMService, ImageGenService
+from dailyai.services.ai_services import LLMService, ImageGenService
 
 
 class OpenAILLMService(LLMService):