khoj 1.42.1.dev10__py3-none-any.whl → 1.42.2.dev16__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 (55) hide show
  1. khoj/configure.py +2 -0
  2. khoj/database/adapters/__init__.py +9 -7
  3. khoj/database/models/__init__.py +9 -9
  4. khoj/interface/compiled/404/index.html +2 -2
  5. khoj/interface/compiled/_next/static/chunks/7127-79a3af5138960272.js +1 -0
  6. khoj/interface/compiled/_next/static/chunks/{5138-2cce449fd2454abf.js → 7211-7fedd2ee3655239c.js} +1 -1
  7. khoj/interface/compiled/_next/static/chunks/app/automations/page-ef89ac958e78aa81.js +1 -0
  8. khoj/interface/compiled/_next/static/chunks/app/chat/page-db0fbea54ccea62f.js +1 -0
  9. khoj/interface/compiled/_next/static/chunks/app/share/chat/{page-9a167dc9b5fcd464.js → page-da90c78180a86040.js} +1 -1
  10. khoj/interface/compiled/_next/static/chunks/{webpack-964e8ed3380daff1.js → webpack-0f15e6b51732b337.js} +1 -1
  11. khoj/interface/compiled/_next/static/css/{9c223d337a984468.css → 7017ee76c2f2cd87.css} +1 -1
  12. khoj/interface/compiled/_next/static/css/9a460202d29476e5.css +1 -0
  13. khoj/interface/compiled/agents/index.html +2 -2
  14. khoj/interface/compiled/agents/index.txt +1 -1
  15. khoj/interface/compiled/automations/index.html +2 -2
  16. khoj/interface/compiled/automations/index.txt +2 -2
  17. khoj/interface/compiled/chat/index.html +2 -2
  18. khoj/interface/compiled/chat/index.txt +2 -2
  19. khoj/interface/compiled/index.html +2 -2
  20. khoj/interface/compiled/index.txt +1 -1
  21. khoj/interface/compiled/search/index.html +2 -2
  22. khoj/interface/compiled/search/index.txt +1 -1
  23. khoj/interface/compiled/settings/index.html +2 -2
  24. khoj/interface/compiled/settings/index.txt +1 -1
  25. khoj/interface/compiled/share/chat/index.html +2 -2
  26. khoj/interface/compiled/share/chat/index.txt +2 -2
  27. khoj/processor/conversation/anthropic/anthropic_chat.py +19 -134
  28. khoj/processor/conversation/anthropic/utils.py +1 -1
  29. khoj/processor/conversation/google/gemini_chat.py +20 -141
  30. khoj/processor/conversation/offline/chat_model.py +23 -153
  31. khoj/processor/conversation/openai/gpt.py +14 -128
  32. khoj/processor/conversation/prompts.py +2 -63
  33. khoj/processor/conversation/utils.py +94 -89
  34. khoj/processor/image/generate.py +16 -11
  35. khoj/processor/operator/__init__.py +2 -3
  36. khoj/processor/operator/operator_agent_binary.py +11 -11
  37. khoj/processor/tools/online_search.py +9 -3
  38. khoj/processor/tools/run_code.py +5 -5
  39. khoj/routers/api.py +5 -527
  40. khoj/routers/api_automation.py +243 -0
  41. khoj/routers/api_chat.py +48 -129
  42. khoj/routers/helpers.py +371 -121
  43. khoj/routers/research.py +11 -43
  44. khoj/utils/helpers.py +0 -6
  45. {khoj-1.42.1.dev10.dist-info → khoj-1.42.2.dev16.dist-info}/METADATA +1 -1
  46. {khoj-1.42.1.dev10.dist-info → khoj-1.42.2.dev16.dist-info}/RECORD +51 -50
  47. khoj/interface/compiled/_next/static/chunks/7127-d3199617463d45f0.js +0 -1
  48. khoj/interface/compiled/_next/static/chunks/app/automations/page-465741d9149dfd48.js +0 -1
  49. khoj/interface/compiled/_next/static/chunks/app/chat/page-898079bcea5376f4.js +0 -1
  50. khoj/interface/compiled/_next/static/css/fca983d49c3dd1a3.css +0 -1
  51. /khoj/interface/compiled/_next/static/{2niR8lV9_OpGs1vdb2yMp → OTsOjbrtuaYMukpuJS4sy}/_buildManifest.js +0 -0
  52. /khoj/interface/compiled/_next/static/{2niR8lV9_OpGs1vdb2yMp → OTsOjbrtuaYMukpuJS4sy}/_ssgManifest.js +0 -0
  53. {khoj-1.42.1.dev10.dist-info → khoj-1.42.2.dev16.dist-info}/WHEEL +0 -0
  54. {khoj-1.42.1.dev10.dist-info → khoj-1.42.2.dev16.dist-info}/entry_points.txt +0 -0
  55. {khoj-1.42.1.dev10.dist-info → khoj-1.42.2.dev16.dist-info}/licenses/LICENSE +0 -0
khoj/routers/helpers.py CHANGED
@@ -1,12 +1,13 @@
1
1
  import base64
2
+ import concurrent.futures
2
3
  import hashlib
3
4
  import json
4
5
  import logging
5
6
  import math
6
7
  import os
7
8
  import re
9
+ import time
8
10
  from datetime import datetime, timedelta, timezone
9
- from functools import partial
10
11
  from random import random
