khoj 1.30.11.dev73__py3-none-any.whl → 1.31.1.dev13__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. khoj/database/adapters/__init__.py +51 -51
  2. khoj/database/admin.py +8 -8
  3. khoj/database/migrations/0077_chatmodel_alter_agent_chat_model_and_more.py +62 -0
  4. khoj/database/models/__init__.py +7 -7
  5. khoj/interface/compiled/404/index.html +1 -1
  6. khoj/interface/compiled/_next/static/chunks/app/agents/layout-1878cc328ea380bd.js +1 -0
  7. khoj/interface/compiled/_next/static/chunks/app/share/chat/layout-592e8c470f2c2084.js +1 -0
  8. khoj/interface/compiled/_next/static/chunks/{webpack-a72e892f6afa93ae.js → webpack-062298330010d2aa.js} +1 -1
  9. khoj/interface/compiled/_next/static/css/f172a0fb3eb177e1.css +1 -0
  10. khoj/interface/compiled/agents/index.html +1 -1
  11. khoj/interface/compiled/agents/index.txt +1 -1
  12. khoj/interface/compiled/automations/index.html +1 -1
  13. khoj/interface/compiled/automations/index.txt +1 -1
  14. khoj/interface/compiled/chat/index.html +1 -1
  15. khoj/interface/compiled/chat/index.txt +1 -1
  16. khoj/interface/compiled/index.html +1 -1
  17. khoj/interface/compiled/index.txt +1 -1
  18. khoj/interface/compiled/search/index.html +1 -1
  19. khoj/interface/compiled/search/index.txt +1 -1
  20. khoj/interface/compiled/settings/index.html +1 -1
  21. khoj/interface/compiled/settings/index.txt +1 -1
  22. khoj/interface/compiled/share/chat/index.html +1 -1
  23. khoj/interface/compiled/share/chat/index.txt +1 -1
  24. khoj/migrations/migrate_server_pg.py +7 -7
  25. khoj/processor/conversation/anthropic/anthropic_chat.py +3 -3
  26. khoj/processor/conversation/google/gemini_chat.py +3 -3
  27. khoj/processor/conversation/offline/chat_model.py +12 -12
  28. khoj/processor/conversation/openai/gpt.py +4 -4
  29. khoj/processor/conversation/openai/utils.py +18 -10
  30. khoj/processor/conversation/utils.py +4 -4
  31. khoj/routers/api.py +22 -27
  32. khoj/routers/api_agents.py +4 -4
  33. khoj/routers/api_chat.py +6 -6
  34. khoj/routers/api_model.py +4 -4
  35. khoj/routers/helpers.py +106 -102
  36. khoj/utils/helpers.py +5 -3
  37. khoj/utils/initialization.py +28 -26
  38. {khoj-1.30.11.dev73.dist-info → khoj-1.31.1.dev13.dist-info}/METADATA +1 -1
  39. {khoj-1.30.11.dev73.dist-info → khoj-1.31.1.dev13.dist-info}/RECORD +44 -43
  40. khoj/interface/compiled/_next/static/chunks/app/agents/layout-f2ea2b26fc0e78b1.js +0 -1
  41. khoj/interface/compiled/_next/static/chunks/app/share/chat/layout-f662c9e5091603cf.js +0 -1
  42. khoj/interface/compiled/_next/static/css/0dff1a59259ff945.css +0 -1
  43. /khoj/interface/compiled/_next/static/{ruEWgIV3qd6RPykTrcoyc → 7K5OLB23CafMhZBFHh4NG}/_buildManifest.js +0 -0
  44. /khoj/interface/compiled/_next/static/{ruEWgIV3qd6RPykTrcoyc → 7K5OLB23CafMhZBFHh4NG}/_ssgManifest.js +0 -0
  45. {khoj-1.30.11.dev73.dist-info → khoj-1.31.1.dev13.dist-info}/WHEEL +0 -0
  46. {khoj-1.30.11.dev73.dist-info → khoj-1.31.1.dev13.dist-info}/entry_points.txt +0 -0
  47. {khoj-1.30.11.dev73.dist-info → khoj-1.31.1.dev13.dist-info}/licenses/LICENSE +0 -0
