klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-06 05:39:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rename OpenAIProcessorConversationConfig DB model to more apt AiModelApi (#998)

Browse Source 
* Rename OpenAIProcessorConversationConfig to more apt AiModelAPI

The DB model name had drifted from what it is being used for,
a general chat api provider that supports other chat api providers like
anthropic and google chat models apart from openai based chat models.

This change renames the DB model and updates the docs to remove this
confusion.

Using Ai Model Api we catch most use-cases including chat, stt, image generation etc.
This commit is contained in:
Debanjum
2024-12-08 18:02:29 -08:00
committed by GitHub
parent df66fb23ab
commit 9dd3782f5c
18 changed files with 108 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/khoj/configure.py
View File

@@ -232,9 +232,9 @@ def configure_server(
config = FullConfig()
state.config = config
if ConversationAdapters.has_valid_openai_conversation_config():
openai_config = ConversationAdapters.get_openai_conversation_config()
state.openai_client = openai.OpenAI(api_key=openai_config.api_key)
if ConversationAdapters.has_valid_ai_model_api():
ai_model_api = ConversationAdapters.get_ai_model_api()
state.openai_client = openai.OpenAI(api_key=ai_model_api.api_key)
# Initialize Search Models from Config and initialize content
try:

30
src/khoj/database/adapters/__init__.py
View File

@@ -35,6 +35,7 @@ from torch import Tensor
from khoj.database.models import (
Agent,
AiModelApi,
ChatModelOptions,
ClientApplication,
Conversation,
@@ -46,7 +47,6 @@ from khoj.database.models import (
KhojApiUser,
KhojUser,
NotionConfig,
OpenAIProcessorConversationConfig,
ProcessLock,
PublicConversation,
ReflectiveQuestion,
@@ -981,7 +981,7 @@ class ConversationAdapters:
@staticmethod
async def aget_all_conversation_configs():
return await sync_to_async(list)(ChatModelOptions.objects.prefetch_related("openai_config").all())
return await sync_to_async(list)(ChatModelOptions.objects.prefetch_related("ai_model_api").all())
@staticmethod
def get_vision_enabled_config():
@@ -1000,12 +1000,12 @@ class ConversationAdapters:
return None
@staticmethod
def get_openai_conversation_config():
return OpenAIProcessorConversationConfig.objects.filter().first()
def get_ai_model_api():
return AiModelApi.objects.filter().first()
@staticmethod
def has_valid_openai_conversation_config():
return OpenAIProcessorConversationConfig.objects.filter().exists()
def has_valid_ai_model_api():
return AiModelApi.objects.filter().exists()
@staticmethod
@arequire_valid_user
@@ -1093,7 +1093,7 @@ class ConversationAdapters:
server_chat_settings: ServerChatSettings = (
await ServerChatSettings.objects.filter()
.prefetch_related(
"chat_default", "chat_default__openai_config", "chat_advanced", "chat_advanced__openai_config"
"chat_default", "chat_default__ai_model_api", "chat_advanced", "chat_advanced__ai_model_api"
)
.afirst()
)
@@ -1109,7 +1109,7 @@ class ConversationAdapters:
# Get the user's chat settings, if the server chat settings are not set
user_chat_settings = (
(await UserConversationConfig.objects.filter(user=user).prefetch_related("setting__openai_config").afirst())
(await UserConversationConfig.objects.filter(user=user).prefetch_related("setting__ai_model_api").afirst())
if user
else None
)
@@ -1117,7 +1117,7 @@ class ConversationAdapters:
return user_chat_settings.setting
# Get the first chat model if even the user chat settings are not set
return await ChatModelOptions.objects.filter().prefetch_related("openai_config").afirst()
return await ChatModelOptions.objects.filter().prefetch_related("ai_model_api").afirst()
@staticmethod
def get_advanced_conversation_config(user: KhojUser):
@@ -1130,7 +1130,7 @@ class ConversationAdapters:
async def aget_advanced_conversation_config(user: KhojUser = None):
server_chat_settings: ServerChatSettings = (
await ServerChatSettings.objects.filter()
.prefetch_related("chat_advanced", "chat_advanced__openai_config")
.prefetch_related("chat_advanced", "chat_advanced__ai_model_api")
.afirst()
)
if server_chat_settings is not None and server_chat_settings.chat_advanced is not None:
@@ -1258,7 +1258,7 @@ class ConversationAdapters:
@staticmethod
async def aget_user_conversation_config(user: KhojUser):
config = (
await UserConversationConfig.objects.filter(user=user).prefetch_related("setting__openai_config").afirst()
await UserConversationConfig.objects.filter(user=user).prefetch_related("setting__ai_model_api").afirst()
)
if not config:
return None
@@ -1313,7 +1313,7 @@ class ConversationAdapters:
ChatModelOptions.ModelType.OPENAI,
ChatModelOptions.ModelType.GOOGLE,
]
) and conversation_config.openai_config:
) and conversation_config.ai_model_api:
return conversation_config
else:
@@ -1321,7 +1321,7 @@ class ConversationAdapters:
@staticmethod
async def aget_text_to_image_model_config():
return await TextToImageModelConfig.objects.filter().prefetch_related("openai_config").afirst()
return await TextToImageModelConfig.objects.filter().prefetch_related("ai_model_api").afirst()
@staticmethod
def get_text_to_image_model_config():
@@ -1343,9 +1343,9 @@ class ConversationAdapters:
@staticmethod
async def aget_user_text_to_image_model(user: KhojUser) -> Optional[TextToImageModelConfig]:
# Create a custom queryset for prefetching settings__openai_config, handling null cases
# Create a custom queryset for prefetching settings__ai_model_api, handling null cases
settings_prefetch = Prefetch(
"setting", queryset=TextToImageModelConfig.objects.prefetch_related("openai_config")
"setting", queryset=TextToImageModelConfig.objects.prefetch_related("ai_model_api")
)
config = await UserTextToImageModelConfig.objects.filter(user=user).prefetch_related(settings_prefetch).afirst()

8
src/khoj/database/admin.py
View File

@@ -1,6 +1,6 @@
import csv
import json
from datetime import date, datetime, timedelta, timezone
from datetime import datetime, timedelta
from apscheduler.job import Job
from django.contrib import admin, messages
@@ -15,6 +15,7 @@ from unfold import admin as unfold_admin
from khoj.database.models import (
Agent,
AiModelApi,
ChatModelOptions,
ClientApplication,
Conversation,
@@ -22,7 +23,6 @@ from khoj.database.models import (
GithubConfig,
KhojUser,
NotionConfig,
OpenAIProcessorConversationConfig,
ProcessLock,
ReflectiveQuestion,
SearchModelConfig,
@@ -232,8 +232,8 @@ class TextToImageModelOptionsAdmin(unfold_admin.ModelAdmin):
search_fields = ("id", "model_name", "model_type")
@admin.register(OpenAIProcessorConversationConfig)
class OpenAIProcessorConversationConfigAdmin(unfold_admin.ModelAdmin):
@admin.register(AiModelApi)
class AiModelApiAdmin(unfold_admin.ModelAdmin):
list_display = (
"id",
"name",

26
src/khoj/database/migrations/0076_rename_openaiprocessorconversationconfig_aimodelapi_and_more.py Normal file
View File

@@ -0,0 +1,26 @@
# Generated by Django 5.0.9 on 2024-12-05 09:00
from django.db import migrations
class Migration(migrations.Migration):
dependencies = [
("database", "0075_migrate_generated_assets_and_validate"),
]
operations = [
migrations.RenameModel(
old_name="OpenAIProcessorConversationConfig",
new_name="AiModelApi",
),
migrations.RenameField(
model_name="chatmodeloptions",
old_name="openai_config",
new_name="ai_model_api",
),
migrations.RenameField(
model_name="texttoimagemodelconfig",
old_name="openai_config",
new_name="ai_model_api",
),
]

20
src/khoj/database/models/__init__.py
View File

@@ -181,7 +181,7 @@ class Subscription(DbBaseModel):
enabled_trial_at = models.DateTimeField(null=True, default=None, blank=True)
class OpenAIProcessorConversationConfig(DbBaseModel):
class AiModelApi(DbBaseModel):
name = models.CharField(max_length=200)
api_key = models.CharField(max_length=200)
api_base_url = models.URLField(max_length=200, default=None, blank=True, null=True)
@@ -200,9 +200,7 @@ class ChatModelOptions(DbBaseModel):
chat_model = models.CharField(max_length=200, default="bartowski/Meta-Llama-3.1-8B-Instruct-GGUF")
model_type = models.CharField(max_length=200, choices=ModelType.choices, default=ModelType.OFFLINE)
vision_enabled = models.BooleanField(default=False)
openai_config = models.ForeignKey(
OpenAIProcessorConversationConfig, on_delete=models.CASCADE, default=None, null=True, blank=True
)
ai_model_api = models.ForeignKey(AiModelApi, on_delete=models.CASCADE, default=None, null=True, blank=True)
class VoiceModelOption(DbBaseModel):
@@ -504,26 +502,24 @@ class TextToImageModelConfig(DbBaseModel):
model_name = models.CharField(max_length=200, default="dall-e-3")
model_type = models.CharField(max_length=200, choices=ModelType.choices, default=ModelType.OPENAI)
api_key = models.CharField(max_length=200, default=None, null=True, blank=True)
openai_config = models.ForeignKey(
OpenAIProcessorConversationConfig, on_delete=models.CASCADE, default=None, null=True, blank=True
)
ai_model_api = models.ForeignKey(AiModelApi, on_delete=models.CASCADE, default=None, null=True, blank=True)
def clean(self):
# Custom validation logic
error = {}
if self.model_type == self.ModelType.OPENAI:
if self.api_key and self.openai_config:
if self.api_key and self.ai_model_api:
error[
"api_key"
] = "Both API key and OpenAI config cannot be set for OpenAI models. Please set only one of them."
] = "Both API key and AI Model API cannot be set for OpenAI models. Please set only one of them."
error[
"openai_config"
"ai_model_api"
] = "Both API key and OpenAI config cannot be set for OpenAI models. Please set only one of them."
if self.model_type != self.ModelType.OPENAI:
if not self.api_key:
error["api_key"] = "The API key field must be set for non OpenAI models."
if self.openai_config:
error["openai_config"] = "OpenAI config cannot be set for non OpenAI models."
if self.ai_model_api:
error["ai_model_api"] = "AI Model API cannot be set for non OpenAI models."
if error:
raise ValidationError(error)

12
src/khoj/migrations/migrate_server_pg.py
View File

@@ -60,11 +60,7 @@ import logging
from packaging import version
from khoj.database.models import (
ChatModelOptions,
OpenAIProcessorConversationConfig,
SearchModelConfig,
)
from khoj.database.models import AiModelApi, ChatModelOptions, SearchModelConfig
from khoj.utils.yaml import load_config_from_file, save_config_to_file
logger = logging.getLogger(__name__)
@@ -121,16 +117,14 @@ def migrate_server_pg(args):
if openai.get("chat-model") is None:
openai["chat-model"] = "gpt-3.5-turbo"
openai_config = OpenAIProcessorConversationConfig.objects.create(
api_key=openai.get("api-key"), name="default"
)
openai_model_api = AiModelApi.objects.create(api_key=openai.get("api-key"), name="default")
ChatModelOptions.objects.create(
chat_model=openai.get("chat-model"),
tokenizer=processor_conversation.get("tokenizer"),
max_prompt_size=processor_conversation.get("max-prompt-size"),
model_type=ChatModelOptions.ModelType.OPENAI,
openai_config=openai_config,
ai_model_api=openai_model_api,
)
save_config_to_file(raw_config, args.config_file)

7
src/khoj/processor/conversation/openai/utils.py
View File

@@ -19,12 +19,7 @@ from khoj.processor.conversation.utils import (
ThreadedGenerator,
commit_conversation_trace,
)
from khoj.utils import state
from khoj.utils.helpers import (
get_chat_usage_metrics,
in_debug_mode,
is_promptrace_enabled,
)
from khoj.utils.helpers import get_chat_usage_metrics, is_promptrace_enabled
logger = logging.getLogger(__name__)

4
src/khoj/processor/image/generate.py
View File

@@ -124,8 +124,8 @@ def generate_image_with_openai(
# Get the API key from the user's configuration
if text_to_image_config.api_key:
api_key = text_to_image_config.api_key
elif text_to_image_config.openai_config:
api_key = text_to_image_config.openai_config.api_key
elif text_to_image_config.ai_model_api:
api_key = text_to_image_config.ai_model_api.api_key
elif state.openai_client:
api_key = state.openai_client.api_key
auth_header = {"Authorization": f"Bearer {api_key}"} if api_key else {}

9
src/khoj/routers/api.py
View File

@@ -430,9 +430,8 @@ async def extract_references_and_questions(
tracer=tracer,
)
elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
openai_chat_config = conversation_config.openai_config
api_key = openai_chat_config.api_key
base_url = openai_chat_config.api_base_url
api_key = conversation_config.ai_model_api.api_key
base_url = conversation_config.ai_model_api.api_base_url
chat_model = conversation_config.chat_model
inferred_queries = extract_questions(
defiltered_query,
@@ -449,7 +448,7 @@ async def extract_references_and_questions(
tracer=tracer,
)
elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
chat_model = conversation_config.chat_model
inferred_queries = extract_questions_anthropic(
defiltered_query,
@@ -465,7 +464,7 @@ async def extract_references_and_questions(
tracer=tracer,
)
elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
chat_model = conversation_config.chat_model
inferred_queries = extract_questions_gemini(
defiltered_query,

22
src/khoj/routers/helpers.py
View File

@@ -136,7 +136,7 @@ def validate_conversation_config(user: KhojUser):
if default_config is None:
raise HTTPException(status_code=500, detail="Contact the server administrator to add a chat model.")
if default_config.model_type == "openai" and not default_config.openai_config:
if default_config.model_type == "openai" and not default_config.ai_model_api:
raise HTTPException(status_code=500, detail="Contact the server administrator to add a chat model.")
@@ -163,7 +163,7 @@ async def is_ready_to_chat(user: KhojUser):
ChatModelOptions.ModelType.GOOGLE,
]
)
and user_conversation_config.openai_config
and user_conversation_config.ai_model_api
):
return True
@@ -990,7 +990,7 @@ async def send_message_to_model_wrapper(
)
elif model_type == ChatModelOptions.ModelType.OPENAI:
openai_chat_config = conversation_config.openai_config
openai_chat_config = conversation_config.ai_model_api
api_key = openai_chat_config.api_key
api_base_url = openai_chat_config.api_base_url
truncated_messages = generate_chatml_messages_with_context(
@@ -1015,7 +1015,7 @@ async def send_message_to_model_wrapper(
tracer=tracer,
)
elif model_type == ChatModelOptions.ModelType.ANTHROPIC:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
truncated_messages = generate_chatml_messages_with_context(
user_message=query,
context_message=context,
@@ -1037,7 +1037,7 @@ async def send_message_to_model_wrapper(
tracer=tracer,
)
elif model_type == ChatModelOptions.ModelType.GOOGLE:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
truncated_messages = generate_chatml_messages_with_context(
user_message=query,
context_message=context,
@@ -1102,7 +1102,7 @@ def send_message_to_model_wrapper_sync(
)
elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
truncated_messages = generate_chatml_messages_with_context(
user_message=message,
system_message=system_message,
@@ -1124,7 +1124,7 @@ def send_message_to_model_wrapper_sync(
return openai_response
elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
truncated_messages = generate_chatml_messages_with_context(
user_message=message,
system_message=system_message,
@@ -1144,7 +1144,7 @@ def send_message_to_model_wrapper_sync(
)
elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
truncated_messages = generate_chatml_messages_with_context(
user_message=message,
system_message=system_message,
@@ -1255,7 +1255,7 @@ def generate_chat_response(
)
elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
openai_chat_config = conversation_config.openai_config
openai_chat_config = conversation_config.ai_model_api
api_key = openai_chat_config.api_key
chat_model = conversation_config.chat_model
chat_response = converse(
@@ -1285,7 +1285,7 @@ def generate_chat_response(
)
elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
chat_response = converse_anthropic(
compiled_references,
query_to_run,
@@ -1311,7 +1311,7 @@ def generate_chat_response(
tracer=tracer,
)
elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
api_key = conversation_config.openai_config.api_key
api_key = conversation_config.ai_model_api.api_key
chat_response = converse_gemini(
compiled_references,
query_to_run,

22
src/khoj/utils/initialization.py
View File

@@ -6,9 +6,9 @@ import openai
from khoj.database.adapters import ConversationAdapters
from khoj.database.models import (
AiModelApi,
ChatModelOptions,
KhojUser,
OpenAIProcessorConversationConfig,
SpeechToTextModelOptions,
TextToImageModelConfig,
)
@@ -98,7 +98,7 @@ def initialization(interactive: bool = True):
TextToImageModelConfig.objects.create(
model_name=openai_text_to_image_model,
model_type=TextToImageModelConfig.ModelType.OPENAI,
openai_config=openai_provider,
ai_model_api=openai_provider,
)
# Set up Google's Gemini online chat models
@@ -177,7 +177,7 @@ def initialization(interactive: bool = True):
vision_enabled: bool = False,
is_offline: bool = False,
provider_name: str = None,
) -> Tuple[bool, OpenAIProcessorConversationConfig]:
) -> Tuple[bool, AiModelApi]:
supported_vision_models = (
default_openai_chat_models + default_anthropic_chat_models + default_gemini_chat_models
)
@@ -192,16 +192,14 @@ def initialization(interactive: bool = True):
logger.info(f"️💬 Setting up your {provider_name} chat configuration")
chat_provider = None
ai_model_api = None
if not is_offline:
if interactive:
user_api_key = input(f"Enter your {provider_name} API key (default: {default_api_key}): ")
api_key = user_api_key if user_api_key != "" else default_api_key
else:
api_key = default_api_key
chat_provider = OpenAIProcessorConversationConfig.objects.create(
api_key=api_key, name=provider_name, api_base_url=api_base_url
)
ai_model_api = AiModelApi.objects.create(api_key=api_key, name=provider_name, api_base_url=api_base_url)
if interactive:
chat_model_names = input(
@@ -223,19 +221,19 @@ def initialization(interactive: bool = True):
"max_prompt_size": default_max_tokens,
"vision_enabled": vision_enabled,
"tokenizer": default_tokenizer,
"openai_config": chat_provider,
"ai_model_api": ai_model_api,
}
ChatModelOptions.objects.create(**chat_model_options)
logger.info(f"🗣️ {provider_name} chat model configuration complete")
return True, chat_provider
return True, ai_model_api
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)
custom_configs = AiModelApi.objects.exclude(api_base_url__isnull=True)
for config in custom_configs:
try:
@@ -247,7 +245,7 @@ def initialization(interactive: bool = True):
# Get existing chat model options for this config
existing_models = ChatModelOptions.objects.filter(
openai_config=config, model_type=ChatModelOptions.ModelType.OPENAI
ai_model_api=config, model_type=ChatModelOptions.ModelType.OPENAI
)
# Add new models
@@ -259,7 +257,7 @@ def initialization(interactive: bool = True):
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,
ai_model_api=config,
)
# Remove models that are no longer available