Fix language model configuration

Author: denkv

learn2rag/pipeline/llm.py

@@ -1,46 +1,69 @@
 import logging
 import os
 from pydantic import SecretStr
+from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_ollama import ChatOllama
 from langchain_openai import ChatOpenAI
-from typing import Callable, Any
 
 
-def ollama_client(*, url: str, token: str | None, model: str, proxy: str | None) -> ChatOllama:
-    return ChatOllama(
-        model=model,
-        temperature=0,
-        base_url=url,
-        client_kwargs={
-            'headers': {'Authorization': f'Bearer {token}'} if token else {},
-            'proxy': proxy,
-        },
-    )
+logger = logging.getLogger(__name__)
 
 
-def openai_client(*, url: str, token: SecretStr, model: str, proxy: str | None) -> ChatOpenAI:
-    return ChatOpenAI(
-        model=model,
-        temperature=0,
-        base_url=url,
-        api_key=token,
-    )
+class LLMClient():
+    # ID is used as a key to store in user data, should not be changed
+    ID: str
+    # LABEL is a display label for user interface
+    LABEL: str
+    chat_model: BaseChatModel
 
 
-# TODO: set up the right llm for user_config
+llms = {}
+def llm_client(cls: type[LLMClient]) -> type[LLMClient]:
+    llms[cls.ID] = cls; return cls
+
+
+# First @llm_client would be the default in UI when adding an external model
+@llm_client
+class OpenAIClient(LLMClient):
+    ID = 'ChatOpenAI'
+    LABEL = 'OpenAI'
+
+    def __init__(self, *, url: str, token: SecretStr, model: str, proxy: str | None) -> None:
+        self.chat_model = ChatOpenAI(
+            model=model,
+            temperature=0,
+            base_url=url,
+            api_key=token,
+        )
+
+
+@llm_client
+class OllamaClient(LLMClient):
+    ID = 'ChatOllama'
+    LABEL = 'Ollama'
+
+    def __init__(self, *, url: str, token: str | None, model: str, proxy: str | None) -> None:
+        self.chat_model = ChatOllama(
+            model=model,
+            temperature=0,
+            base_url=url,
+            client_kwargs={
+                'headers': {'Authorization': f'Bearer {token}'} if token else {},
+                'proxy': proxy,
+            },
+        )
+
+
+default_llm = OpenAIClient
+llm_id = os.environ.get('LLM_API_TYPE', default_llm.ID)
+logger.debug('Using LLM: %s', llm_id)
 
 llm_kwargs = {
     'url': os.environ.get('LLM_API_URL'),
     'token': os.environ.get('LLM_API_TOKEN') or None,
     'model': os.environ.get('LLM_API_MODEL'),
     'proxy': os.environ.get('LLM_API_PROXY') or None,
 }
-logging.info('LLM args: %s', llm_kwargs)
-
-# the keys are written by the configurator UI
-llms: dict[str, Callable[..., Any]] = {
-    'ChatOllama': ollama_client,
-    'ChatOpenAI': openai_client,
-}
+logger.debug('Using LLM args: %s', llm_kwargs)
 
-llm = llms[os.environ.get('LLM_API_TYPE', 'ChatOllama')](**llm_kwargs)
+llm = llms[llm_id](**llm_kwargs).chat_model

learn2rag/ui/__init__.py

@@ -26,6 +26,7 @@
 
 from learn2rag.compose import Project
 import learn2rag.data
+import learn2rag.pipeline.llm
 
 from datetime import datetime  # <-- ADD THIS
 
@@ -179,6 +180,7 @@ def inject_info() -> dict[str, Any]:
             'firststeps_storage_path': app.instance_path + '/storage/example',
             'debug_logging': config.get('logging', {}).get('debug', False),
             'current_timestamp': math.floor(time.time()),
+            'llm': learn2rag.pipeline.llm,
         }
 
     @app.context_processor
@@ -243,7 +245,7 @@ def model_create() -> 'str | werkzeug.wrappers.response.Response':
         ok = True
         model = request.form['model']
         api = request.form['api']
-        if api == 'ollama_clientent':
+        if api == learn2rag.pipeline.llm.OllamaClient.ID:
             url = request.form.get('url') or 'http://127.0.0.1:' + str(app.config['OLLAMA']['port']) + '/'
             # TODO setup tokens for locally running ollama
             token = request.form.get('token') or ''