@@ -62,7 +62,7 @@ async def all_agents(
62
62
  "color": agent.style_color,
63
63
  "icon": agent.style_icon,
64
64
  "privacy_level": agent.privacy_level,
65
- "chat_model": agent.chat_model.chat_model,
65
+ "chat_model": agent.chat_model.name,
66
66
  "files": file_names,
67
67
  "input_tools": agent.input_tools,
68
68
  "output_modes": agent.output_modes,
@@ -150,7 +150,7 @@ async def get_agent(
150
150
  "color": agent.style_color,
151
151
  "icon": agent.style_icon,
152
152
  "privacy_level": agent.privacy_level,
153
- "chat_model": agent.chat_model.chat_model,
153
+ "chat_model": agent.chat_model.name,
154
154
  "files": file_names,
155
155
  "input_tools": agent.input_tools,
156
156
  "output_modes": agent.output_modes,
@@ -225,7 +225,7 @@ async def create_agent(
225
225
  "color": agent.style_color,
226
226
  "icon": agent.style_icon,
227
227
  "privacy_level": agent.privacy_level,
228
- "chat_model": agent.chat_model.chat_model,
228
+ "chat_model": agent.chat_model.name,
229
229
  "files": body.files,
230
230
  "input_tools": agent.input_tools,
231
231
  "output_modes": agent.output_modes,
@@ -286,7 +286,7 @@ async def update_agent(
286
286
  "color": agent.style_color,
287
287
  "icon": agent.style_icon,
288
288
  "privacy_level": agent.privacy_level,
289
- "chat_model": agent.chat_model.chat_model,
289
+ "chat_model": agent.chat_model.name,
290
290
  "files": body.files,
291
291
  "input_tools": agent.input_tools,
292
292
  "output_modes": agent.output_modes,
khoj/routers/api_chat.py CHANGED
@@ -58,7 +58,7 @@ from khoj.routers.helpers import (
58
58
  is_ready_to_chat,
59
59
  read_chat_stream,
60
60
  update_telemetry_state,
61
- validate_conversation_config,
61
+ validate_chat_model,
62
62
  )
63
63
  from khoj.routers.research import (
64
64
  InformationCollectionIteration,
@@ -205,7 +205,7 @@ def chat_history(
205
205
  n: Optional[int] = None,
206
206
  ):
207
207
  user = request.user.object
208
- validate_conversation_config(user)
208
+ validate_chat_model(user)
209
209
 
210
210
  # Load Conversation History
211
211
  conversation = ConversationAdapters.get_conversation_by_user(
@@ -898,10 +898,10 @@ async def chat(
898
898
  custom_filters = []
899
899
  if conversation_commands == [ConversationCommand.Help]:
900
900
  if not q:
901
- conversation_config = await ConversationAdapters.aget_user_conversation_config(user)
902
- if conversation_config == None:
903
- conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
904
- model_type = conversation_config.model_type
901
+ chat_model = await ConversationAdapters.aget_user_chat_model(user)
902
+ if chat_model == None:
903
+ chat_model = await ConversationAdapters.aget_default_chat_model(user)
904
+ model_type = chat_model.model_type
905
905
  formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
906
906
  async for result in send_llm_response(formatted_help, tracer.get("usage")):
907
907
  yield result
khoj/routers/api_model.py CHANGED
@@ -24,7 +24,7 @@ def get_chat_model_options(
24
24
 
25
25
  all_conversation_options = list()
26
26
  for conversation_option in conversation_options:
27
- all_conversation_options.append({"chat_model": conversation_option.chat_model, "id": conversation_option.id})
27
+ all_conversation_options.append({"chat_model": conversation_option.name, "id": conversation_option.id})
28
28
 
29
29
  return Response(content=json.dumps(all_conversation_options), media_type="application/json", status_code=200)
30
30
 
@@ -37,12 +37,12 @@ def get_user_chat_model(
37
37
  ):
38
38
  user = request.user.object
39
39
 
40
- chat_model = ConversationAdapters.get_conversation_config(user)
40
+ chat_model = ConversationAdapters.get_chat_model(user)
41
41
 
42
42
  if chat_model is None:
43
- chat_model = ConversationAdapters.get_default_conversation_config(user)
43
+ chat_model = ConversationAdapters.get_default_chat_model(user)
44
44
 
45
- return Response(status_code=200, content=json.dumps({"id": chat_model.id, "chat_model": chat_model.chat_model}))
45
+ return Response(status_code=200, content=json.dumps({"id": chat_model.id, "chat_model": chat_model.name}))
46
46
 
47
47
 
48
48
  @api_model.post("/chat", status_code=200)
khoj/routers/helpers.py CHANGED
@@ -56,7 +56,7 @@ from khoj.database.adapters import (
56
56
  )
57
57
  from khoj.database.models import (
58
58
  Agent,
59
- ChatModelOptions,
59
+ ChatModel,
60
60
  ClientApplication,
61
61
  Conversation,
62
62
  GithubConfig,
@@ -133,40 +133,40 @@ def is_query_empty(query: str) -> bool:
133
133
  return is_none_or_empty(query.strip())
134
134
 
135
135
 
136
- def validate_conversation_config(user: KhojUser):
137
- default_config = ConversationAdapters.get_default_conversation_config(user)
136
+ def validate_chat_model(user: KhojUser):
137
+ default_chat_model = ConversationAdapters.get_default_chat_model(user)
138
138
 
139
- if default_config is None:
139
+ if default_chat_model is None:
140
140
  raise HTTPException(status_code=500, detail="Contact the server administrator to add a chat model.")
141
141
 
142
- if default_config.model_type == "openai" and not default_config.ai_model_api:
142
+ if default_chat_model.model_type == "openai" and not default_chat_model.ai_model_api:
143
143
  raise HTTPException(status_code=500, detail="Contact the server administrator to add a chat model.")
144
144
 
145
145
 
146
146
  async def is_ready_to_chat(user: KhojUser):
147
- user_conversation_config = await ConversationAdapters.aget_user_conversation_config(user)
148
- if user_conversation_config == None:
149
- user_conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
147
+ user_chat_model = await ConversationAdapters.aget_user_chat_model(user)
148
+ if user_chat_model == None:
149
+ user_chat_model = await ConversationAdapters.aget_default_chat_model(user)
150
150
 
151
- if user_conversation_config and user_conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:
152
- chat_model = user_conversation_config.chat_model
153
- max_tokens = user_conversation_config.max_prompt_size
151
+ if user_chat_model and user_chat_model.model_type == ChatModel.ModelType.OFFLINE:
152
+ chat_model_name = user_chat_model.name
153
+ max_tokens = user_chat_model.max_prompt_size
154
154
  if state.offline_chat_processor_config is None:
155
155
  logger.info("Loading Offline Chat Model...")
156
- state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model, max_tokens)
156
+ state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model_name, max_tokens)
157
157
  return True
158
158
 
159
159
  if (
160
- user_conversation_config
160
+ user_chat_model
161
161
  and (
162
- user_conversation_config.model_type
162
+ user_chat_model.model_type
163
163
  in [
164
- ChatModelOptions.ModelType.OPENAI,
165
- ChatModelOptions.ModelType.ANTHROPIC,
166
- ChatModelOptions.ModelType.GOOGLE,
164
+ ChatModel.ModelType.OPENAI,
165
+ ChatModel.ModelType.ANTHROPIC,
166
+ ChatModel.ModelType.GOOGLE,
167
167
  ]
168
168
  )
169
- and user_conversation_config.ai_model_api
169
+ and user_chat_model.ai_model_api
170
170
  ):
171
171
  return True
172
172
 
@@ -942,120 +942,124 @@ async def send_message_to_model_wrapper(
942
942
  query_files: str = None,
943
943
  tracer: dict = {},
944
944
  ):
945
- conversation_config: ChatModelOptions = await ConversationAdapters.aget_default_conversation_config(user)
946
- vision_available = conversation_config.vision_enabled
945
+ chat_model: ChatModel = await ConversationAdapters.aget_default_chat_model(user)
946
+ vision_available = chat_model.vision_enabled
947
947
  if not vision_available and query_images:
948
- logger.warning(f"Vision is not enabled for default model: {conversation_config.chat_model}.")
948
+ logger.warning(f"Vision is not enabled for default model: {chat_model.name}.")
949
949
  vision_enabled_config = await ConversationAdapters.aget_vision_enabled_config()
950
950
  if vision_enabled_config:
951
- conversation_config = vision_enabled_config
951
+ chat_model = vision_enabled_config
952
952
  vision_available = True
953
953
  if vision_available and query_images:
954
- logger.info(f"Using {conversation_config.chat_model} model to understand {len(query_images)} images.")
954
+ logger.info(f"Using {chat_model.name} model to understand {len(query_images)} images.")
955
955
 
956
956
  subscribed = await ais_user_subscribed(user)
957
- chat_model = conversation_config.chat_model
957
+ chat_model_name = chat_model.name
958
958
  max_tokens = (
959
- conversation_config.subscribed_max_prompt_size
960
- if subscribed and conversation_config.subscribed_max_prompt_size
961
- else conversation_config.max_prompt_size
959
+ chat_model.subscribed_max_prompt_size
960
+ if subscribed and chat_model.subscribed_max_prompt_size
961
+ else chat_model.max_prompt_size
962
962
  )
963
- tokenizer = conversation_config.tokenizer
964
- model_type = conversation_config.model_type
965
- vision_available = conversation_config.vision_enabled
963
+ tokenizer = chat_model.tokenizer
964
+ model_type = chat_model.model_type
965
+ vision_available = chat_model.vision_enabled
966
966
 
967
- if model_type == ChatModelOptions.ModelType.OFFLINE:
967
+ if model_type == ChatModel.ModelType.OFFLINE:
968
968
  if state.offline_chat_processor_config is None or state.offline_chat_processor_config.loaded_model is None:
969
- state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model, max_tokens)
969
+ state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model_name, max_tokens)
970
970
 
971
971
  loaded_model = state.offline_chat_processor_config.loaded_model
972
972
  truncated_messages = generate_chatml_messages_with_context(
973
973
  user_message=query,
974
974
  context_message=context,
975
975
  system_message=system_message,
976
- model_name=chat_model,
976
+ model_name=chat_model_name,
977
977
  loaded_model=loaded_model,
978
978
  tokenizer_name=tokenizer,
979
979
  max_prompt_size=max_tokens,
980
980
  vision_enabled=vision_available,
981
- model_type=conversation_config.model_type,
981
+ model_type=chat_model.model_type,
982
982
  query_files=query_files,
983
983
  )
984
984
 
985
985
  return send_message_to_model_offline(
986
986
  messages=truncated_messages,
987
987
  loaded_model=loaded_model,
988
- model=chat_model,
988
+ model_name=chat_model_name,
989
989
  max_prompt_size=max_tokens,
990
990
  streaming=False,
991
991
  response_type=response_type,
992
992
  tracer=tracer,
993
993
  )
994
994
 
995
- elif model_type == ChatModelOptions.ModelType.OPENAI:
996
- openai_chat_config = conversation_config.ai_model_api
995
+ elif model_type == ChatModel.ModelType.OPENAI:
996
+ openai_chat_config = chat_model.ai_model_api
997
997
  api_key = openai_chat_config.api_key
998
998
  api_base_url = openai_chat_config.api_base_url
999
999
  truncated_messages = generate_chatml_messages_with_context(
1000
1000
  user_message=query,
1001
1001
  context_message=context,
1002
1002
  system_message=system_message,
1003
- model_name=chat_model,
1003
+ model_name=chat_model_name,
1004
1004
  max_prompt_size=max_tokens,
1005
1005
  tokenizer_name=tokenizer,
1006
1006
  vision_enabled=vision_available,
1007
1007
  query_images=query_images,
1008
- model_type=conversation_config.model_type,
1008
+ model_type=chat_model.model_type,
1009
1009
  query_files=query_files,
1010
1010
  )
1011
1011
 
1012
1012
  return send_message_to_model(
1013
1013
  messages=truncated_messages,
1014
1014
  api_key=api_key,
1015
- model=chat_model,
1015
+ model=chat_model_name,
1016
1016
  response_type=response_type,
1017
1017
  api_base_url=api_base_url,
1018
1018
  tracer=tracer,
1019
1019
  )
1020
- elif model_type == ChatModelOptions.ModelType.ANTHROPIC:
1021
- api_key = conversation_config.ai_model_api.api_key
1020
+ elif model_type == ChatModel.ModelType.ANTHROPIC:
1021
+ api_key = chat_model.ai_model_api.api_key
1022
1022
  truncated_messages = generate_chatml_messages_with_context(
1023
1023
  user_message=query,
1024
1024
  context_message=context,
1025
1025
  system_message=system_message,
1026
- model_name=chat_model,
1026
+ model_name=chat_model_name,
1027
1027
  max_prompt_size=max_tokens,
1028
1028
  tokenizer_name=tokenizer,
1029
1029
  vision_enabled=vision_available,
1030
1030
  query_images=query_images,
1031
- model_type=conversation_config.model_type,
1031
+ model_type=chat_model.model_type,
1032
1032
  query_files=query_files,
1033
1033
  )
1034
1034
 
1035
1035
  return anthropic_send_message_to_model(
1036
1036
  messages=truncated_messages,
1037
1037
  api_key=api_key,
1038
- model=chat_model,
1038
+ model=chat_model_name,
1039
1039
  response_type=response_type,
1040
1040
  tracer=tracer,
1041
1041
  )
1042
- elif model_type == ChatModelOptions.ModelType.GOOGLE:
1043
- api_key = conversation_config.ai_model_api.api_key
1042
+ elif model_type == ChatModel.ModelType.GOOGLE:
1043
+ api_key = chat_model.ai_model_api.api_key
1044
1044
  truncated_messages = generate_chatml_messages_with_context(
1045
1045
  user_message=query,
1046
1046
  context_message=context,
1047
1047
  system_message=system_message,
1048
- model_name=chat_model,
1048
+ model_name=chat_model_name,
1049
1049
  max_prompt_size=max_tokens,
1050
1050
  tokenizer_name=tokenizer,
1051
1051
  vision_enabled=vision_available,
1052
1052
  query_images=query_images,
1053
- model_type=conversation_config.model_type,
1053
+ model_type=chat_model.model_type,
1054
1054
  query_files=query_files,
1055
1055
  )
1056
1056
 
1057
1057
  return gemini_send_message_to_model(
1058
- messages=truncated_messages, api_key=api_key, model=chat_model, response_type=response_type, tracer=tracer
1058
+ messages=truncated_messages,
1059
+ api_key=api_key,
1060
+ model=chat_model_name,
1061
+ response_type=response_type,
1062
+ tracer=tracer,
1059
1063
  )
1060
1064
  else:
1061
1065
  raise HTTPException(status_code=500, detail="Invalid conversation config")
@@ -1069,99 +1073,99 @@ def send_message_to_model_wrapper_sync(
1069
1073
  query_files: str = "",
1070
1074
  tracer: dict = {},
1071
1075
  ):
1072
- conversation_config: ChatModelOptions = ConversationAdapters.get_default_conversation_config(user)
1076
+ chat_model: ChatModel = ConversationAdapters.get_default_chat_model(user)
1073
1077
 
1074
- if conversation_config is None:
1078
+ if chat_model is None:
1075
1079
  raise HTTPException(status_code=500, detail="Contact the server administrator to set a default chat model.")
1076
1080
 
1077
- chat_model = conversation_config.chat_model
1078
- max_tokens = conversation_config.max_prompt_size
1079
- vision_available = conversation_config.vision_enabled
1081
+ chat_model_name = chat_model.name
1082
+ max_tokens = chat_model.max_prompt_size
1083
+ vision_available = chat_model.vision_enabled
1080
1084
 
1081
- if conversation_config.model_type == ChatModelOptions.ModelType.OFFLINE:
1085
+ if chat_model.model_type == ChatModel.ModelType.OFFLINE:
1082
1086
  if state.offline_chat_processor_config is None or state.offline_chat_processor_config.loaded_model is None:
1083
- state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model, max_tokens)
1087
+ state.offline_chat_processor_config = OfflineChatProcessorModel(chat_model_name, max_tokens)
1084
1088
 
1085
1089
  loaded_model = state.offline_chat_processor_config.loaded_model
1086
1090
  truncated_messages = generate_chatml_messages_with_context(
1087
1091
  user_message=message,
1088
1092
  system_message=system_message,
1089
- model_name=chat_model,
1093
+ model_name=chat_model_name,
1090
1094
  loaded_model=loaded_model,
1091
1095
  max_prompt_size=max_tokens,
1092
1096
  vision_enabled=vision_available,
1093
- model_type=conversation_config.model_type,
1097
+ model_type=chat_model.model_type,
1094
1098
  query_files=query_files,
1095
1099
  )
1096
1100
 
1097
1101
  return send_message_to_model_offline(
1098
1102
  messages=truncated_messages,
1099
1103
  loaded_model=loaded_model,
1100
- model=chat_model,
1104
+ model_name=chat_model_name,
1101
1105
  max_prompt_size=max_tokens,
1102
1106
  streaming=False,
1103
1107
  response_type=response_type,
1104
1108
  tracer=tracer,
1105
1109
  )
1106
1110
 
1107
- elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
1108
- api_key = conversation_config.ai_model_api.api_key
1111
+ elif chat_model.model_type == ChatModel.ModelType.OPENAI:
1112
+ api_key = chat_model.ai_model_api.api_key
1109
1113
  truncated_messages = generate_chatml_messages_with_context(
1110
1114
  user_message=message,
1111
1115
  system_message=system_message,
1112
- model_name=chat_model,
1116
+ model_name=chat_model_name,
1113
1117
  max_prompt_size=max_tokens,
1114
1118
  vision_enabled=vision_available,
1115
- model_type=conversation_config.model_type,
1119
+ model_type=chat_model.model_type,
1116
1120
  query_files=query_files,
1117
1121
  )
1118
1122
 
1119
1123
  openai_response = send_message_to_model(
1120
1124
  messages=truncated_messages,
1121
1125
  api_key=api_key,
1122
- model=chat_model,
1126
+ model=chat_model_name,
1123
1127
  response_type=response_type,
1124
1128
  tracer=tracer,
1125
1129
  )
1126
1130
 
1127
1131
  return openai_response
1128
1132
 
1129
- elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
1130
- api_key = conversation_config.ai_model_api.api_key
1133
+ elif chat_model.model_type == ChatModel.ModelType.ANTHROPIC:
1134
+ api_key = chat_model.ai_model_api.api_key
1131
1135
  truncated_messages = generate_chatml_messages_with_context(
1132
1136
  user_message=message,
1133
1137
  system_message=system_message,
1134
- model_name=chat_model,
1138
+ model_name=chat_model_name,
1135
1139
  max_prompt_size=max_tokens,
1136
1140
  vision_enabled=vision_available,
1137
- model_type=conversation_config.model_type,
1141
+ model_type=chat_model.model_type,
1138
1142
  query_files=query_files,
1139
1143
  )
1140
1144
 
1141
1145
  return anthropic_send_message_to_model(
1142
1146
  messages=truncated_messages,
1143
1147
  api_key=api_key,
1144
- model=chat_model,
1148
+ model=chat_model_name,
1145
1149
  response_type=response_type,
1146
1150
  tracer=tracer,
1147
1151
  )
1148
1152
 
1149
- elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
1150
- api_key = conversation_config.ai_model_api.api_key
1153
+ elif chat_model.model_type == ChatModel.ModelType.GOOGLE:
1154
+ api_key = chat_model.ai_model_api.api_key
1151
1155
  truncated_messages = generate_chatml_messages_with_context(
1152
1156
  user_message=message,
1153
1157
  system_message=system_message,
1154
- model_name=chat_model,
1158
+ model_name=chat_model_name,
1155
1159
  max_prompt_size=max_tokens,
1156
1160
  vision_enabled=vision_available,
1157
- model_type=conversation_config.model_type,
1161
+ model_type=chat_model.model_type,
1158
1162
  query_files=query_files,
1159
1163
  )
1160
1164
 
1161
1165
  return gemini_send_message_to_model(
1162
1166
  messages=truncated_messages,
1163
1167
  api_key=api_key,
1164
- model=chat_model,
1168
+ model=chat_model_name,
1165
1169
  response_type=response_type,
1166
1170
  tracer=tracer,
1167
1171
  )
@@ -1229,15 +1233,15 @@ def generate_chat_response(
1229
1233
  online_results = {}
1230
1234
  code_results = {}
1231
1235
 
1232
- conversation_config = ConversationAdapters.get_valid_conversation_config(user, conversation)
1233
- vision_available = conversation_config.vision_enabled
1236
+ chat_model = ConversationAdapters.get_valid_chat_model(user, conversation)
1237
+ vision_available = chat_model.vision_enabled
1234
1238
  if not vision_available and query_images:
1235
1239
  vision_enabled_config = ConversationAdapters.get_vision_enabled_config()
1236
1240
  if vision_enabled_config:
1237
- conversation_config = vision_enabled_config
1241
+ chat_model = vision_enabled_config
1238
1242
  vision_available = True
1239
1243
 
1240
- if conversation_config.model_type == "offline":
1244
+ if chat_model.model_type == "offline":
1241
1245
  loaded_model = state.offline_chat_processor_config.loaded_model
1242
1246
  chat_response = converse_offline(
1243
1247
  user_query=query_to_run,
@@ -1247,9 +1251,9 @@ def generate_chat_response(
1247
1251
  conversation_log=meta_log,
1248
1252
  completion_func=partial_completion,
1249
1253
  conversation_commands=conversation_commands,
1250
- model=conversation_config.chat_model,
1251
- max_prompt_size=conversation_config.max_prompt_size,
1252
- tokenizer_name=conversation_config.tokenizer,
1254
+ model_name=chat_model.name,
1255
+ max_prompt_size=chat_model.max_prompt_size,
1256
+ tokenizer_name=chat_model.tokenizer,
1253
1257
  location_data=location_data,
1254
1258
  user_name=user_name,
1255
1259
  agent=agent,
@@ -1259,10 +1263,10 @@ def generate_chat_response(
1259
1263
  tracer=tracer,
1260
1264
  )
1261
1265
 
1262
- elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
1263
- openai_chat_config = conversation_config.ai_model_api
1266
+ elif chat_model.model_type == ChatModel.ModelType.OPENAI:
1267
+ openai_chat_config = chat_model.ai_model_api
1264
1268
  api_key = openai_chat_config.api_key
1265
- chat_model = conversation_config.chat_model
1269
+ chat_model_name = chat_model.name
1266
1270
  chat_response = converse_openai(
1267
1271
  compiled_references,
1268
1272
  query_to_run,
@@ -1270,13 +1274,13 @@ def generate_chat_response(
1270
1274
  online_results=online_results,
1271
1275
  code_results=code_results,
1272
1276
  conversation_log=meta_log,
1273
- model=chat_model,
1277
+ model=chat_model_name,
1274
1278
  api_key=api_key,
1275
1279
  api_base_url=openai_chat_config.api_base_url,
1276
1280
  completion_func=partial_completion,
1277
1281
  conversation_commands=conversation_commands,
1278
- max_prompt_size=conversation_config.max_prompt_size,
1279
- tokenizer_name=conversation_config.tokenizer,
1282
+ max_prompt_size=chat_model.max_prompt_size,
1283
+ tokenizer_name=chat_model.tokenizer,
1280
1284
  location_data=location_data,
1281
1285
  user_name=user_name,
1282
1286
  agent=agent,
@@ -1288,8 +1292,8 @@ def generate_chat_response(
1288
1292
  tracer=tracer,
1289
1293
  )
1290
1294
 
1291
- elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
1292
- api_key = conversation_config.ai_model_api.api_key
1295
+ elif chat_model.model_type == ChatModel.ModelType.ANTHROPIC:
1296
+ api_key = chat_model.ai_model_api.api_key
1293
1297
  chat_response = converse_anthropic(
1294
1298
  compiled_references,
1295
1299
  query_to_run,
@@ -1297,12 +1301,12 @@ def generate_chat_response(
1297
1301
  online_results=online_results,
1298
1302
  code_results=code_results,
1299
1303
  conversation_log=meta_log,
1300
- model=conversation_config.chat_model,
1304
+ model=chat_model.name,
1301
1305
  api_key=api_key,
1302
1306
  completion_func=partial_completion,
1303
1307
  conversation_commands=conversation_commands,
1304
- max_prompt_size=conversation_config.max_prompt_size,
1305
- tokenizer_name=conversation_config.tokenizer,
1308
+ max_prompt_size=chat_model.max_prompt_size,
1309
+ tokenizer_name=chat_model.tokenizer,
1306
1310
  location_data=location_data,
1307
1311
  user_name=user_name,
1308
1312
  agent=agent,
@@ -1313,20 +1317,20 @@ def generate_chat_response(
1313
1317
  program_execution_context=program_execution_context,
1314
1318
  tracer=tracer,
1315
1319
  )
1316
- elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
1317
- api_key = conversation_config.ai_model_api.api_key
1320
+ elif chat_model.model_type == ChatModel.ModelType.GOOGLE:
1321
+ api_key = chat_model.ai_model_api.api_key
1318
1322
  chat_response = converse_gemini(
1319
1323
  compiled_references,
1320
1324
  query_to_run,
1321
1325
  online_results,
1322
1326
  code_results,
1323
1327
  meta_log,
1324
- model=conversation_config.chat_model,
1328
+ model=chat_model.name,
1325
1329
  api_key=api_key,
1326
1330
  completion_func=partial_completion,
1327
1331
  conversation_commands=conversation_commands,
1328
- max_prompt_size=conversation_config.max_prompt_size,
1329
- tokenizer_name=conversation_config.tokenizer,
1332
+ max_prompt_size=chat_model.max_prompt_size,
1333
+ tokenizer_name=chat_model.tokenizer,
1330
1334
  location_data=location_data,
1331
1335
  user_name=user_name,
1332
1336
  agent=agent,
@@ -1339,7 +1343,7 @@ def generate_chat_response(
1339
1343
  tracer=tracer,
1340
1344
  )
1341
1345
 
1342
- metadata.update({"chat_model": conversation_config.chat_model})
1346
+ metadata.update({"chat_model": chat_model.name})
1343
1347
 
1344
1348
  except Exception as e:
1345
1349
  logger.error(e, exc_info=True)
@@ -1939,13 +1943,13 @@ def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False)
1939
1943
  current_notion_config = get_user_notion_config(user)
1940
1944
  notion_token = current_notion_config.token if current_notion_config else ""
1941
1945
 
1942
- selected_chat_model_config = ConversationAdapters.get_conversation_config(
1946
+ selected_chat_model_config = ConversationAdapters.get_chat_model(
1943
1947
  user
1944
- ) or ConversationAdapters.get_default_conversation_config(user)
1948
+ ) or ConversationAdapters.get_default_chat_model(user)
1945
1949
  chat_models = ConversationAdapters.get_conversation_processor_options().all()
1946
1950
  chat_model_options = list()
1947
1951
  for chat_model in chat_models:
1948
- chat_model_options.append({"name": chat_model.chat_model, "id": chat_model.id})
1952
+ chat_model_options.append({"name": chat_model.name, "id": chat_model.id})
1949
1953
 
1950
1954
  selected_paint_model_config = ConversationAdapters.get_user_text_to_image_model_config(user)
1951
1955
  paint_model_options = ConversationAdapters.get_text_to_image_model_options().all()
khoj/utils/helpers.py CHANGED
@@ -584,13 +584,15 @@ def get_cost_of_chat_message(model_name: str, input_tokens: int = 0, output_toke
584
584
  return input_cost + output_cost + prev_cost
585
585
 
586
586
 
587
- def get_chat_usage_metrics(model_name: str, input_tokens: int = 0, output_tokens: int = 0, usage: dict = {}):
587
+ def get_chat_usage_metrics(
588
+ model_name: str, input_tokens: int = 0, output_tokens: int = 0, usage: dict = {}, cost: float = None
589
+ ):
588
590
  """
589
- Get usage metrics for chat message based on input and output tokens
591
+ Get usage metrics for chat message based on input and output tokens and cost
590
592
  """
591
593
  prev_usage = usage or {"input_tokens": 0, "output_tokens": 0, "cost": 0.0}
592
594
  return {
593
595
  "input_tokens": prev_usage["input_tokens"] + input_tokens,
594
596
  "output_tokens": prev_usage["output_tokens"] + output_tokens,
595
- "cost": get_cost_of_chat_message(model_name, input_tokens, output_tokens, prev_cost=prev_usage["cost"]),
597
+ "cost": cost or get_cost_of_chat_message(model_name, input_tokens, output_tokens, prev_cost=prev_usage["cost"]),
596
598
  }