From 59602e15a7a8ad25e324d9a2c78e3e2e600f031d Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Mon, 19 May 2025 23:25:19 +0200
Subject: [PATCH] Ollama: Add basic configuration for LiteLLM proxy

---
 etc/litellm-config.yaml         | 43 +++++++++++++++++++++++++++++++++
 src/cratedb_about/query/core.py | 12 ++++-----
 tests/test_query.py             | 21 ----------------
 3 files changed, 49 insertions(+), 27 deletions(-)
 create mode 100644 etc/litellm-config.yaml

diff --git a/etc/litellm-config.yaml b/etc/litellm-config.yaml
new file mode 100644
index 0000000..139b4fb
--- /dev/null
+++ b/etc/litellm-config.yaml
@@ -0,0 +1,43 @@
+# LiteLLM supports all models from Ollama.
+# https://docs.litellm.ai/docs/providers/ollama
+#
+# uvx litellm --config etc/litellm-config.yaml --detailed_debug
+#
+# TODO: Use `ollama_chat`
+#       We recommend using ollama_chat for better responses.
+---
+
+model_list:
+  - model_name: "llama3.2"
+    litellm_params:
+      model: "ollama/llama3.2"
+      api_base: "http://localhost:11434"
+
+  - model_name: "gemma3:1b"
+    litellm_params:
+      model: "ollama/gemma3:1b"
+      api_base: "http://localhost:11434"
+
+  - model_name: "qwen3:0.6b"
+    litellm_params:
+      model: "ollama/qwen3:0.6b"
+      api_base: "http://localhost:11434"
+
+  - model_name: "deepseek-r1:7b"
+    litellm_params:
+      model: "ollama/deepseek-r1:7b"
+      api_base: "http://localhost:11434"
+
+# https://github.com/BerriAI/litellm/issues/1517#issuecomment-1922022209
+#model_list:
+#  - model_name: ollama-codellama
+#    litellm_params:
+#      model: ollama/codellama:70b
+#     api_base: http://0.0.0.0:11434
+#      rpm: 1440
+#    model_info:
+#      version: 2
+
+#litellm_settings:
+#  drop_params: True
+#  set_verbose: True
diff --git a/src/cratedb_about/query/core.py b/src/cratedb_about/query/core.py
index d6928bc..848453b 100644
--- a/src/cratedb_about/query/core.py
+++ b/src/cratedb_about/query/core.py
@@ -50,10 +50,6 @@ def __post_init__(self):
             raise ImportError("The 'openai' package is required when using the OpenAI backend")
         if self.backend == "claude" and not CLAUDE_AVAILABLE:
             raise ImportError("The 'claudette' package is required when using the Claude backend")
-        if self.backend == "openai" and not os.environ.get("OPENAI_API_KEY"):
-            raise ValueError(
-                "OPENAI_API_KEY environment variable is required when using 'openai' backend"
-            )
         if self.backend == "claude" and not os.environ.get("ANTHROPIC_API_KEY"):
             raise ValueError(
                 "ANTHROPIC_API_KEY environment variable is required when using 'claude' backend"
@@ -110,7 +106,7 @@ def ask_gpt(self, question: str) -> str:
         - https://community.openai.com/t/how-is-developer-message-better-than-system-prompt/1062784
         """
 
-        client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+        client = OpenAI(api_key="n/a", base_url="http://localhost:4000")
 
         input_messages: t.List[Message] = []
 
@@ -140,9 +136,13 @@ def ask_gpt(self, question: str) -> str:
         )
 
         # model = "gpt-4o"  # noqa: ERA001
-        model = "gpt-4.1"  # noqa: ERA001
+        # model = "gpt-4.1"  # noqa: ERA001
         # model = "o4-mini"  # noqa: ERA001
         # model = "o3"  # noqa: ERA001
+        model = "llama3.2"  # noqa: ERA001
+        # model = "gemma3:1b"  # noqa: ERA001
+        # model = "qwen3:0.6b"  # noqa: ERA001
+        # model = "deepseek-r1:7b"  # noqa: ERA001
         reasoning = None
         if model == "o4-mini":
             reasoning = Reasoning(
diff --git a/tests/test_query.py b/tests/test_query.py
index e9fe243..182d571 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -2,7 +2,6 @@
 
 import hishel
 import httpx
-import openai
 import pytest
 
 from cratedb_about import CrateDbKnowledgeConversation
@@ -54,26 +53,6 @@ def test_example_question():
     assert "How to enumerate active jobs?" in Example.questions
 
 
-def test_ask_openai_no_api_key():
-    """
-    Validate inquiry with OpenAI, failing without an API key.
-    """
-    with pytest.raises(ValueError) as ex:
-        CrateDbKnowledgeConversation()
-    assert ex.match("OPENAI_API_KEY environment variable is required when using 'openai' backend")
-
-
-def test_ask_openai_invalid_api_key(mocker):
-    """
-    Validate inquiry with OpenAI, failing when using an invalid API key.
-    """
-    mocker.patch.dict("os.environ", {"OPENAI_API_KEY": "foo"})
-    knowledge = CrateDbKnowledgeConversation()
-    with pytest.raises(openai.AuthenticationError) as ex:
-        knowledge.ask("CrateDB does not seem to provide an AUTOINCREMENT feature?")
-    assert ex.match("Incorrect API key provided: foo")
-
-
 @pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10 or higher")
 def test_ask_claude_no_api_key():
     """