11
12
  from typing import (
12
13
  Annotated,
@@ -46,6 +47,7 @@ from khoj.database.adapters import (
46
47
  aget_user_by_email,
47
48
  ais_user_subscribed,
48
49
  create_khoj_token,
50
+ get_default_search_model,
49
51
  get_khoj_tokens,
50
52
  get_user_name,
51
53
  get_user_notion_config,
@@ -55,6 +57,7 @@ from khoj.database.adapters import (
55
57
  )
56
58
  from khoj.database.models import (
57
59
  Agent,
60
+ ChatMessageModel,
58
61
  ChatModel,
59
62
  ClientApplication,
60
63
  Conversation,
@@ -100,12 +103,16 @@ from khoj.processor.conversation.utils import (
100
103
  clean_json,
101
104
  clean_mermaidjs,
102
105
  construct_chat_history,
106
+ construct_question_history,
107
+ defilter_query,
103
108
  generate_chatml_messages_with_context,
104
- save_to_conversation_log,
105
109
  )
106
110
  from khoj.processor.speech.text_to_speech import is_eleven_labs_enabled
107
111
  from khoj.routers.email import is_resend_enabled, send_task_email
108
112
  from khoj.routers.twilio import is_twilio_enabled
113
+ from khoj.search_filter.date_filter import DateFilter
114
+ from khoj.search_filter.file_filter import FileFilter
115
+ from khoj.search_filter.word_filter import WordFilter
109
116
  from khoj.search_type import text_search
110
117
  from khoj.utils import state
111
118
  from khoj.utils.config import OfflineChatProcessorModel
@@ -122,7 +129,13 @@ from khoj.utils.helpers import (
122
129
  timer,
123
130
  tool_descriptions_for_llm,
124
131
  )
125
- from khoj.utils.rawconfig import ChatRequestBody, FileAttachment, FileData, LocationData
132
+ from khoj.utils.rawconfig import (
133
+ ChatRequestBody,
134
+ FileAttachment,
135
+ FileData,
136
+ LocationData,
137
+ SearchResponse,
138
+ )
126
139
 
127
140
  logger = logging.getLogger(__name__)
128
141
 
@@ -236,8 +249,6 @@ def get_next_url(request: Request) -> str:
236
249
  def get_conversation_command(query: str) -> ConversationCommand:
237
250
  if query.startswith("/notes"):
238
251
  return ConversationCommand.Notes
239
- elif query.startswith("/help"):
240
- return ConversationCommand.Help
241
252
  elif query.startswith("/general"):
242
253
  return ConversationCommand.General
243
254
  elif query.startswith("/online"):
@@ -248,8 +259,6 @@ def get_conversation_command(query: str) -> ConversationCommand:
248
259
  return ConversationCommand.Image
249
260
  elif query.startswith("/automated_task"):
250
261
  return ConversationCommand.AutomatedTask
251
- elif query.startswith("/summarize"):
252
- return ConversationCommand.Summarize
253
262
  elif query.startswith("/diagram"):
254
263
  return ConversationCommand.Diagram
255
264
  elif query.startswith("/code"):
@@ -285,7 +294,7 @@ async def acreate_title_from_history(
285
294
  """
286
295
  Create a title from the given conversation history
287
296
  """
288
- chat_history = construct_chat_history(conversation.conversation_log)
297
+ chat_history = construct_chat_history(conversation.messages)
289
298
 
290
299
  title_generation_prompt = prompts.conversation_title_generation.format(chat_history=chat_history)
291
300
 
@@ -345,7 +354,7 @@ async def acheck_if_safe_prompt(system_prompt: str, user: KhojUser = None, lax:
345
354
 
346
355
  async def aget_data_sources_and_output_format(
347
356
  query: str,
348
- conversation_history: dict,
357
+ chat_history: list[ChatMessageModel],
349
358
  is_task: bool,
350
359
  user: KhojUser,
351
360
  query_images: List[str] = None,
@@ -379,14 +388,11 @@ async def aget_data_sources_and_output_format(
379
388
  agent_outputs = agent.output_modes if agent else []
380
389
 
381
390
  for output, description in mode_descriptions_for_llm.items():
382
- # Do not allow tasks to schedule another task
383
- if is_task and output == ConversationCommand.Automation:
384
- continue
385
391
  output_options[output.value] = description
386
392
  if len(agent_outputs) == 0 or output.value in agent_outputs:
387
393
  output_options_str += f'- "{output.value}": "{description}"\n'
388
394
 
389
- chat_history = construct_chat_history(conversation_history, n=6)
395
+ chat_history_str = construct_chat_history(chat_history, n=6)
390
396
 
391
397
  if query_images:
392
398
  query = f"[placeholder for {len(query_images)} user attached images]\n{query}"
@@ -399,7 +405,7 @@ async def aget_data_sources_and_output_format(
399
405
  query=query,
400
406
  sources=source_options_str,
401
407
  outputs=output_options_str,
402
- chat_history=chat_history,
408
+ chat_history=chat_history_str,
403
409
  personality_context=personality_context,
404
410
  )
405
411
 
@@ -462,7 +468,7 @@ async def aget_data_sources_and_output_format(
462
468
  async def infer_webpage_urls(
463
469
  q: str,
464
470
  max_webpages: int,
465
- conversation_history: dict,
471
+ chat_history: List[ChatMessageModel],
466
472
  location_data: LocationData,
467
473
  user: KhojUser,
468
474
  query_images: List[str] = None,
@@ -475,7 +481,7 @@ async def infer_webpage_urls(
475
481
  """
476
482
  location = f"{location_data}" if location_data else "Unknown"
477
483
  username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
478
- chat_history = construct_chat_history(conversation_history)
484
+ chat_history_str = construct_chat_history(chat_history)
479
485
 
480
486
  utc_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
481
487
  personality_context = (
@@ -485,7 +491,7 @@ async def infer_webpage_urls(
485
491
  online_queries_prompt = prompts.infer_webpages_to_read.format(
486
492
  query=q,
487
493
  max_webpages=max_webpages,
488
- chat_history=chat_history,
494
+ chat_history=chat_history_str,
489
495
  current_date=utc_date,
490
496
  location=location,
491
497
  username=username,
@@ -526,7 +532,7 @@ async def infer_webpage_urls(
526
532
 
527
533
  async def generate_online_subqueries(
528
534
  q: str,
529
- conversation_history: dict,
535
+ chat_history: List[ChatMessageModel],
530
536
  location_data: LocationData,
531
537
  user: KhojUser,
532
538
  query_images: List[str] = None,
@@ -540,7 +546,7 @@ async def generate_online_subqueries(
540
546
  """
541
547
  location = f"{location_data}" if location_data else "Unknown"
542
548
  username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
543
- chat_history = construct_chat_history(conversation_history)
549
+ chat_history_str = construct_chat_history(chat_history)
544
550
 
545
551
  utc_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
546
552
  personality_context = (
@@ -549,7 +555,7 @@ async def generate_online_subqueries(
549
555
 
550
556
  online_queries_prompt = prompts.online_search_conversation_subqueries.format(
551
557
  query=q,
552
- chat_history=chat_history,
558
+ chat_history=chat_history_str,
553
559
  max_queries=max_queries,
554
560
  current_date=utc_date,
555
561
  location=location,
@@ -591,16 +597,16 @@ async def generate_online_subqueries(
591
597
 
592
598
 
593
599
  def schedule_query(
594
- q: str, conversation_history: dict, user: KhojUser, query_images: List[str] = None, tracer: dict = {}
600
+ q: str, chat_history: List[ChatMessageModel], user: KhojUser, query_images: List[str] = None, tracer: dict = {}
595
601
  ) -> Tuple[str, str, str]:
596
602
  """
597
603
  Schedule the date, time to run the query. Assume the server timezone is UTC.
598
604
  """
599
- chat_history = construct_chat_history(conversation_history)
605
+ chat_history_str = construct_chat_history(chat_history)
600
606
 
601
607
  crontime_prompt = prompts.crontime_prompt.format(
602
608
  query=q,
603
- chat_history=chat_history,
609
+ chat_history=chat_history_str,
604
610
  )
605
611
 
606
612
  raw_response = send_message_to_model_wrapper_sync(
@@ -619,16 +625,16 @@ def schedule_query(
619
625
 
620
626
 
621
627
  async def aschedule_query(
622
- q: str, conversation_history: dict, user: KhojUser, query_images: List[str] = None, tracer: dict = {}
628
+ q: str, chat_history: List[ChatMessageModel], user: KhojUser, query_images: List[str] = None, tracer: dict = {}
623
629
  ) -> Tuple[str, str, str]:
624
630
  """
625
631
  Schedule the date, time to run the query. Assume the server timezone is UTC.
626
632
  """
627
- chat_history = construct_chat_history(conversation_history)
633
+ chat_history_str = construct_chat_history(chat_history)
628
634
 
629
635
  crontime_prompt = prompts.crontime_prompt.format(
630
636
  query=q,
631
- chat_history=chat_history,
637
+ chat_history=chat_history_str,
632
638
  )
633
639
 
634
640
  raw_response = await send_message_to_model_wrapper(
@@ -681,7 +687,7 @@ async def extract_relevant_info(
681
687
  async def extract_relevant_summary(
682
688
  q: str,
683
689
  corpus: str,
684
- conversation_history: dict,
690
+ chat_history: List[ChatMessageModel] = [],
685
691
  query_images: List[str] = None,
686
692
  user: KhojUser = None,
687
693
  agent: Agent = None,
@@ -698,11 +704,11 @@ async def extract_relevant_summary(
698
704
  prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
699
705
  )
700
706
 
701
- chat_history = construct_chat_history(conversation_history)
707
+ chat_history_str = construct_chat_history(chat_history)
702
708
 
703
709
  extract_relevant_information = prompts.extract_relevant_summary.format(
704
710
  query=q,
705
- chat_history=chat_history,
711
+ chat_history=chat_history_str,
706
712
  corpus=corpus.strip(),
707
713
  personality_context=personality_context,
708
714
  )
@@ -725,7 +731,7 @@ async def generate_summary_from_files(
725
731
  q: str,
726
732
  user: KhojUser,
727
733
  file_filters: List[str],
728
- meta_log: dict,
734
+ chat_history: List[ChatMessageModel] = [],
729
735
  query_images: List[str] = None,
730
736
  agent: Agent = None,
731
737
  send_status_func: Optional[Callable] = None,
@@ -766,7 +772,7 @@ async def generate_summary_from_files(
766
772
  response = await extract_relevant_summary(
767
773
  q,
768
774
  contextual_data,
769
- conversation_history=meta_log,
775
+ chat_history=chat_history,
770
776
  query_images=query_images,
771
777
  user=user,
772
778
  agent=agent,
@@ -782,7 +788,7 @@ async def generate_summary_from_files(
782
788
 
783
789
  async def generate_excalidraw_diagram(
784
790
  q: str,
785
- conversation_history: Dict[str, Any],
791
+ chat_history: List[ChatMessageModel],
786
792
  location_data: LocationData,
787
793
  note_references: List[Dict[str, Any]],
788
794
  online_results: Optional[dict] = None,
@@ -799,7 +805,7 @@ async def generate_excalidraw_diagram(
799
805
 
800
806
  better_diagram_description_prompt = await generate_better_diagram_description(
801
807
  q=q,
802
- conversation_history=conversation_history,
808
+ chat_history=chat_history,
803
809
  location_data=location_data,
804
810
  note_references=note_references,
805
811
  online_results=online_results,
@@ -834,7 +840,7 @@ async def generate_excalidraw_diagram(
834
840
 
835
841
  async def generate_better_diagram_description(
836
842
  q: str,
837
- conversation_history: Dict[str, Any],
843
+ chat_history: List[ChatMessageModel],
838
844
  location_data: LocationData,
839
845
  note_references: List[Dict[str, Any]],
840
846
  online_results: Optional[dict] = None,
@@ -857,7 +863,7 @@ async def generate_better_diagram_description(
857
863
 
858
864
  user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
859
865
 
860
- chat_history = construct_chat_history(conversation_history)
866
+ chat_history_str = construct_chat_history(chat_history)
861
867
 
862
868
  simplified_online_results = {}
863
869
 
@@ -870,7 +876,7 @@ async def generate_better_diagram_description(
870
876
 
871
877
  improve_diagram_description_prompt = prompts.improve_excalidraw_diagram_description_prompt.format(
872
878
  query=q,
873
- chat_history=chat_history,
879
+ chat_history=chat_history_str,
874
880
  location=location,
875
881
  current_date=today_date,
876
882
  references=user_references,
@@ -939,7 +945,7 @@ async def generate_excalidraw_diagram_from_description(
939
945
 
940
946
  async def generate_mermaidjs_diagram(
941
947
  q: str,
942
- conversation_history: Dict[str, Any],
948
+ chat_history: List[ChatMessageModel],
943
949
  location_data: LocationData,
944
950
  note_references: List[Dict[str, Any]],
945
951
  online_results: Optional[dict] = None,
@@ -956,7 +962,7 @@ async def generate_mermaidjs_diagram(
956
962
 
957
963
  better_diagram_description_prompt = await generate_better_mermaidjs_diagram_description(
958
964
  q=q,
959
- conversation_history=conversation_history,
965
+ chat_history=chat_history,
960
966
  location_data=location_data,
961
967
  note_references=note_references,
962
968
  online_results=online_results,
@@ -985,7 +991,7 @@ async def generate_mermaidjs_diagram(
985
991
 
986
992
  async def generate_better_mermaidjs_diagram_description(
987
993
  q: str,
988
- conversation_history: Dict[str, Any],
994
+ chat_history: List[ChatMessageModel],
989
995
  location_data: LocationData,
990
996
  note_references: List[Dict[str, Any]],
991
997
  online_results: Optional[dict] = None,
@@ -1008,7 +1014,7 @@ async def generate_better_mermaidjs_diagram_description(
1008
1014
 
1009
1015
  user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
1010
1016
 
1011
- chat_history = construct_chat_history(conversation_history)
1017
+ chat_history_str = construct_chat_history(chat_history)
1012
1018
 
1013
1019
  simplified_online_results = {}
1014
1020
 
@@ -1021,7 +1027,7 @@ async def generate_better_mermaidjs_diagram_description(
1021
1027
 
1022
1028
  improve_diagram_description_prompt = prompts.improve_mermaid_js_diagram_description_prompt.format(
1023
1029
  query=q,
1024
- chat_history=chat_history,
1030
+ chat_history=chat_history_str,
1025
1031
  location=location,
1026
1032
  current_date=today_date,
1027
1033
  references=user_references,
@@ -1150,6 +1156,276 @@ async def generate_better_image_prompt(
1150
1156
  return response
1151
1157
 
1152
1158
 
1159
+ async def search_documents(
1160
+ user: KhojUser,
1161
+ chat_history: list[ChatMessageModel],
1162
+ q: str,
1163
+ n: int,
1164
+ d: float,
1165
+ conversation_id: str,
1166
+ conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
1167
+ location_data: LocationData = None,
1168
+ send_status_func: Optional[Callable] = None,
1169
+ query_images: Optional[List[str]] = None,
1170
+ previous_inferred_queries: Set = set(),
1171
+ agent: Agent = None,
1172
+ query_files: str = None,
1173
+ tracer: dict = {},
1174
+ ):
1175
+ # Initialize Variables
1176
+ compiled_references: List[dict[str, str]] = []
1177
+ inferred_queries: List[str] = []
1178
+
1179
+ agent_has_entries = False
1180
+
1181
+ if agent:
1182
+ agent_has_entries = await sync_to_async(EntryAdapters.agent_has_entries)(agent=agent)
1183
+
1184
+ if (
1185
+ not ConversationCommand.Notes in conversation_commands
1186
+ and not ConversationCommand.Default in conversation_commands
1187
+ and not agent_has_entries
1188
+ ):
1189
+ yield compiled_references, inferred_queries, q
1190
+ return
1191
+
1192
+ # If Notes or Default is not in the conversation command, then the search should be restricted to the agent's knowledge base
1193
+ should_limit_to_agent_knowledge = (
1194
+ ConversationCommand.Notes not in conversation_commands
1195
+ and ConversationCommand.Default not in conversation_commands
1196
+ )
1197
+
1198
+ if not await sync_to_async(EntryAdapters.user_has_entries)(user=user):
1199
+ if not agent_has_entries:
1200
+ logger.debug("No documents in knowledge base. Use a Khoj client to sync and chat with your docs.")
1201
+ yield compiled_references, inferred_queries, q
1202
+ return
1203
+
1204
+ # Extract filter terms from user message
1205
+ defiltered_query = defilter_query(q)
1206
+ filters_in_query = q.replace(defiltered_query, "").strip()
1207
+ conversation = await sync_to_async(ConversationAdapters.get_conversation_by_id)(conversation_id)
1208
+
1209
+ if not conversation:
1210
+ logger.error(f"Conversation with id {conversation_id} not found when extracting references.")
1211
+ yield compiled_references, inferred_queries, defiltered_query
1212
+ return
1213
+
1214
+ filters_in_query += " ".join([f'file:"{filter}"' for filter in conversation.file_filters])
1215
+ if is_none_or_empty(filters_in_query):
1216
+ logger.debug(f"Filters in query: {filters_in_query}")
1217
+
1218
+ personality_context = prompts.personality_context.format(personality=agent.personality) if agent else ""
1219
+
1220
+ # Infer search queries from user message
1221
+ with timer("Extracting search queries took", logger):
1222
+ inferred_queries = await extract_questions(
1223
+ query=defiltered_query,
1224
+ user=user,
1225
+ personality_context=personality_context,
1226
+ chat_history=chat_history,
1227
+ location_data=location_data,
1228
+ query_images=query_images,
1229
+ query_files=query_files,
1230
+ tracer=tracer,
1231
+ )
1232
+
1233
+ # Collate search results as context for the LLM
1234
+ inferred_queries = list(set(inferred_queries) - previous_inferred_queries)
1235
+ with timer("Searching knowledge base took", logger):
1236
+ search_results = []
1237
+ logger.info(f"🔍 Searching knowledge base with queries: {inferred_queries}")
1238
+ if send_status_func:
1239
+ inferred_queries_str = "\n- " + "\n- ".join(inferred_queries)
1240
+ async for event in send_status_func(f"**Searching Documents for:** {inferred_queries_str}"):
1241
+ yield {ChatEvent.STATUS: event}
1242
+ for query in inferred_queries:
1243
+ search_results.extend(
1244
+ await execute_search(
1245
+ user if not should_limit_to_agent_knowledge else None,
1246
+ f"{query} {filters_in_query}",
1247
+ n=n,
1248
+ t=state.SearchType.All,
1249
+ r=True,
1250
+ max_distance=d,
1251
+ dedupe=False,
1252
+ agent=agent,
1253
+ )
1254
+ )
1255
+ search_results = text_search.deduplicated_search_responses(search_results)
1256
+ compiled_references = [
1257
+ {"query": q, "compiled": item.additional["compiled"], "file": item.additional["file"]}
1258
+ for q, item in zip(inferred_queries, search_results)
1259
+ ]
1260
+
1261
+ yield compiled_references, inferred_queries, defiltered_query
1262
+
1263
+
1264
+ async def extract_questions(
1265
+ query: str,
1266
+ user: KhojUser,
1267
+ personality_context: str = "",
1268
+ chat_history: List[ChatMessageModel] = [],
1269
+ location_data: LocationData = None,
1270
+ query_images: Optional[List[str]] = None,
1271
+ query_files: str = None,
1272
+ tracer: dict = {},
1273
+ ):
1274
+ """
1275
+ Infer document search queries from user message and provided context
1276
+ """
1277
+ # Shared context setup
1278
+ location = f"{location_data}" if location_data else "N/A"
1279
+ username = prompts.user_name.format(name=user.get_full_name()) if user and user.get_full_name() else ""
1280
+
1281
+ # Date variables for prompt formatting
1282
+ today = datetime.today()
1283
+ current_new_year = today.replace(month=1, day=1)
1284
+ last_new_year = current_new_year.replace(year=today.year - 1)
1285
+ yesterday = (today - timedelta(days=1)).strftime("%Y-%m-%d")
1286
+
1287
+ # Common prompt setup for API-based models (using Anthropic prompts for consistency)
1288
+ chat_history_str = construct_question_history(chat_history, query_prefix="User", agent_name="Assistant")
1289
+
1290
+ system_prompt = prompts.extract_questions_system_prompt.format(
1291
+ current_date=today.strftime("%Y-%m-%d"),
1292
+ day_of_week=today.strftime("%A"),
1293
+ current_month=today.strftime("%Y-%m"),
1294
+ last_new_year=last_new_year.strftime("%Y"),
1295
+ last_new_year_date=last_new_year.strftime("%Y-%m-%d"),
1296
+ current_new_year_date=current_new_year.strftime("%Y-%m-%d"),
1297
+ yesterday_date=yesterday,
1298
+ location=location,
1299
+ username=username,
1300
+ personality_context=personality_context,
1301
+ )
1302
+
1303
+ prompt = prompts.extract_questions_user_message.format(text=query, chat_history=chat_history_str)
1304
+
1305
+ class DocumentQueries(BaseModel):
1306
+ """Choose searches to run on user documents."""
1307
+
1308
+ queries: List[str] = Field(..., min_items=1, description="List of search queries to run on user documents.")
1309
+
1310
+ raw_response = await send_message_to_model_wrapper(
1311
+ system_message=system_prompt,
1312
+ query=prompt,
1313
+ query_images=query_images,
1314
+ query_files=query_files,
1315
+ chat_history=chat_history,
1316
+ response_type="json_object",
1317
+ response_schema=DocumentQueries,
1318
+ user=user,
1319
+ tracer=tracer,
1320
+ )
1321
+
1322
+ # Extract questions from the response
1323
+ try:
1324
+ response = clean_json(raw_response)
1325
+ response = pyjson5.loads(response)
1326
+ queries = [q.strip() for q in response["queries"] if q.strip()]
1327
+ if not isinstance(queries, list) or not queries:
1328
+ logger.error(f"Invalid response for constructing subqueries: {response}")
1329
+ return [query]
1330
+ return queries
1331
+ except:
1332
+ logger.warning(f"LLM returned invalid JSON. Falling back to using user message as search query.")
1333
+ return [query]
1334
+
1335
+
1336
+ async def execute_search(
1337
+ user: KhojUser,
1338
+ q: str,
1339
+ n: Optional[int] = 5,
1340
+ t: Optional[state.SearchType] = None,
1341
+ r: Optional[bool] = False,
1342
+ max_distance: Optional[Union[float, None]] = None,
1343
+ dedupe: Optional[bool] = True,
1344
+ agent: Optional[Agent] = None,
1345
+ ):
1346
+ # Run validation checks
1347
+ results: List[SearchResponse] = []
1348
+
1349
+ start_time = time.time()
1350
+
1351
+ # Ensure the agent, if present, is accessible by the user
1352
+ if user and agent and not await AgentAdapters.ais_agent_accessible(agent, user):
1353
+ logger.error(f"Agent {agent.slug} is not accessible by user {user}")
1354
+ return results
1355
+
1356
+ if q is None or q == "":
1357
+ logger.warning(f"No query param (q) passed in API call to initiate search")
1358
+ return results
1359
+
1360
+ # initialize variables
1361
+ user_query = q.strip()
1362
+ results_count = n or 5
1363
+ t = t or state.SearchType.All
1364
+ search_futures: List[concurrent.futures.Future] = []
1365
+
1366
+ # return cached results, if available
1367
+ if user:
1368
+ query_cache_key = f"{user_query}-{n}-{t}-{r}-{max_distance}-{dedupe}"
1369
+ if query_cache_key in state.query_cache[user.uuid]:
1370
+ logger.debug(f"Return response from query cache")
1371
+ return state.query_cache[user.uuid][query_cache_key]
1372
+
1373
+ # Encode query with filter terms removed
1374
+ defiltered_query = user_query
1375
+ for filter in [DateFilter(), WordFilter(), FileFilter()]:
1376
+ defiltered_query = filter.defilter(defiltered_query)
1377
+
1378
+ encoded_asymmetric_query = None
1379
+ if t != state.SearchType.Image:
1380
+ with timer("Encoding query took", logger=logger):
1381
+ search_model = await sync_to_async(get_default_search_model)()
1382
+ encoded_asymmetric_query = state.embeddings_model[search_model.name].embed_query(defiltered_query)
1383
+
1384
+ with concurrent.futures.ThreadPoolExecutor() as executor:
1385
+ if t in [
1386
+ state.SearchType.All,
1387
+ state.SearchType.Org,
1388
+ state.SearchType.Markdown,
1389
+ state.SearchType.Github,
1390
+ state.SearchType.Notion,
1391
+ state.SearchType.Plaintext,
1392
+ state.SearchType.Pdf,
1393
+ ]:
1394
+ # query markdown notes
1395
+ search_futures += [
1396
+ executor.submit(
1397
+ text_search.query,
1398
+ user_query,
1399
+ user,
1400
+ t,
1401
+ question_embedding=encoded_asymmetric_query,
1402
+ max_distance=max_distance,
1403
+ agent=agent,
1404
+ )
1405
+ ]
1406
+
1407
+ # Query across each requested content types in parallel
1408
+ with timer("Query took", logger):
1409
+ for search_future in concurrent.futures.as_completed(search_futures):
1410
+ hits = await search_future.result()
1411
+ # Collate results
1412
+ results += text_search.collate_results(hits, dedupe=dedupe)
1413
+
1414
+ # Sort results across all content types and take top results
1415
+ results = text_search.rerank_and_sort_results(
1416
+ results, query=defiltered_query, rank_results=r, search_model_name=search_model.name
1417
+ )[:results_count]
1418
+
1419
+ # Cache results
1420
+ if user:
1421
+ state.query_cache[user.uuid][query_cache_key] = results
1422
+
1423
+ end_time = time.time()
1424
+ logger.debug(f"🔍 Search took: {end_time - start_time:.3f} seconds")
1425
+
1426
+ return results
1427
+
1428
+
1153
1429
  async def send_message_to_model_wrapper(
1154
1430
  query: str,
1155
1431
  system_message: str = "",
@@ -1160,7 +1436,7 @@ async def send_message_to_model_wrapper(
1160
1436
  query_images: List[str] = None,
1161
1437
  context: str = "",
1162
1438
  query_files: str = None,
1163
- conversation_log: dict = {},
1439
+ chat_history: list[ChatMessageModel] = [],
1164
1440
  agent_chat_model: ChatModel = None,
1165
1441
  tracer: dict = {},
1166
1442
  ):
@@ -1193,7 +1469,7 @@ async def send_message_to_model_wrapper(
1193
1469
  user_message=query,
1194
1470
  context_message=context,
1195
1471
  system_message=system_message,
1196
- conversation_log=conversation_log,
1472
+ chat_history=chat_history,
1197
1473
  model_name=chat_model_name,
1198
1474
  loaded_model=loaded_model,
1199
1475
  tokenizer_name=tokenizer,
@@ -1260,7 +1536,7 @@ def send_message_to_model_wrapper_sync(
1260
1536
  user: KhojUser = None,
1261
1537
  query_images: List[str] = None,
1262
1538
  query_files: str = "",
1263
- conversation_log: dict = {},
1539
+ chat_history: List[ChatMessageModel] = [],
1264
1540
  tracer: dict = {},
1265
1541
  ):
1266
1542
  chat_model: ChatModel = ConversationAdapters.get_default_chat_model(user)
@@ -1284,7 +1560,7 @@ def send_message_to_model_wrapper_sync(
1284
1560
  truncated_messages = generate_chatml_messages_with_context(
1285
1561
  user_message=message,
1286
1562
  system_message=system_message,
1287
- conversation_log=conversation_log,
1563
+ chat_history=chat_history,
1288
1564
  model_name=chat_model_name,
1289
1565
  loaded_model=loaded_model,
1290
1566
  max_prompt_size=max_tokens,
@@ -1342,61 +1618,31 @@ def send_message_to_model_wrapper_sync(
1342
1618
 
1343
1619
  async def agenerate_chat_response(
1344
1620
  q: str,
1345
- meta_log: dict,
1621
+ chat_history: List[ChatMessageModel],
1346
1622
  conversation: Conversation,
1347
1623
  compiled_references: List[Dict] = [],
1348
1624
  online_results: Dict[str, Dict] = {},
1349
1625
  code_results: Dict[str, Dict] = {},
1350
1626
  operator_results: List[OperatorRun] = [],
1351
1627
  research_results: List[ResearchIteration] = [],
1352
- inferred_queries: List[str] = [],
1353
- conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
1354
1628
  user: KhojUser = None,
1355
- client_application: ClientApplication = None,
1356
1629
  location_data: LocationData = None,
1357
1630
  user_name: Optional[str] = None,
1358
1631
  query_images: Optional[List[str]] = None,
1359
- train_of_thought: List[Any] = [],
1360
1632
  query_files: str = None,
1361
- raw_query_files: List[FileAttachment] = None,
1362
- generated_images: List[str] = None,
1363
1633
  raw_generated_files: List[FileAttachment] = [],
1364
- generated_mermaidjs_diagram: str = None,
1365
1634
  program_execution_context: List[str] = [],
1366
1635
  generated_asset_results: Dict[str, Dict] = {},
1367
1636
  is_subscribed: bool = False,
1368
1637
  tracer: dict = {},
1369
- ) -> Tuple[AsyncGenerator[str | ResponseWithThought, None], Dict[str, str]]:
1638
+ ) -> Tuple[AsyncGenerator[ResponseWithThought, None], Dict[str, str]]:
1370
1639
  # Initialize Variables
1371
- chat_response_generator: AsyncGenerator[str | ResponseWithThought, None] = None
1372
- logger.debug(f"Conversation Types: {conversation_commands}")
1640
+ chat_response_generator: AsyncGenerator[ResponseWithThought, None] = None
1373
1641
 
1374
1642
  metadata = {}
1375
1643
  agent = await AgentAdapters.aget_conversation_agent_by_id(conversation.agent.id) if conversation.agent else None
1376
1644
 
1377
1645
  try:
1378
- partial_completion = partial(
1379
- save_to_conversation_log,
1380
- q,
1381
- user=user,
1382
- meta_log=meta_log,
1383
- compiled_references=compiled_references,
1384
- online_results=online_results,
1385
- code_results=code_results,
1386
- operator_results=operator_results,
1387
- research_results=research_results,
1388
- inferred_queries=inferred_queries,
1389
- client_application=client_application,
1390
- conversation_id=str(conversation.id),
1391
- query_images=query_images,
1392
- train_of_thought=train_of_thought,
1393
- raw_query_files=raw_query_files,
1394
- generated_images=generated_images,
1395
- raw_generated_files=raw_generated_files,
1396
- generated_mermaidjs_diagram=generated_mermaidjs_diagram,
1397
- tracer=tracer,
1398
- )
1399
-
1400
1646
  query_to_run = q
1401
1647
  deepthought = False
1402
1648
  if research_results:
@@ -1420,22 +1666,23 @@ async def agenerate_chat_response(
1420
1666
  if chat_model.model_type == "offline":
1421
1667
  loaded_model = state.offline_chat_processor_config.loaded_model
1422
1668
  chat_response_generator = converse_offline(
1669
+ # Query
1423
1670
  user_query=query_to_run,
1671
+ # Context
1424
1672
  references=compiled_references,
1425
1673
  online_results=online_results,
1674
+ generated_files=raw_generated_files,
1675
+ generated_asset_results=generated_asset_results,
1676
+ location_data=location_data,
1677
+ user_name=user_name,
1678
+ query_files=query_files,
1679
+ chat_history=chat_history,
1680
+ # Model
1426
1681
  loaded_model=loaded_model,
1427
- conversation_log=meta_log,
1428
- completion_func=partial_completion,
1429
- conversation_commands=conversation_commands,
1430
1682
  model_name=chat_model.name,
1431
1683
  max_prompt_size=chat_model.max_prompt_size,
1432
1684
  tokenizer_name=chat_model.tokenizer,
1433
- location_data=location_data,
1434
- user_name=user_name,
1435
1685
  agent=agent,
1436
- query_files=query_files,
1437
- generated_files=raw_generated_files,
1438
- generated_asset_results=generated_asset_results,
1439
1686
  tracer=tracer,
1440
1687
  )
1441
1688
 
@@ -1444,28 +1691,29 @@ async def agenerate_chat_response(
1444
1691
  api_key = openai_chat_config.api_key
1445
1692
  chat_model_name = chat_model.name
1446
1693
  chat_response_generator = converse_openai(
1694
+ # Query
1447
1695
  query_to_run,
1448
- compiled_references,
1449
- query_images=query_images,
1696
+ # Context
1697
+ references=compiled_references,
1450
1698
  online_results=online_results,
1451
1699
  code_results=code_results,
1452
1700
  operator_results=operator_results,
1453
- conversation_log=meta_log,
1701
+ query_images=query_images,
1702
+ query_files=query_files,
1703
+ generated_files=raw_generated_files,
1704
+ generated_asset_results=generated_asset_results,
1705
+ program_execution_context=program_execution_context,
1706
+ location_data=location_data,
1707
+ user_name=user_name,
1708
+ chat_history=chat_history,
1709
+ # Model
1454
1710
  model=chat_model_name,
1455
1711
  api_key=api_key,
1456
1712
  api_base_url=openai_chat_config.api_base_url,
1457
- completion_func=partial_completion,
1458
- conversation_commands=conversation_commands,
1459
1713
  max_prompt_size=chat_model.max_prompt_size,
1460
1714
  tokenizer_name=chat_model.tokenizer,
1461
- location_data=location_data,
1462
- user_name=user_name,
1463
1715
  agent=agent,
1464
1716
  vision_available=vision_available,
1465
- query_files=query_files,
1466
- generated_files=raw_generated_files,
1467
- generated_asset_results=generated_asset_results,
1468
- program_execution_context=program_execution_context,
1469
1717
  deepthought=deepthought,
1470
1718
  tracer=tracer,
1471
1719
  )
@@ -1474,28 +1722,29 @@ async def agenerate_chat_response(
1474
1722
  api_key = chat_model.ai_model_api.api_key
1475
1723
  api_base_url = chat_model.ai_model_api.api_base_url
1476
1724
  chat_response_generator = converse_anthropic(
1725
+ # Query
1477
1726
  query_to_run,
1478
- compiled_references,
1479
- query_images=query_images,
1727
+ # Context
1728
+ references=compiled_references,
1480
1729
  online_results=online_results,
1481
1730
  code_results=code_results,
1482
1731
  operator_results=operator_results,
1483
- conversation_log=meta_log,
1732
+ query_images=query_images,
1733
+ query_files=query_files,
1734
+ generated_files=raw_generated_files,
1735
+ generated_asset_results=generated_asset_results,
1736
+ program_execution_context=program_execution_context,
1737
+ location_data=location_data,
1738
+ user_name=user_name,
1739
+ chat_history=chat_history,
1740
+ # Model
1484
1741
  model=chat_model.name,
1485
1742
  api_key=api_key,
1486
1743
  api_base_url=api_base_url,
1487
- completion_func=partial_completion,
1488
- conversation_commands=conversation_commands,
1489
1744
  max_prompt_size=chat_model.max_prompt_size,
1490
1745
  tokenizer_name=chat_model.tokenizer,
1491
- location_data=location_data,
1492
- user_name=user_name,
1493
1746
  agent=agent,
1494
1747
  vision_available=vision_available,
1495
- query_files=query_files,
1496
- generated_files=raw_generated_files,
1497
- generated_asset_results=generated_asset_results,
1498
- program_execution_context=program_execution_context,
1499
1748
  deepthought=deepthought,
1500
1749
  tracer=tracer,
1501
1750
  )
@@ -1503,28 +1752,29 @@ async def agenerate_chat_response(
1503
1752
  api_key = chat_model.ai_model_api.api_key
1504
1753
  api_base_url = chat_model.ai_model_api.api_base_url
1505
1754
  chat_response_generator = converse_gemini(
1755
+ # Query
1506
1756
  query_to_run,
1507
- compiled_references,
1757
+ # Context
1758
+ references=compiled_references,
1508
1759
  online_results=online_results,
1509
1760
  code_results=code_results,
1510
1761
  operator_results=operator_results,
1511
- conversation_log=meta_log,
1762
+ query_images=query_images,
1763
+ query_files=query_files,
1764
+ generated_files=raw_generated_files,
1765
+ generated_asset_results=generated_asset_results,
1766
+ program_execution_context=program_execution_context,
1767
+ location_data=location_data,
1768
+ user_name=user_name,
1769
+ chat_history=chat_history,
1770
+ # Model
1512
1771
  model=chat_model.name,
1513
1772
  api_key=api_key,
1514
1773
  api_base_url=api_base_url,
1515
- completion_func=partial_completion,
1516
- conversation_commands=conversation_commands,
1517
1774
  max_prompt_size=chat_model.max_prompt_size,
1518
1775
  tokenizer_name=chat_model.tokenizer,
1519
- location_data=location_data,
1520
- user_name=user_name,
1521
1776
  agent=agent,
1522
- query_images=query_images,
1523
1777
  vision_available=vision_available,
1524
- query_files=query_files,
1525
- generated_files=raw_generated_files,
1526
- generated_asset_results=generated_asset_results,
1527
- program_execution_context=program_execution_context,
1528
1778
  deepthought=deepthought,
1529
1779
  tracer=tracer,
1530
1780
  )
@@ -2005,11 +2255,11 @@ async def create_automation(
2005
2255
  timezone: str,
2006
2256
  user: KhojUser,
2007
2257
  calling_url: URL,
2008
- meta_log: dict = {},
2258
+ chat_history: List[ChatMessageModel] = [],
2009
2259
  conversation_id: str = None,
2010
2260
  tracer: dict = {},
2011
2261
  ):
2012
- crontime, query_to_run, subject = await aschedule_query(q, meta_log, user, tracer=tracer)
2262
+ crontime, query_to_run, subject = await aschedule_query(q, chat_history, user, tracer=tracer)
2013
2263
  job = await aschedule_automation(query_to_run, subject, crontime, timezone, q, user, calling_url, conversation_id)
2014
2264
  return job, crontime, query_to_run, subject
2015
2265