From d98f6469291004b3218ab025cd1256c0949de717 Mon Sep 17 00:00:00 2001
From: Johannes Binder <j.binder@zoho.com>
Date: Wed, 28 Aug 2024 01:51:14 +0200
Subject: [PATCH] IPv6 support (#262)

* Support IPv6 host addresses

* Add missing tests

* Integrate feedback review to make it generally more stable and remove redundancies

* Remove unused import

* Fix formatting
---
 ollama/_client.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/ollama/_client.py b/ollama/_client.py
index 4b913d7..3124a58 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -1,3 +1,4 @@
+import ipaddress
 import os
 import io
 import json
@@ -995,6 +996,28 @@ def _parse_host(host: Optional[str]) -> str:
   'https://example.com:56789/path'
   >>> _parse_host('example.com:56789/path/')
   'http://example.com:56789/path'
+  >>> _parse_host('[0001:002:003:0004::1]')
+  'http://[0001:002:003:0004::1]:11434'
+  >>> _parse_host('[0001:002:003:0004::1]:56789')
+  'http://[0001:002:003:0004::1]:56789'
+  >>> _parse_host('http://[0001:002:003:0004::1]')
+  'http://[0001:002:003:0004::1]:80'
+  >>> _parse_host('https://[0001:002:003:0004::1]')
+  'https://[0001:002:003:0004::1]:443'
+  >>> _parse_host('https://[0001:002:003:0004::1]:56789')
+  'https://[0001:002:003:0004::1]:56789'
+  >>> _parse_host('[0001:002:003:0004::1]/')
+  'http://[0001:002:003:0004::1]:11434'
+  >>> _parse_host('[0001:002:003:0004::1]:56789/')
+  'http://[0001:002:003:0004::1]:56789'
+  >>> _parse_host('[0001:002:003:0004::1]/path')
+  'http://[0001:002:003:0004::1]:11434/path'
+  >>> _parse_host('[0001:002:003:0004::1]:56789/path')
+  'http://[0001:002:003:0004::1]:56789/path'
+  >>> _parse_host('https://[0001:002:003:0004::1]:56789/path')
+  'https://[0001:002:003:0004::1]:56789/path'
+  >>> _parse_host('[0001:002:003:0004::1]:56789/path/')
+  'http://[0001:002:003:0004::1]:56789/path'
   """
 
   host, port = host or '', 11434
@@ -1010,6 +1033,13 @@ def _parse_host(host: Optional[str]) -> str:
   host = split.hostname or '127.0.0.1'
   port = split.port or port
 
+  # Fix missing square brackets for IPv6 from urlsplit
+  try:
+    if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address):
+      host = f'[{host}]'
+  except ValueError:
+    ...
+
   if path := split.path.strip('/'):
     return f'{scheme}://{host}:{port}/{path}'