ollama: Support multiple ollama-connector instances

gbenson · Feb 2, 2025 · 303ebbe · 303ebbe
1 parent e92372e
commit 303ebbe
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 2 deletions.
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -14,6 +14,9 @@ include:
       - $HOME/.config/hive/vane-webui.env
 
   - path: services/ollama/docker-compose.yml
+    env_file:
+      - $HOME/.config/hive/ollama.env
+
   - path: services/service-monitor/docker-compose.yml
   - path: services/vane-webui-api/docker-compose.yml
 

diff --git a/services/ollama/docker-compose.yml b/services/ollama/docker-compose.yml
@@ -8,6 +8,8 @@ services:
       - ollama
     secrets:
       - rabbitmq.env
+    command:
+      - --request-queue-prefix=${OLLAMA_REQUEST_QUEUE_PREFIX}
 
   ollama:
     image: ollama/ollama

diff --git a/services/ollama/hive/ollama/service.py b/services/ollama/hive/ollama/service.py
@@ -9,6 +9,7 @@
 
 import requests
 
+from hive.common import ArgumentParser
 from hive.messaging import Channel, Message
 from hive.service import HiveService
 
@@ -18,8 +19,44 @@
 
 @dataclass
 class Service(HiveService):
-    ollama_api_url: str = "http://ollama:11434"
-    request_queue: str = "ollama.api.requests"
+    DEFAULT_API_URL: ClassVar[str] = "http://ollama:11434"
+    ollama_api_url: Optional[str] = None
+
+    DEFAULT_QUEUE: ClassVar[str] = "ollama.api.requests"
+    request_queue: Optional[str] = None
+
+    def make_argument_parser(self) -> ArgumentParser:
+        parser = super().make_argument_parser()
+        parser.add_argument(
+            "--ollama-api-url",
+            metavar="URL",
+            default=self.DEFAULT_API_URL,
+            help=(f"URL to proxy requests to"
+                  f" [default: {self.DEFAULT_API_URL}]"),
+        )
+        parser.add_argument(
+            "--request-queue",
+            metavar="QUEUE",
+            default=self.DEFAULT_QUEUE,
+            help=(f"Queue to consume requests from"
+                  f" [default: {self.DEFAULT_QUEUE}]"),
+        )
+        parser.add_argument(
+            "--request-queue-prefix",
+            metavar="PREFIX",
+            help=("Prefix to prepend to request queue name"
+                  " [default: no prefix]"),
+        )
+        return parser
+
+    def __post_init__(self):
+        super().__post_init__()
+        if not self.ollama_api_url:
+            self.ollama_api_url = self.args.ollama_api_url
+        if not self.request_queue:
+            self.request_queue = self.args.request_queue
+        if (prefix := self.args.request_queue_prefix):
+            self.request_queue = f"{prefix}{self.request_queue}"
 
     def run(self):
         with self.blocking_connection() as conn:

diff --git a/services/ollama/pyproject.toml b/services/ollama/pyproject.toml
@@ -5,6 +5,7 @@ authors = [{ name = "Gary Benson", email = "gary@gbenson.net" }]
 description = "Ollama connector service for Hive"
 requires-python = ">=3.10"  # match..case
 dependencies = [
+    "hive-common",
     "hive-messaging",
     "hive-service",
     "requests",