@@ -252,7 +254,7 @@ def model_create() -> 'str | werkzeug.wrappers.response.Response':
                     model += ':latest'
                 start_project('ollama_download', components_template_path / 'ollama-download.yml', Path(), {'model': model})
                 return flask_redirect(url_for('model_pulling', model=model))
-        elif api == 'openai_clientent':
+        elif api == learn2rag.pipeline.llm.OpenAIClient.ID:
             url = request.form['url']
             token = request.form['token']
         else:
@@ -284,7 +286,7 @@ def model_pulling() -> 'str | werkzeug.wrappers.response.Response':
                     'url': 'http://127.0.0.1:' + str(app.config['OLLAMA']['port']) + '/',
                     'token': '',
                     'model': model,
-                    'api': 'ollama_clientent',
+                    'api': learn2rag.pipeline.llm.OllamaClient.ID,
                 })
                 flash(pgettext('flash', 'Downloaded a language model: %(model)s', model=model))
                 res = make_response(render_template('model_pulling_success.html'))

learn2rag/ui/config.yml

@@ -6,7 +6,7 @@ SUGGESTED_MODELS:
     link: https://ollama.com/library/gemma3
     ollama: pull
     config:
-      api: ollama_client
+      api: ChatOllama
       model: gemma3:27b
   llama3.3_70b:
     label: Meta Llama 3.3
@@ -15,13 +15,13 @@ SUGGESTED_MODELS:
     link: https://ollama.com/library/llama3.3
     ollama: pull
     config:
-      api: ollama_client
+      api: ChatOllama
       model: llama3.3:70b
   tinyllama:
     label: TinyLlama
     image: models/tinyllama.png
     link: https://github.com/jzhang38/TinyLlama
     ollama: pull
     config:
-      api: ollama_client
+      api: ChatOllama
       model: tinyllama:latest

learn2rag/ui/templates/models_list.html

@@ -30,8 +30,14 @@ <h2 class="accordion-header">
               <div class="mb-3">
                 <label for="api" class="form-label">{{pgettext('form_label', 'API type')}}</label>
                 <select required="required" class="form-select" name="api">
-                  <option selected="selected">ChatOpenAI</option>
-                  <option>ChatOllama</option>
+                  {% for llm_client in llm.llms.values() %}
+                  <option
+                    value="{{ llm_client.ID }}"
+                    {% if loop.first %}
+                    selected="selected"
+                    {% endif %}
+                  >{{ llm_client.LABEL }}</option>
+                  {% endfor %}
                 </select>
               </div>
               <div class="mb-3">
@@ -103,7 +109,7 @@ <h5 class="card-title">{{gettext('Another model')}}</h5>
                     <label for="model" class="form-label">{{pgettext('form_label', 'Language model')}}</label>
                     <input class="form-control" name="model" required="required" title="{{pgettext('tooltip', 'For example: %(model)s', model=firststeps_model.get('config', {}).get('model'))}}">
                   </div>
-                  <input type="hidden" name="api" value="ollama_clientent"/>
+                  <input type="hidden" name="api" value="{{ llm.OllamaClient.ID }}"/>
                   <input type="hidden" name="ollama" value="pull"/>
                   <button type="submit" class="btn btn-primary">Save</button>
                 </form>
@@ -140,7 +146,7 @@ <h2 class="accordion-header">
                   {% endfor %}
                 </select>
               </div>
-              <input type="hidden" name="api" value="ollama_clientent"/>
+              <input type="hidden" name="api" value="{{ llm.OllamaClient.ID }}"/>
               <input type="hidden" name="ollama" value="use"/>
               <button type="submit" class="btn btn-primary">{{pgettext('button', 'Save')}}</button>
             </form>
@@ -168,7 +174,14 @@ <h2 class="accordion-header">
     {% for name, model in models.items() %}
     <tr>
       <td title="{{ name }}">{{ model.label }}</td>
-      <td>{{ model.api }}</td>
+      <td>
+        {% if model.api in llm.llms %}
+        {{ llm.llms[model.api].LABEL }}
+        {% else %}
+        <code>{{ model.api }}</code>
+        <span title="{{pgettext('tooltip', 'Unknown value')}}" style="cursor: default;">⚠️</span>
+        {% endif %}
+      </td>
       <td>{{ model.url }}</td>
       <td>{{ model.token }}</td>
       <td>{{ model.model }}</td>