khoj 1.30.1.dev9__py3-none-any.whl → 1.30.2__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 (98) hide show
  1. khoj/configure.py +25 -0
  2. khoj/database/adapters/__init__.py +1 -1
  3. khoj/database/admin.py +39 -0
  4. khoj/interface/compiled/404/index.html +1 -1
  5. khoj/interface/compiled/_next/static/chunks/1210.ef7a0f9a7e43da1d.js +1 -0
  6. khoj/interface/compiled/_next/static/chunks/1279-4cb23143aa2c0228.js +1 -0
  7. khoj/interface/compiled/_next/static/chunks/1603-ba5f9f05e92c8412.js +1 -0
  8. khoj/interface/compiled/_next/static/chunks/1970-1b63ac1497b03a10.js +1 -0
  9. khoj/interface/compiled/_next/static/chunks/2646-92ba433951d02d52.js +20 -0
  10. khoj/interface/compiled/_next/static/chunks/3072-be830e4f8412b9d2.js +1 -0
  11. khoj/interface/compiled/_next/static/chunks/3463-081c031e873b7966.js +3 -0
  12. khoj/interface/compiled/_next/static/chunks/3690-51312931ba1eae30.js +1 -0
  13. khoj/interface/compiled/_next/static/chunks/{6297-d1c842ed3f714ab0.js → 3717-b46079dbe9f55694.js} +1 -1
  14. khoj/interface/compiled/_next/static/chunks/4504-62ac13e7d94c52f9.js +1 -0
  15. khoj/interface/compiled/_next/static/chunks/4752-554a3db270186ce3.js +1 -0
  16. khoj/interface/compiled/_next/static/chunks/5512-7cc62049bbe60e11.js +1 -0
  17. khoj/interface/compiled/_next/static/chunks/5538-e5f3c9f4d67a64b9.js +1 -0
  18. khoj/interface/compiled/_next/static/chunks/7592-a09c39a38e60634b.js +1 -0
  19. khoj/interface/compiled/_next/static/chunks/8423-1dda16bc56236523.js +1 -0
  20. khoj/interface/compiled/_next/static/chunks/app/agents/{page-4353b1a532795ad1.js → page-f5c0801b27a8e95e.js} +1 -1
  21. khoj/interface/compiled/_next/static/chunks/app/automations/{layout-27c28e923c9b1ff0.js → layout-7f1b79a2c67af0b4.js} +1 -1
  22. khoj/interface/compiled/_next/static/chunks/app/automations/{page-c9f13c865e739607.js → page-0393501fad5f8e3d.js} +1 -1
  23. khoj/interface/compiled/_next/static/chunks/app/chat/{page-2790303dee566590.js → page-f2539e3197d03c0d.js} +1 -1
  24. khoj/interface/compiled/_next/static/chunks/app/{page-e83f8a77f5c3caef.js → page-5cc56a8db5d21b38.js} +1 -1
  25. khoj/interface/compiled/_next/static/chunks/app/search/{page-8e28deacb61f75aa.js → page-e8b578d155550386.js} +1 -1
  26. khoj/interface/compiled/_next/static/chunks/app/settings/{layout-254eaaf916449a60.js → layout-1f4d76a8b09517b1.js} +1 -1
  27. khoj/interface/compiled/_next/static/chunks/app/settings/page-b6c835050c970be7.js +1 -0
  28. khoj/interface/compiled/_next/static/chunks/app/share/chat/{page-07e1a8a345e768de.js → page-635635e4fb39fe29.js} +1 -1
  29. khoj/interface/compiled/_next/static/chunks/webpack-5dbccc5145b80b64.js +1 -0
  30. khoj/interface/compiled/_next/static/css/4cae6c0e5c72fb2d.css +1 -0
  31. khoj/interface/compiled/_next/static/css/5d8d85d3f2e95bae.css +25 -0
  32. khoj/interface/compiled/_next/static/css/63e106a52a0ec4ca.css +1 -0
  33. khoj/interface/compiled/agents/index.html +1 -1
  34. khoj/interface/compiled/agents/index.txt +2 -2
  35. khoj/interface/compiled/automations/index.html +1 -1
  36. khoj/interface/compiled/automations/index.txt +3 -3
  37. khoj/interface/compiled/chat/index.html +1 -1
  38. khoj/interface/compiled/chat/index.txt +2 -2
  39. khoj/interface/compiled/index.html +1 -1
  40. khoj/interface/compiled/index.txt +2 -2
  41. khoj/interface/compiled/search/index.html +1 -1
  42. khoj/interface/compiled/search/index.txt +2 -2
  43. khoj/interface/compiled/settings/index.html +1 -1
  44. khoj/interface/compiled/settings/index.txt +3 -3
  45. khoj/interface/compiled/share/chat/index.html +1 -1
  46. khoj/interface/compiled/share/chat/index.txt +2 -2
  47. khoj/main.py +4 -0
  48. khoj/processor/conversation/anthropic/anthropic_chat.py +8 -2
  49. khoj/processor/conversation/anthropic/utils.py +22 -3
  50. khoj/processor/conversation/google/gemini_chat.py +8 -2
  51. khoj/processor/conversation/google/utils.py +19 -3
  52. khoj/processor/conversation/offline/chat_model.py +12 -4
  53. khoj/processor/conversation/openai/gpt.py +9 -2
  54. khoj/processor/conversation/openai/utils.py +39 -21
  55. khoj/processor/conversation/prompts.py +40 -21
  56. khoj/processor/conversation/utils.py +15 -9
  57. khoj/processor/tools/run_code.py +1 -25
  58. khoj/routers/api_chat.py +41 -16
  59. khoj/routers/api_subscription.py +9 -2
  60. khoj/routers/auth.py +2 -2
  61. khoj/routers/helpers.py +20 -5
  62. khoj/routers/research.py +2 -1
  63. khoj/utils/cli.py +2 -0
  64. khoj/utils/constants.py +17 -0
  65. khoj/utils/helpers.py +55 -1
  66. khoj/utils/state.py +1 -0
  67. {khoj-1.30.1.dev9.dist-info → khoj-1.30.2.dist-info}/METADATA +9 -4
  68. {khoj-1.30.1.dev9.dist-info → khoj-1.30.2.dist-info}/RECORD +77 -79
  69. khoj/interface/compiled/_next/static/chunks/1210.132a7e1910006bbb.js +0 -1
  70. khoj/interface/compiled/_next/static/chunks/1279-f37ee4a388ebf544.js +0 -1
  71. khoj/interface/compiled/_next/static/chunks/1603-dc5fd983dbcd070d.js +0 -1
  72. khoj/interface/compiled/_next/static/chunks/1970-c78f6acc8e16e30b.js +0 -1
  73. khoj/interface/compiled/_next/static/chunks/2261-748f7c327df3c8c1.js +0 -1
  74. khoj/interface/compiled/_next/static/chunks/3062-71ed4b46ac2bb87c.js +0 -1
  75. khoj/interface/compiled/_next/static/chunks/3124-a4cea2eda163128d.js +0 -1
  76. khoj/interface/compiled/_next/static/chunks/3803-d74118a2d0182c52.js +0 -1
  77. khoj/interface/compiled/_next/static/chunks/4504-1629487c8bc82203.js +0 -1
  78. khoj/interface/compiled/_next/static/chunks/5512-94c7c2bbcf58c19d.js +0 -1
  79. khoj/interface/compiled/_next/static/chunks/5538-b87b60ecc0c27ceb.js +0 -1
  80. khoj/interface/compiled/_next/static/chunks/7883-b1305ec254213afe.js +0 -20
  81. khoj/interface/compiled/_next/static/chunks/796-68f9e87f9cdfda1d.js +0 -3
  82. khoj/interface/compiled/_next/static/chunks/8423-c0123d454681e03a.js +0 -1
  83. khoj/interface/compiled/_next/static/chunks/9001-3b27af6d5f21df44.js +0 -1
  84. khoj/interface/compiled/_next/static/chunks/9417-32c4db52ca42e681.js +0 -1
  85. khoj/interface/compiled/_next/static/chunks/app/settings/page-610d33158b233b34.js +0 -1
  86. khoj/interface/compiled/_next/static/chunks/webpack-6e43825796b7dfa6.js +0 -1
  87. khoj/interface/compiled/_next/static/css/2d097a35da6bfe8d.css +0 -1
  88. khoj/interface/compiled/_next/static/css/80bd6301fc657983.css +0 -1
  89. khoj/interface/compiled/_next/static/css/ed437164d77aa600.css +0 -25
  90. /khoj/interface/compiled/_next/static/{iSHZR6RZZ76lh-q0YbAa5 → UR4enQiSbkZKb3SDFX2tx}/_buildManifest.js +0 -0
  91. /khoj/interface/compiled/_next/static/{iSHZR6RZZ76lh-q0YbAa5 → UR4enQiSbkZKb3SDFX2tx}/_ssgManifest.js +0 -0
  92. /khoj/interface/compiled/_next/static/chunks/{4602-8eeb4b76385ad159.js → 4602-460621c3241e0d13.js} +0 -0
  93. /khoj/interface/compiled/_next/static/chunks/{7023-a5bf5744d19b3bd3.js → 7023-e8de2bded4df6539.js} +0 -0
  94. /khoj/interface/compiled/_next/static/chunks/app/_not-found/{page-07ff4ab42b07845e.js → page-cfba071f5a657256.js} +0 -0
  95. /khoj/interface/compiled/_next/static/chunks/{fd9d1056-2b978342deb60015.js → fd9d1056-2e6c8140e79afc3b.js} +0 -0
  96. {khoj-1.30.1.dev9.dist-info → khoj-1.30.2.dist-info}/WHEEL +0 -0
  97. {khoj-1.30.1.dev9.dist-info → khoj-1.30.2.dist-info}/entry_points.txt +0 -0
  98. {khoj-1.30.1.dev9.dist-info → khoj-1.30.2.dist-info}/licenses/LICENSE +0 -0
