Skip to content

Commit

Permalink
Sync chat models in Khoj with OpenAI proxies (e.g Ollama) on startup
Browse files Browse the repository at this point in the history
- 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.
  • Loading branch information
debanjum committed Nov 17, 2024
1 parent 4a7f5d1 commit 285006d
Showing 1 changed file with 43 additions and 0 deletions.
43 changes: 43 additions & 0 deletions src/khoj/utils/initialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,46 @@ def _setup_chat_model_provider(
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:
Expand All @@ -252,3 +292,6 @@ def _setup_chat_model_provider(
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")

0 comments on commit 285006d

Please sign in to comment.