Sync chat models in Khoj with OpenAI proxies (e.g Ollama) on startup

- Allows managing chat models in the OpenAI proxy service like Ollama.
- Removes need to manually add, remove chat models from Khoj Admin Panel
  for these OpenAI compatible API services when enabled.
- Khoj still mantains the chat models configs within Khoj, so they can
  be configured via the Khoj admin panel as usual.
This commit is contained in:
Debanjum
2024-11-17 14:47:45 -08:00
parent 4a7f5d1abe
commit 285006d6c9

View File

@@ -231,6 +231,46 @@ def initialization(interactive: bool = True):
logger.info(f"🗣️ {provider_name} chat model configuration complete")
return True, chat_provider
def _update_chat_model_options():
"""Update available chat models for OpenAI-compatible APIs"""
try:
# Get OpenAI configs with custom base URLs
custom_configs = OpenAIProcessorConversationConfig.objects.exclude(api_base_url__isnull=True)
for config in custom_configs:
try:
# Create OpenAI client with custom base URL
openai_client = openai.OpenAI(api_key=config.api_key, base_url=config.api_base_url)
# Get available models
available_models = [model.id for model in openai_client.models.list()]
# Get existing chat model options for this config
existing_models = ChatModelOptions.objects.filter(
openai_config=config, model_type=ChatModelOptions.ModelType.OPENAI
)
# Add new models
for model in available_models:
if not existing_models.filter(chat_model=model).exists():
ChatModelOptions.objects.create(
chat_model=model,
model_type=ChatModelOptions.ModelType.OPENAI,
max_prompt_size=model_to_prompt_size.get(model),
vision_enabled=model in default_openai_chat_models,
tokenizer=model_to_tokenizer.get(model),
openai_config=config,
)
# Remove models that are no longer available
existing_models.exclude(chat_model__in=available_models).delete()
except Exception as e:
logger.warning(f"Failed to update models for {config.name}: {str(e)}")
except Exception as e:
logger.error(f"Failed to update chat model options: {str(e)}")
admin_user = KhojUser.objects.filter(is_staff=True).first()
if admin_user is None:
while True:
@@ -252,3 +292,6 @@ def initialization(interactive: bool = True):
return
except Exception as e:
logger.error(f"🚨 Failed to create chat configuration: {e}", exc_info=True)
else:
_update_chat_model_options()
logger.info("🗣️ Chat model configuration updated")