@@ -183,20 +183,23 @@ Improved Prompt:
183
183
 
184
184
  improve_diagram_description_prompt = PromptTemplate.from_template(
185
185
  """
186
- you are an architect working with a novice artist using a diagramming tool.
186
+ you are an architect working with a novice digital artist using a diagramming software.
187
187
  {personality_context}
188
188
 
189
189
  you need to convert the user's query to a description format that the novice artist can use very well. you are allowed to use primitives like
190
190
  - text
191
191
  - rectangle
192
- - diamond
193
192
  - ellipse
194
193
  - line
195
194
  - arrow
196
195
 
197
196
  use these primitives to describe what sort of diagram the drawer should create. the artist must recreate the diagram every time, so include all relevant prior information in your description.
198
197
 
199
- use simple, concise language.
198
+ - include the full, exact description. the artist does not have much experience, so be precise.
199
+ - describe the layout.
200
+ - you can only use straight lines.
201
+ - use simple, concise language.
202
+ - keep it simple and easy to understand. the artist is easily distracted.
200
203
 
201
204
  Today's Date: {current_date}
202
205
  User's Location: {location}
@@ -218,19 +221,23 @@ Query: {query}
218
221
 
219
222
  excalidraw_diagram_generation_prompt = PromptTemplate.from_template(
220
223
  """
221
- You are a program manager with the ability to describe diagrams to compose in professional, fine detail.
224
+ You are a program manager with the ability to describe diagrams to compose in professional, fine detail. You LOVE getting into the details and making tedious labels, lines, and shapes look beautiful. You make everything look perfect.
222
225
  {personality_context}
223
226
 
224
- You need to create a declarative description of the diagram and relevant components, using this base schema. Use the `label` property to specify the text to be rendered in the respective elements. Always use light colors for the `backgroundColor` property, like white, or light blue, green, red. "type", "x", "y", "id", are required properties for all elements.
227
+ You need to create a declarative description of the diagram and relevant components, using this base schema.
228
+ - `label`: specify the text to be rendered in the respective elements.
229
+ - Always use light colors for the `backgroundColor` property, like white, or light blue, green, red
230
+ - **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
231
+ - Be very generous with spacing and composition. Use ample space between elements.
225
232
 
226
233
  {{
227
234
  type: string,
228
235
  x: number,
229
236
  y: number,
230
- strokeColor: string,
231
- backgroundColor: string,
232
237
  width: number,
233
238
  height: number,
239
+ strokeColor: string,
240
+ backgroundColor: string,
234
241
  id: string,
235
242
  label: {{
236
243
  text: string,
@@ -240,28 +247,30 @@ You need to create a declarative description of the diagram and relevant compone
240
247
  Valid types:
241
248
  - text
242
249
  - rectangle
243
- - diamond
244
250
  - ellipse
245
251
  - line
246
252
  - arrow
247
253
 
248
- For arrows and lines, you can use the `points` property to specify the start and end points of the arrow. You may also use the `label` property to specify the text to be rendered. You may use the `start` and `end` properties to connect the linear elements to other elements. The start and end point can either be the ID to map to an existing object, or the `type` to create a new object. Mapping to an existing object is useful if you want to connect it to multiple objects. Lines and arrows can only start and end at rectangle, text, diamond, or ellipse elements.
254
+ For arrows and lines,
255
+ - `points`: specify the start and end points of the arrow
256
+ - **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
257
+ - `start` and `end` properties: connect the linear elements to other elements. The start and end point can either be the ID to map to an existing object, or the `type` and `text` to create a new object. Mapping to an existing object is useful if you want to connect it to multiple objects. Lines and arrows can only start and end at rectangle, text, or ellipse elements. Even if you're using the `start` and `end` properties, you still need to specify the `x` and `y` properties for the start and end points.
249
258
 
250
259
  {{
251
260
  type: "arrow",
252
261
  id: string,
253
262
  x: number,
254
263
  y: number,
255
- width: number,
256
- height: number,
257
264
  strokeColor: string,
258
265
  start: {{
259
266
  id: string,
260
267
  type: string,
268
+ text: string,
261
269
  }},
262
270
  end: {{
263
271
  id: string,
264
272
  type: string,
273
+ text: string,
265
274
  }},
266
275
  label: {{
267
276
  text: string,
@@ -272,7 +281,11 @@ For arrows and lines, you can use the `points` property to specify the start and
272
281
  ]
273
282
  }}
274
283
 
275
- For text, you must use the `text` property to specify the text to be rendered. You may also use `fontSize` property to specify the font size of the text. Only use the `text` element for titles, subtitles, and overviews. For labels, use the `label` property in the respective elements.
284
+ For text,
285
+ - `text`: specify the text to be rendered
286
+ - **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
287
+ - `fontSize`: optional property to specify the font size of the text
288
+ - Use this element only for titles, subtitles, and overviews. For labels, use the `label` property in the respective elements.
276
289
 
277
290
  {{
278
291
  type: "text",
@@ -287,19 +300,25 @@ Here's an example of a valid diagram:
287
300
 
288
301
  Design Description: Create a diagram describing a circular development process with 3 stages: design, implementation and feedback. The design stage is connected to the implementation stage and the implementation stage is connected to the feedback stage and the feedback stage is connected to the design stage. Each stage should be labeled with the stage name.
289
302
 
290
- Response:
291
-
292
- [
293
- {{"type":"text","x":-150,"y":50,"width":300,"height":40,"id":"title_text","text":"Circular Development Process","fontSize":24}},
294
- {{"type":"ellipse","x":-169,"y":113,"width":188,"height":202,"id":"design_ellipse", "label": {{"text": "Design"}}}},
295
- {{"type":"ellipse","x":62,"y":394,"width":186,"height":188,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
296
- {{"type":"ellipse","x":-348,"y":430,"width":184,"height":170,"id":"feedback_ellipse", "label": {{"text": "Feedback"}}}},
303
+ Example Response:
304
+ ```json
305
+ {{
306
+ "scratchpad": "The diagram represents a circular development process with 3 stages: design, implementation and feedback. Each stage is connected to the next stage using an arrow, forming a circular process.",
307
+ "elements": [
308
+ {{"type":"text","x":-150,"y":50,"id":"title_text","text":"Circular Development Process","fontSize":24}},
309
+ {{"type":"ellipse","x":-169,"y":113,"id":"design_ellipse", "label": {{"text": "Design"}}}},
310
+ {{"type":"ellipse","x":62,"y":394,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
311
+ {{"type":"ellipse","x":-348,"y":430,"id":"feedback_ellipse", "label": {{"text": "Feedback"}}}},
297
312
  {{"type":"arrow","x":21,"y":273,"id":"design_to_implement_arrow","points":[[0,0],[86,105]],"start":{{"id":"design_ellipse"}}, "end":{{"id":"implement_ellipse"}}}},
298
313
  {{"type":"arrow","x":50,"y":519,"id":"implement_to_feedback_arrow","points":[[0,0],[-198,-6]],"start":{{"id":"implement_ellipse"}}, "end":{{"id":"feedback_ellipse"}}}},
299
314
  {{"type":"arrow","x":-228,"y":417,"id":"feedback_to_design_arrow","points":[[0,0],[85,-123]],"start":{{"id":"feedback_ellipse"}}, "end":{{"id":"design_ellipse"}}}},
300
- ]
315
+ ]
316
+ }}
317
+ ```
318
+
319
+ Think about spacing and composition. Use ample space between elements. Double the amount of space you think you need. Create a detailed diagram from the provided context and user prompt below.
301
320
 
302
- Create a detailed diagram from the provided context and user prompt below. Return a valid JSON object:
321
+ Return a valid JSON object, where the drawing is in `elements` and your thought process is in `scratchpad`. If you can't make the whole diagram in one response, you can split it into multiple responses. If you need to simplify for brevity, simply do so in the `scratchpad` field. DO NOT add additional info in the `elements` field.
303
322
 
304
323
  Diagram Description: {query}
305
324
 
@@ -5,7 +5,6 @@ import math
5
5
  import mimetypes
6
6
  import os
7
7
  import queue
8
- import re
9
8
  import uuid
10
9
  from dataclasses import dataclass
11
10
  from datetime import datetime
@@ -35,6 +34,7 @@ from khoj.utils.helpers import (
35
34
  ConversationCommand,
36
35
  in_debug_mode,
37
36
  is_none_or_empty,
37
+ is_promptrace_enabled,
38
38
  merge_dicts,
39
39
  )
40
40
  from khoj.utils.rawconfig import FileAttachment
@@ -57,7 +57,7 @@ model_to_prompt_size = {
57
57
  "gemini-1.5-flash": 20000,
58
58
  "gemini-1.5-pro": 20000,
59
59
  # Anthropic Models
60
- "claude-3-5-sonnet-20240620": 20000,
60
+ "claude-3-5-sonnet-20241022": 20000,
61
61
  "claude-3-5-haiku-20241022": 20000,
62
62
  # Offline Models
63
63
  "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": 20000,
@@ -213,6 +213,8 @@ class ChatEvent(Enum):
213
213
  REFERENCES = "references"
214
214
  STATUS = "status"
215
215
  METADATA = "metadata"
216
+ USAGE = "usage"
217
+ END_RESPONSE = "end_response"
216
218
 
217
219
 
218
220
  def message_to_log(
@@ -291,7 +293,7 @@ def save_to_conversation_log(
291
293
  user_message=q,
292
294
  )
293
295
 
294
- if in_debug_mode() or state.verbose > 1:
296
+ if is_promptrace_enabled():
295
297
  merge_message_into_conversation_trace(q, chat_response, tracer)
296
298
 
297
299
  logger.info(
@@ -578,7 +580,7 @@ def commit_conversation_trace(
578
580
  response: str | list[dict],
579
581
  tracer: dict,
580
582
  system_message: str | list[dict] = "",
581
- repo_path: str = "/tmp/promptrace",
583
+ repo_path: str = None,
582
584
  ) -> str:
583
585
  """
584
586
  Save trace of conversation step using git. Useful to visualize, compare and debug traces.
@@ -589,6 +591,11 @@ def commit_conversation_trace(
589
591
  except ImportError:
590
592
  return None
591
593
 
594
+ # Infer repository path from environment variable or provided path
595
+ repo_path = repo_path if not is_none_or_empty(repo_path) else os.getenv("PROMPTRACE_DIR")
596
+ if not repo_path:
597
+ return None
598
+
592
599
  # Serialize session, system message and response to yaml
593
600
  system_message_yaml = json.dumps(system_message, ensure_ascii=False, sort_keys=False)
594
601
  response_yaml = json.dumps(response, ensure_ascii=False, sort_keys=False)
@@ -601,9 +608,6 @@ def commit_conversation_trace(
601
608
  # Extract chat metadata for session
602
609
  uid, cid, mid = tracer.get("uid", "main"), tracer.get("cid", "main"), tracer.get("mid")
603
610
 
604
- # Infer repository path from environment variable or provided path
605
- repo_path = os.getenv("PROMPTRACE_DIR", repo_path)
606
-
607
611
  try:
608
612
  # Prepare git repository
609
613
  os.makedirs(repo_path, exist_ok=True)
@@ -683,7 +687,7 @@ Metadata
683
687
  return None
684
688
 
685
689
 
686
- def merge_message_into_conversation_trace(query: str, response: str, tracer: dict, repo_path="/tmp/promptrace") -> bool:
690
+ def merge_message_into_conversation_trace(query: str, response: str, tracer: dict, repo_path=None) -> bool:
687
691
  """
688
692
  Merge the message branch into its parent conversation branch.
689
693
 
@@ -706,7 +710,9 @@ def merge_message_into_conversation_trace(query: str, response: str, tracer: dic
706
710
  conv_branch = f"c_{tracer['cid']}"
707
711
 
708
712
  # Infer repository path from environment variable or provided path
709
- repo_path = os.getenv("PROMPTRACE_DIR", repo_path)
713
+ repo_path = repo_path if not is_none_or_empty(repo_path) else os.getenv("PROMPTRACE_DIR")
714
+ if not repo_path:
715
+ return None
710
716
  repo = Repo(repo_path)
711
717
 
712
718
  # Checkout conversation branch
@@ -1,5 +1,4 @@
1
1
  import base64
2
- import copy
3
2
  import datetime
4
3
  import json
5
4
  import logging
@@ -20,7 +19,7 @@ from khoj.processor.conversation.utils import (
20
19
  construct_chat_history,
21
20
  )
22
21
  from khoj.routers.helpers import send_message_to_model_wrapper
23
- from khoj.utils.helpers import is_none_or_empty, timer
22
+ from khoj.utils.helpers import is_none_or_empty, timer, truncate_code_context
24
23
  from khoj.utils.rawconfig import LocationData
25
24
 
26
25
  logger = logging.getLogger(__name__)
@@ -180,26 +179,3 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
180
179
  "std_err": f"Failed to execute code with {response.status}",
181
180
  "output_files": [],
182
181
  }
183
-
184
-
185
- def truncate_code_context(original_code_results: dict[str, Any], max_chars=10000) -> dict[str, Any]:
186
- """
187
- Truncate large output files and drop image file data from code results.
188
- """
189
- # Create a deep copy of the code results to avoid modifying the original data
190
- code_results = copy.deepcopy(original_code_results)
191
- for code_result in code_results.values():
192
- for idx, output_file in enumerate(code_result["results"]["output_files"]):
193
- # Drop image files from code results
194
- if Path(output_file["filename"]).suffix in {".png", ".jpg", ".jpeg", ".webp"}:
195
- code_result["results"]["output_files"][idx] = {
196
- "filename": output_file["filename"],
197
- "b64_data": "[placeholder for generated image data for brevity]",
198
- }
199
- # Truncate large output files
200
- elif len(output_file["b64_data"]) > max_chars:
201
- code_result["results"]["output_files"][idx] = {
202
- "filename": output_file["filename"],
203
- "b64_data": output_file["b64_data"][:max_chars] + "...",
204
- }
205
- return code_results
khoj/routers/api_chat.py CHANGED
@@ -432,7 +432,15 @@ def chat_sessions(
432
432
  conversations = conversations[:8]
433
433
 
434
434
  sessions = conversations.values_list(
435
- "id", "slug", "title", "agent__slug", "agent__name", "created_at", "updated_at"
435
+ "id",
436
+ "slug",
437
+ "title",
438
+ "agent__slug",
439
+ "agent__name",
440
+ "created_at",
441
+ "updated_at",
442
+ "agent__style_icon",
443
+ "agent__style_color",
436
444
  )
437
445
 
438
446
  session_values = [
@@ -442,6 +450,8 @@ def chat_sessions(
442
450
  "agent_name": session[4],
443
451
  "created": session[5].strftime("%Y-%m-%d %H:%M:%S"),
444
452
  "updated": session[6].strftime("%Y-%m-%d %H:%M:%S"),
453
+ "agent_icon": session[7],
454
+ "agent_color": session[8],
445
455
  }
446
456
  for session in sessions
447
457
  ]
@@ -667,27 +677,37 @@ async def chat(
667
677
  finally:
668
678
  yield event_delimiter
669
679
 
670
- async def send_llm_response(response: str):
680
+ async def send_llm_response(response: str, usage: dict = None):
681
+ # Send Chat Response
671
682
  async for result in send_event(ChatEvent.START_LLM_RESPONSE, ""):
672
683
  yield result
673
684
  async for result in send_event(ChatEvent.MESSAGE, response):
674
685
  yield result
675
686
  async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
676
687
  yield result
688
+ # Send Usage Metadata once llm interactions are complete
689
+ if usage:
690
+ async for event in send_event(ChatEvent.USAGE, usage):
691
+ yield event
692
+ async for result in send_event(ChatEvent.END_RESPONSE, ""):
693
+ yield result
677
694
 
678
695
  def collect_telemetry():
679
696
  # Gather chat response telemetry
680
697
  nonlocal chat_metadata
681
698
  latency = time.perf_counter() - start_time
682
699
  cmd_set = set([cmd.value for cmd in conversation_commands])
700
+ cost = (tracer.get("usage", {}) or {}).get("cost", 0)
683
701
  chat_metadata = chat_metadata or {}
684
702
  chat_metadata["conversation_command"] = cmd_set
685
- chat_metadata["agent"] = conversation.agent.slug if conversation.agent else None
703
+ chat_metadata["agent"] = conversation.agent.slug if conversation and conversation.agent else None
686
704
  chat_metadata["latency"] = f"{latency:.3f}"
687
705
  chat_metadata["ttft_latency"] = f"{ttft:.3f}"
706
+ chat_metadata["usage"] = tracer.get("usage")
688
707
 
689
708
  logger.info(f"Chat response time to first token: {ttft:.3f} seconds")
690
709
  logger.info(f"Chat response total time: {latency:.3f} seconds")
710
+ logger.info(f"Chat response cost: ${cost:.5f}")
691
711
  update_telemetry_state(
692
712
  request=request,
693
713
  telemetry_type="api",
@@ -699,7 +719,7 @@ async def chat(
699
719
  )
700
720
 
701
721
  if is_query_empty(q):
702
- async for result in send_llm_response("Please ask your query to get started."):
722
+ async for result in send_llm_response("Please ask your query to get started.", tracer.get("usage")):
703
723
  yield result
704
724
  return
705
725
 
@@ -713,7 +733,7 @@ async def chat(
713
733
  create_new=body.create_new,
714
734
  )
715
735
  if not conversation:
716
- async for result in send_llm_response(f"Conversation {conversation_id} not found"):
736
+ async for result in send_llm_response(f"Conversation {conversation_id} not found", tracer.get("usage")):
717
737
  yield result
718
738
  return
719
739
  conversation_id = conversation.id
@@ -777,7 +797,7 @@ async def chat(
777
797
  await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
778
798
  q = q.replace(f"/{cmd.value}", "").strip()
779
799
  except HTTPException as e:
780
- async for result in send_llm_response(str(e.detail)):
800
+ async for result in send_llm_response(str(e.detail), tracer.get("usage")):
781
801
  yield result
782
802
  return
783
803
 
@@ -834,7 +854,7 @@ async def chat(
834
854
  agent_has_entries = await EntryAdapters.aagent_has_entries(agent)
835
855
  if len(file_filters) == 0 and not agent_has_entries:
836
856
  response_log = "No files selected for summarization. Please add files using the section on the left."
837
- async for result in send_llm_response(response_log):
857
+ async for result in send_llm_response(response_log, tracer.get("usage")):
838
858
  yield result
839
859
  else:
840
860
  async for response in generate_summary_from_files(
@@ -853,7 +873,7 @@ async def chat(
853
873
  else:
854
874
  if isinstance(response, str):
855
875
  response_log = response
856
- async for result in send_llm_response(response):
876
+ async for result in send_llm_response(response, tracer.get("usage")):
857
877
  yield result
858
878
 
859
879
  await sync_to_async(save_to_conversation_log)(
@@ -880,7 +900,7 @@ async def chat(
880
900
  conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
881
901
  model_type = conversation_config.model_type
882
902
  formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
883
- async for result in send_llm_response(formatted_help):
903
+ async for result in send_llm_response(formatted_help, tracer.get("usage")):
884
904
  yield result
885
905
  return
886
906
  # Adding specification to search online specifically on khoj.dev pages.
@@ -895,7 +915,7 @@ async def chat(
895
915
  except Exception as e:
896
916
  logger.error(f"Error scheduling task {q} for {user.email}: {e}")
897
917
  error_message = f"Unable to create automation. Ensure the automation doesn't already exist."
898
- async for result in send_llm_response(error_message):
918
+ async for result in send_llm_response(error_message, tracer.get("usage")):
899
919
  yield result
900
920
  return
901
921
 
@@ -916,7 +936,7 @@ async def chat(
916
936
  raw_query_files=raw_query_files,
917
937
  tracer=tracer,
918
938
  )
919
- async for result in send_llm_response(llm_response):
939
+ async for result in send_llm_response(llm_response, tracer.get("usage")):
920
940
  yield result
921
941
  return
922
942
 
@@ -963,7 +983,7 @@ async def chat(
963
983
  yield result
964
984
 
965
985
  if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
966
- async for result in send_llm_response(f"{no_entries_found.format()}"):
986
+ async for result in send_llm_response(f"{no_entries_found.format()}", tracer.get("usage")):
967
987
  yield result
968
988
  return
969
989
 
@@ -1105,7 +1125,7 @@ async def chat(
1105
1125
  "detail": improved_image_prompt,
1106
1126
  "image": None,
1107
1127
  }
1108
- async for result in send_llm_response(json.dumps(content_obj)):
1128
+ async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
1109
1129
  yield result
1110
1130
  return
1111
1131
 
@@ -1132,7 +1152,7 @@ async def chat(
1132
1152
  "inferredQueries": [improved_image_prompt],
1133
1153
  "image": generated_image,
1134
1154
  }
1135
- async for result in send_llm_response(json.dumps(content_obj)):
1155
+ async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
1136
1156
  yield result
1137
1157
  return
1138
1158
 
@@ -1166,7 +1186,7 @@ async def chat(
1166
1186
  diagram_description = excalidraw_diagram_description
1167
1187
  else:
1168
1188
  error_message = "Failed to generate diagram. Please try again later."
1169
- async for result in send_llm_response(error_message):
1189
+ async for result in send_llm_response(error_message, tracer.get("usage")):
1170
1190
  yield result
1171
1191
 
1172
1192
  await sync_to_async(save_to_conversation_log)(
@@ -1213,7 +1233,7 @@ async def chat(
1213
1233
  tracer=tracer,
1214
1234
  )
1215
1235
 
1216
- async for result in send_llm_response(json.dumps(content_obj)):
1236
+ async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
1217
1237
  yield result
1218
1238
  return
1219
1239
 
@@ -1252,6 +1272,11 @@ async def chat(
1252
1272
  if item is None:
1253
1273
  async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
1254
1274
  yield result
1275
+ # Send Usage Metadata once llm interactions are complete
1276
+ async for event in send_event(ChatEvent.USAGE, tracer.get("usage")):
1277
+ yield event
1278
+ async for result in send_event(ChatEvent.END_RESPONSE, ""):
1279
+ yield result
1255
1280
  logger.debug("Finished streaming response")
1256
1281
  return
1257
1282
  if not connection_alive or not continue_stream:
@@ -66,16 +66,23 @@ async def subscribe(request: Request):
66
66
  success = user is not None
67
67
  elif event_type in {"customer.subscription.updated"}:
68
68
  user_subscription = await sync_to_async(adapters.get_user_subscription)(customer_email)
69
+
70
+ renewal_date = None
71
+ if subscription["current_period_end"]:
72
+ renewal_date = datetime.fromtimestamp(subscription["current_period_end"], tz=timezone.utc)
73
+
69
74
  # Allow updating subscription status if paid user
70
75
  if user_subscription and user_subscription.renewal_date:
71
76
  # Mark user as unsubscribed or resubscribed
72
77
  is_recurring = not subscription["cancel_at_period_end"]
73
- user, is_new = await adapters.set_user_subscription(customer_email, is_recurring=is_recurring)
78
+ user, is_new = await adapters.set_user_subscription(
79
+ customer_email, is_recurring=is_recurring, renewal_date=renewal_date
80
+ )
74
81
  success = user is not None
75
82
  elif event_type in {"customer.subscription.deleted"}:
76
83
  # Reset the user to trial state
77
84
  user, is_new = await adapters.set_user_subscription(
78
- customer_email, is_recurring=False, renewal_date=False, type=Subscription.Type.TRIAL
85
+ customer_email, is_recurring=False, renewal_date=None, type=Subscription.Type.TRIAL
79
86
  )
80
87
  success = user is not None
81
88
 
khoj/routers/auth.py CHANGED
@@ -89,7 +89,7 @@ async def login_magic_link(request: Request, form: MagicLinkForm):
89
89
  update_telemetry_state(
90
90
  request=request,
91
91
  telemetry_type="api",
92
- api="create_user",
92
+ api="create_user__email",
93
93
  metadata={"server_id": str(user.uuid)},
94
94
  )
95
95
  logger.log(logging.INFO, f"🥳 New User Created: {user.uuid}")
@@ -174,7 +174,7 @@ async def auth(request: Request):
174
174
  update_telemetry_state(
175
175
  request=request,
176
176
  telemetry_type="api",
177
- api="create_user",
177
+ api="create_user__google",
178
178
  metadata={"server_id": str(khoj_user.uuid)},
179
179
  )
180
180
  logger.log(logging.INFO, f"🥳 New User Created: {khoj_user.uuid}")
khoj/routers/helpers.py CHANGED
@@ -411,7 +411,7 @@ async def aget_data_sources_and_output_format(
411
411
  f"Invalid response for determining relevant tools: {selected_sources}. Raw Response: {response}"
412
412
  )
413
413
 
414
- result: Dict = {"sources": [], "output": None} if not is_task else {"output": ConversationCommand.AutomatedTask}
414
+ result: Dict = {"sources": [], "output": None if not is_task else ConversationCommand.AutomatedTask}
415
415
  for selected_source in selected_sources:
416
416
  # Add a double check to verify it's in the agent list, because the LLM sometimes gets confused by the tool options.
417
417
  if (
@@ -753,7 +753,11 @@ async def generate_excalidraw_diagram(
753
753
  yield None, None
754
754
  return
755
755
 
756
- yield better_diagram_description_prompt, excalidraw_diagram_description
756
+ scratchpad = excalidraw_diagram_description.get("scratchpad")
757
+
758
+ inferred_queries = f"Instruction: {better_diagram_description_prompt}\n\nScratchpad: {scratchpad}"
759
+
760
+ yield inferred_queries, excalidraw_diagram_description.get("elements")
757
761
 
758
762
 
759
763
  async def generate_better_diagram_description(
@@ -822,7 +826,7 @@ async def generate_excalidraw_diagram_from_description(
822
826
  user: KhojUser = None,
823
827
  agent: Agent = None,
824
828
  tracer: dict = {},
825
- ) -> str:
829
+ ) -> Dict[str, Any]:
826
830
  personality_context = (
827
831
  prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
828
832
  )
@@ -838,10 +842,18 @@ async def generate_excalidraw_diagram_from_description(
838
842
  )
839
843
  raw_response = clean_json(raw_response)
840
844
  try:
845
+ # Expect response to have `elements` and `scratchpad` keys
841
846
  response: Dict[str, str] = json.loads(raw_response)
847
+ if (
848
+ not response
849
+ or not isinstance(response, Dict)
850
+ or not response.get("elements")
851
+ or not response.get("scratchpad")
852
+ ):
853
+ raise AssertionError(f"Invalid response for generating Excalidraw diagram: {response}")
842
854
  except Exception:
843
855
  raise AssertionError(f"Invalid response for generating Excalidraw diagram: {raw_response}")
844
- if not response or not isinstance(response, List) or not isinstance(response[0], Dict):
856
+ if not response or not isinstance(response["elements"], List) or not isinstance(response["elements"][0], Dict):
845
857
  # TODO Some additional validation here that it's a valid Excalidraw diagram
846
858
  raise AssertionError(f"Invalid response for improving diagram description: {response}")
847
859
 
@@ -1770,6 +1782,7 @@ Manage your automations [here](/automations).
1770
1782
  class MessageProcessor:
1771
1783
  def __init__(self):
1772
1784
  self.references = {}
1785
+ self.usage = {}
1773
1786
  self.raw_response = ""
1774
1787
 
1775
1788
  def convert_message_chunk_to_json(self, raw_chunk: str) -> Dict[str, Any]:
@@ -1793,6 +1806,8 @@ class MessageProcessor:
1793
1806
  chunk_type = ChatEvent(chunk["type"])
1794
1807
  if chunk_type == ChatEvent.REFERENCES:
1795
1808
  self.references = chunk["data"]
1809
+ elif chunk_type == ChatEvent.USAGE:
1810
+ self.usage = chunk["data"]
1796
1811
  elif chunk_type == ChatEvent.MESSAGE:
1797
1812
  chunk_data = chunk["data"]
1798
1813
  if isinstance(chunk_data, dict):
@@ -1837,7 +1852,7 @@ async def read_chat_stream(response_iterator: AsyncGenerator[str, None]) -> Dict
1837
1852
  if buffer:
1838
1853
  processor.process_message_chunk(buffer)
1839
1854
 
1840
- return {"response": processor.raw_response, "references": processor.references}
1855
+ return {"response": processor.raw_response, "references": processor.references, "usage": processor.usage}
1841
1856
 
1842
1857
 
1843
1858
  def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False):
khoj/routers/research.py CHANGED
@@ -16,7 +16,7 @@ from khoj.processor.conversation.utils import (
16
16
  construct_tool_chat_history,
17
17
  )
18
18
  from khoj.processor.tools.online_search import read_webpages, search_online
19
- from khoj.processor.tools.run_code import run_code, truncate_code_context
19
+ from khoj.processor.tools.run_code import run_code
20
20
  from khoj.routers.api import extract_references_and_questions
21
21
  from khoj.routers.helpers import (
22
22
  ChatEvent,
@@ -28,6 +28,7 @@ from khoj.utils.helpers import (
28
28
  function_calling_description_for_llm,
29
29
  is_none_or_empty,
30
30
  timer,
31
+ truncate_code_context,
31
32
  )
32
33
  from khoj.utils.rawconfig import LocationData
33
34
 
khoj/utils/cli.py CHANGED
@@ -40,6 +40,8 @@ def cli(args=None):
40
40
  type=pathlib.Path,
41
41
  help="Path to UNIX socket for server. Use to run server behind reverse proxy. Default: /tmp/uvicorn.sock",
42
42
  )
43
+ parser.add_argument("--sslcert", type=str, help="Path to SSL certificate file")
44
+ parser.add_argument("--sslkey", type=str, help="Path to SSL key file")
43
45
  parser.add_argument("--version", "-V", action="store_true", help="Print the installed Khoj version and exit")
44
46
  parser.add_argument(
45
47
  "--disable-chat-on-gpu", action="store_true", default=False, help="Disable using GPU for the offline chat model"
khoj/utils/constants.py CHANGED
@@ -1,4 +1,5 @@
1
1
  from pathlib import Path
2
+ from typing import Dict
2
3
 
3
4
  app_root_directory = Path(__file__).parent.parent.parent
4
5
  web_directory = app_root_directory / "khoj/interface/web/"
@@ -31,3 +32,19 @@ default_config = {
31
32
  "image": {"encoder": "sentence-transformers/clip-ViT-B-32", "model_directory": "~/.khoj/search/image/"},
32
33
  },
33
34
  }
35
+
36
+ model_to_cost: Dict[str, Dict[str, float]] = {
37
+ # OpenAI Pricing: https://openai.com/api/pricing/
38
+ "gpt-4o": {"input": 2.50, "output": 10.00},
39
+ "gpt-4o-mini": {"input": 0.15, "output": 0.60},
40
+ "o1-preview": {"input": 15.0, "output": 60.00},
41
+ "o1-mini": {"input": 3.0, "output": 12.0},
42
+ # Gemini Pricing: https://ai.google.dev/pricing
43
+ "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
44
+ "gemini-1.5-flash-002": {"input": 0.075, "output": 0.30},
45
+ "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
46
+ "gemini-1.5-pro-002": {"input": 1.25, "output": 5.00},
47
+ # Anthropic Pricing: https://www.anthropic.com/pricing#anthropic-api_
48
+ "claude-3-5-sonnet-20241022": {"input": 3.0, "output": 15.0},
49
+ "claude-3-5-haiku-20241022": {"input": 1.0, "output": 5.0},
50
+ }