khoj 1.30.1.dev9__py3-none-any.whl → 1.30.2__py3-none-any.whl

khoj/processor/conversation/prompts.py

@@ -183,20 +183,23 @@ Improved Prompt:
 
 improve_diagram_description_prompt = PromptTemplate.from_template(
     """
-you are an architect working with a novice artist using a diagramming tool.
+you are an architect working with a novice digital artist using a diagramming software.
 {personality_context}
 
 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
 - text
 - rectangle
-- diamond
 - ellipse
 - line
 - arrow
 
 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.
 
-use simple, concise language.
+- include the full, exact description. the artist does not have much experience, so be precise.
+- describe the layout.
+- you can only use straight lines.
+- use simple, concise language.
+- keep it simple and easy to understand. the artist is easily distracted.
 
 Today's Date: {current_date}
 User's Location: {location}
@@ -218,19 +221,23 @@ Query: {query}
 
 excalidraw_diagram_generation_prompt = PromptTemplate.from_template(
     """
-You are a program manager with the ability to describe diagrams to compose in professional, fine detail.
+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.
 {personality_context}
 
-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.
+You need to create a declarative description of the diagram and relevant components, using this base schema.
+- `label`: 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
+- **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
+- Be very generous with spacing and composition. Use ample space between elements.
 
 {{
     type: string,
     x: number,
     y: number,
-    strokeColor: string,
-    backgroundColor: string,
     width: number,
     height: number,
+    strokeColor: string,
+    backgroundColor: string,
     id: string,
     label: {{
         text: string,
@@ -240,28 +247,30 @@ You need to create a declarative description of the diagram and relevant compone
 Valid types:
 - text
 - rectangle
-- diamond
 - ellipse
 - line
 - arrow
 
-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.
+For arrows and lines,
+- `points`: specify the start and end points of the arrow
+- **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
+- `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.
 
 {{
     type: "arrow",
     id: string,
     x: number,
     y: number,
-    width: number,
-    height: number,
     strokeColor: string,
     start: {{
         id: string,
         type: string,
+        text: string,
     }},
     end: {{
         id: string,
         type: string,
+        text: string,
     }},
     label: {{
         text: string,
@@ -272,7 +281,11 @@ For arrows and lines, you can use the `points` property to specify the start and
     ]
 }}
 
-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.
+For text,
+- `text`: specify the text to be rendered
+- **ALWAYS Required properties for ALL elements**: `type`, `x`, `y`, `id`.
+- `fontSize`: optional property to specify the font size of the text
+- Use this element only for titles, subtitles, and overviews. For labels, use the `label` property in the respective elements.
 
 {{
     type: "text",
@@ -287,19 +300,25 @@ Here's an example of a valid diagram:
 
 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.
 
-Response:
-
-[
-    {{"type":"text","x":-150,"y":50,"width":300,"height":40,"id":"title_text","text":"Circular Development Process","fontSize":24}},
-    {{"type":"ellipse","x":-169,"y":113,"width":188,"height":202,"id":"design_ellipse", "label": {{"text": "Design"}}}},
-    {{"type":"ellipse","x":62,"y":394,"width":186,"height":188,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
-    {{"type":"ellipse","x":-348,"y":430,"width":184,"height":170,"id":"feedback_ellipse", "label": {{"text": "Feedback"}}}},
+Example Response:
+```json
+{{
+    "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.",
+    "elements": [
+    {{"type":"text","x":-150,"y":50,"id":"title_text","text":"Circular Development Process","fontSize":24}},
+    {{"type":"ellipse","x":-169,"y":113,"id":"design_ellipse", "label": {{"text": "Design"}}}},
+    {{"type":"ellipse","x":62,"y":394,"id":"implement_ellipse", "label": {{"text": "Implement"}}}},
+    {{"type":"ellipse","x":-348,"y":430,"id":"feedback_ellipse", "label": {{"text": "Feedback"}}}},
     {{"type":"arrow","x":21,"y":273,"id":"design_to_implement_arrow","points":[[0,0],[86,105]],"start":{{"id":"design_ellipse"}}, "end":{{"id":"implement_ellipse"}}}},
     {{"type":"arrow","x":50,"y":519,"id":"implement_to_feedback_arrow","points":[[0,0],[-198,-6]],"start":{{"id":"implement_ellipse"}}, "end":{{"id":"feedback_ellipse"}}}},
     {{"type":"arrow","x":-228,"y":417,"id":"feedback_to_design_arrow","points":[[0,0],[85,-123]],"start":{{"id":"feedback_ellipse"}}, "end":{{"id":"design_ellipse"}}}},
-]
+    ]
+}}
+```
+
+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.
 
-Create a detailed diagram from the provided context and user prompt below. Return a valid JSON object:
+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.
 
 Diagram Description: {query}
 

khoj/processor/conversation/utils.py

@@ -5,7 +5,6 @@ import math
 import mimetypes
 import os
 import queue
-import re
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
@@ -35,6 +34,7 @@ from khoj.utils.helpers import (
     ConversationCommand,
     in_debug_mode,
     is_none_or_empty,
+    is_promptrace_enabled,
     merge_dicts,
 )
 from khoj.utils.rawconfig import FileAttachment
@@ -57,7 +57,7 @@ model_to_prompt_size = {
     "gemini-1.5-flash": 20000,
     "gemini-1.5-pro": 20000,
     # Anthropic Models
-    "claude-3-5-sonnet-20240620": 20000,
+    "claude-3-5-sonnet-20241022": 20000,
     "claude-3-5-haiku-20241022": 20000,
     # Offline Models
     "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": 20000,
@@ -213,6 +213,8 @@ class ChatEvent(Enum):
     REFERENCES = "references"
     STATUS = "status"
     METADATA = "metadata"
+    USAGE = "usage"
+    END_RESPONSE = "end_response"
 
 
 def message_to_log(
@@ -291,7 +293,7 @@ def save_to_conversation_log(
         user_message=q,
     )
 
-    if in_debug_mode() or state.verbose > 1:
+    if is_promptrace_enabled():
         merge_message_into_conversation_trace(q, chat_response, tracer)
 
     logger.info(
@@ -578,7 +580,7 @@ def commit_conversation_trace(
     response: str | list[dict],
     tracer: dict,
     system_message: str | list[dict] = "",
-    repo_path: str = "/tmp/promptrace",
+    repo_path: str = None,
 ) -> str:
     """
     Save trace of conversation step using git. Useful to visualize, compare and debug traces.
@@ -589,6 +591,11 @@ def commit_conversation_trace(
     except ImportError:
         return None
 
+    # Infer repository path from environment variable or provided path
+    repo_path = repo_path if not is_none_or_empty(repo_path) else os.getenv("PROMPTRACE_DIR")
+    if not repo_path:
+        return None
+
     # Serialize session, system message and response to yaml
     system_message_yaml = json.dumps(system_message, ensure_ascii=False, sort_keys=False)
     response_yaml = json.dumps(response, ensure_ascii=False, sort_keys=False)
@@ -601,9 +608,6 @@ def commit_conversation_trace(
     # Extract chat metadata for session
     uid, cid, mid = tracer.get("uid", "main"), tracer.get("cid", "main"), tracer.get("mid")
 
-    # Infer repository path from environment variable or provided path
-    repo_path = os.getenv("PROMPTRACE_DIR", repo_path)
-
     try:
         # Prepare git repository
         os.makedirs(repo_path, exist_ok=True)
@@ -683,7 +687,7 @@ Metadata
         return None
 
 
-def merge_message_into_conversation_trace(query: str, response: str, tracer: dict, repo_path="/tmp/promptrace") -> bool:
+def merge_message_into_conversation_trace(query: str, response: str, tracer: dict, repo_path=None) -> bool:
     """
     Merge the message branch into its parent conversation branch.
 
@@ -706,7 +710,9 @@ def merge_message_into_conversation_trace(query: str, response: str, tracer: dic
         conv_branch = f"c_{tracer['cid']}"
 
         # Infer repository path from environment variable or provided path
-        repo_path = os.getenv("PROMPTRACE_DIR", repo_path)
+        repo_path = repo_path if not is_none_or_empty(repo_path) else os.getenv("PROMPTRACE_DIR")
+        if not repo_path:
+            return None
         repo = Repo(repo_path)
 
         # Checkout conversation branch

khoj/processor/tools/run_code.py

@@ -1,5 +1,4 @@
 import base64
-import copy
 import datetime
 import json
 import logging
@@ -20,7 +19,7 @@ from khoj.processor.conversation.utils import (
     construct_chat_history,
 )
 from khoj.routers.helpers import send_message_to_model_wrapper
-from khoj.utils.helpers import is_none_or_empty, timer
+from khoj.utils.helpers import is_none_or_empty, timer, truncate_code_context
 from khoj.utils.rawconfig import LocationData
 
 logger = logging.getLogger(__name__)
@@ -180,26 +179,3 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
                     "std_err": f"Failed to execute code with {response.status}",
                     "output_files": [],
                 }
-
-
-def truncate_code_context(original_code_results: dict[str, Any], max_chars=10000) -> dict[str, Any]:
-    """
-    Truncate large output files and drop image file data from code results.
-    """
-    # Create a deep copy of the code results to avoid modifying the original data
-    code_results = copy.deepcopy(original_code_results)
-    for code_result in code_results.values():
-        for idx, output_file in enumerate(code_result["results"]["output_files"]):
-            # Drop image files from code results
-            if Path(output_file["filename"]).suffix in {".png", ".jpg", ".jpeg", ".webp"}:
-                code_result["results"]["output_files"][idx] = {
-                    "filename": output_file["filename"],
-                    "b64_data": "[placeholder for generated image data for brevity]",
-                }
-            # Truncate large output files
-            elif len(output_file["b64_data"]) > max_chars:
-                code_result["results"]["output_files"][idx] = {
-                    "filename": output_file["filename"],
-                    "b64_data": output_file["b64_data"][:max_chars] + "...",
-                }
-    return code_results

khoj/routers/api_chat.py

@@ -432,7 +432,15 @@ def chat_sessions(
         conversations = conversations[:8]
 
     sessions = conversations.values_list(
-        "id", "slug", "title", "agent__slug", "agent__name", "created_at", "updated_at"
+        "id",
+        "slug",
+        "title",
+        "agent__slug",
+        "agent__name",
+        "created_at",
+        "updated_at",
+        "agent__style_icon",
+        "agent__style_color",
     )
 
     session_values = [
@@ -442,6 +450,8 @@ def chat_sessions(
             "agent_name": session[4],
             "created": session[5].strftime("%Y-%m-%d %H:%M:%S"),
             "updated": session[6].strftime("%Y-%m-%d %H:%M:%S"),
+            "agent_icon": session[7],
+            "agent_color": session[8],
         }
         for session in sessions
     ]
@@ -667,27 +677,37 @@ async def chat(
             finally:
                 yield event_delimiter
 
-        async def send_llm_response(response: str):
+        async def send_llm_response(response: str, usage: dict = None):
+            # Send Chat Response
             async for result in send_event(ChatEvent.START_LLM_RESPONSE, ""):
                 yield result
             async for result in send_event(ChatEvent.MESSAGE, response):
                 yield result
             async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
                 yield result
+            # Send Usage Metadata once llm interactions are complete
+            if usage:
+                async for event in send_event(ChatEvent.USAGE, usage):
+                    yield event
+            async for result in send_event(ChatEvent.END_RESPONSE, ""):
+                yield result
 
         def collect_telemetry():
             # Gather chat response telemetry
             nonlocal chat_metadata
             latency = time.perf_counter() - start_time
             cmd_set = set([cmd.value for cmd in conversation_commands])
+            cost = (tracer.get("usage", {}) or {}).get("cost", 0)
             chat_metadata = chat_metadata or {}
             chat_metadata["conversation_command"] = cmd_set
-            chat_metadata["agent"] = conversation.agent.slug if conversation.agent else None
+            chat_metadata["agent"] = conversation.agent.slug if conversation and conversation.agent else None
             chat_metadata["latency"] = f"{latency:.3f}"
             chat_metadata["ttft_latency"] = f"{ttft:.3f}"
+            chat_metadata["usage"] = tracer.get("usage")
 
             logger.info(f"Chat response time to first token: {ttft:.3f} seconds")
             logger.info(f"Chat response total time: {latency:.3f} seconds")
+            logger.info(f"Chat response cost: ${cost:.5f}")
             update_telemetry_state(
                 request=request,
                 telemetry_type="api",
@@ -699,7 +719,7 @@ async def chat(
             )
 
         if is_query_empty(q):
-            async for result in send_llm_response("Please ask your query to get started."):
+            async for result in send_llm_response("Please ask your query to get started.", tracer.get("usage")):
                 yield result
             return
 
@@ -713,7 +733,7 @@ async def chat(
             create_new=body.create_new,
         )
         if not conversation:
-            async for result in send_llm_response(f"Conversation {conversation_id} not found"):
+            async for result in send_llm_response(f"Conversation {conversation_id} not found", tracer.get("usage")):
                 yield result
             return
         conversation_id = conversation.id
@@ -777,7 +797,7 @@ async def chat(
                 await conversation_command_rate_limiter.update_and_check_if_valid(request, cmd)
                 q = q.replace(f"/{cmd.value}", "").strip()
             except HTTPException as e:
-                async for result in send_llm_response(str(e.detail)):
+                async for result in send_llm_response(str(e.detail), tracer.get("usage")):
                     yield result
                 return
 
@@ -834,7 +854,7 @@ async def chat(
             agent_has_entries = await EntryAdapters.aagent_has_entries(agent)
             if len(file_filters) == 0 and not agent_has_entries:
                 response_log = "No files selected for summarization. Please add files using the section on the left."
-                async for result in send_llm_response(response_log):
+                async for result in send_llm_response(response_log, tracer.get("usage")):
                     yield result
             else:
                 async for response in generate_summary_from_files(
@@ -853,7 +873,7 @@ async def chat(
                     else:
                         if isinstance(response, str):
                             response_log = response
-                            async for result in send_llm_response(response):
+                            async for result in send_llm_response(response, tracer.get("usage")):
                                 yield result
 
             await sync_to_async(save_to_conversation_log)(
@@ -880,7 +900,7 @@ async def chat(
                     conversation_config = await ConversationAdapters.aget_default_conversation_config(user)
                 model_type = conversation_config.model_type
                 formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
-                async for result in send_llm_response(formatted_help):
+                async for result in send_llm_response(formatted_help, tracer.get("usage")):
                     yield result
                 return
             # Adding specification to search online specifically on khoj.dev pages.
@@ -895,7 +915,7 @@ async def chat(
             except Exception as e:
                 logger.error(f"Error scheduling task {q} for {user.email}: {e}")
                 error_message = f"Unable to create automation. Ensure the automation doesn't already exist."
-                async for result in send_llm_response(error_message):
+                async for result in send_llm_response(error_message, tracer.get("usage")):
                     yield result
                 return
 
@@ -916,7 +936,7 @@ async def chat(
                 raw_query_files=raw_query_files,
                 tracer=tracer,
             )
-            async for result in send_llm_response(llm_response):
+            async for result in send_llm_response(llm_response, tracer.get("usage")):
                 yield result
             return
 
@@ -963,7 +983,7 @@ async def chat(
                     yield result
 
             if conversation_commands == [ConversationCommand.Notes] and not await EntryAdapters.auser_has_entries(user):
-                async for result in send_llm_response(f"{no_entries_found.format()}"):
+                async for result in send_llm_response(f"{no_entries_found.format()}", tracer.get("usage")):
                     yield result
                 return
 
@@ -1105,7 +1125,7 @@ async def chat(
                     "detail": improved_image_prompt,
                     "image": None,
                 }
-                async for result in send_llm_response(json.dumps(content_obj)):
+                async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
                     yield result
                 return
 
@@ -1132,7 +1152,7 @@ async def chat(
                 "inferredQueries": [improved_image_prompt],
                 "image": generated_image,
             }
-            async for result in send_llm_response(json.dumps(content_obj)):
+            async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
                 yield result
             return
 
@@ -1166,7 +1186,7 @@ async def chat(
                         diagram_description = excalidraw_diagram_description
                     else:
                         error_message = "Failed to generate diagram. Please try again later."
-                        async for result in send_llm_response(error_message):
+                        async for result in send_llm_response(error_message, tracer.get("usage")):
                             yield result
 
                         await sync_to_async(save_to_conversation_log)(
@@ -1213,7 +1233,7 @@ async def chat(
                 tracer=tracer,
             )
 
-            async for result in send_llm_response(json.dumps(content_obj)):
+            async for result in send_llm_response(json.dumps(content_obj), tracer.get("usage")):
                 yield result
             return
 
@@ -1252,6 +1272,11 @@ async def chat(
             if item is None:
                 async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):
                     yield result
+                # Send Usage Metadata once llm interactions are complete
+                async for event in send_event(ChatEvent.USAGE, tracer.get("usage")):
+                    yield event
+                async for result in send_event(ChatEvent.END_RESPONSE, ""):
+                    yield result
                 logger.debug("Finished streaming response")
                 return
             if not connection_alive or not continue_stream:

khoj/routers/api_subscription.py

@@ -66,16 +66,23 @@ async def subscribe(request: Request):
         success = user is not None
     elif event_type in {"customer.subscription.updated"}:
         user_subscription = await sync_to_async(adapters.get_user_subscription)(customer_email)
+
+        renewal_date = None
+        if subscription["current_period_end"]:
+            renewal_date = datetime.fromtimestamp(subscription["current_period_end"], tz=timezone.utc)
+
         # Allow updating subscription status if paid user
         if user_subscription and user_subscription.renewal_date:
             # Mark user as unsubscribed or resubscribed
             is_recurring = not subscription["cancel_at_period_end"]
-            user, is_new = await adapters.set_user_subscription(customer_email, is_recurring=is_recurring)
+            user, is_new = await adapters.set_user_subscription(
+                customer_email, is_recurring=is_recurring, renewal_date=renewal_date
+            )
             success = user is not None
     elif event_type in {"customer.subscription.deleted"}:
         # Reset the user to trial state
         user, is_new = await adapters.set_user_subscription(
-            customer_email, is_recurring=False, renewal_date=False, type=Subscription.Type.TRIAL
+            customer_email, is_recurring=False, renewal_date=None, type=Subscription.Type.TRIAL
         )
         success = user is not None
 

khoj/routers/auth.py

@@ -89,7 +89,7 @@ async def login_magic_link(request: Request, form: MagicLinkForm):
             update_telemetry_state(
                 request=request,
                 telemetry_type="api",
-                api="create_user",
+                api="create_user__email",
                 metadata={"server_id": str(user.uuid)},
             )
             logger.log(logging.INFO, f"🥳 New User Created: {user.uuid}")
@@ -174,7 +174,7 @@ async def auth(request: Request):
             update_telemetry_state(
                 request=request,
                 telemetry_type="api",
-                api="create_user",
+                api="create_user__google",
                 metadata={"server_id": str(khoj_user.uuid)},
             )
             logger.log(logging.INFO, f"🥳 New User Created: {khoj_user.uuid}")

khoj/routers/helpers.py

@@ -411,7 +411,7 @@ async def aget_data_sources_and_output_format(
                 f"Invalid response for determining relevant tools: {selected_sources}. Raw Response: {response}"
             )
 
-        result: Dict = {"sources": [], "output": None} if not is_task else {"output": ConversationCommand.AutomatedTask}
+        result: Dict = {"sources": [], "output": None if not is_task else ConversationCommand.AutomatedTask}
         for selected_source in selected_sources:
             # Add a double check to verify it's in the agent list, because the LLM sometimes gets confused by the tool options.
             if (
@@ -753,7 +753,11 @@ async def generate_excalidraw_diagram(
         yield None, None
         return
 
-    yield better_diagram_description_prompt, excalidraw_diagram_description
+    scratchpad = excalidraw_diagram_description.get("scratchpad")
+
+    inferred_queries = f"Instruction: {better_diagram_description_prompt}\n\nScratchpad: {scratchpad}"
+
+    yield inferred_queries, excalidraw_diagram_description.get("elements")
 
 
 async def generate_better_diagram_description(
@@ -822,7 +826,7 @@ async def generate_excalidraw_diagram_from_description(
     user: KhojUser = None,
     agent: Agent = None,
     tracer: dict = {},
-) -> str:
+) -> Dict[str, Any]:
     personality_context = (
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
     )
@@ -838,10 +842,18 @@ async def generate_excalidraw_diagram_from_description(
         )
         raw_response = clean_json(raw_response)
         try:
+            # Expect response to have `elements` and `scratchpad` keys
             response: Dict[str, str] = json.loads(raw_response)
+            if (
+                not response
+                or not isinstance(response, Dict)
+                or not response.get("elements")
+                or not response.get("scratchpad")
+            ):
+                raise AssertionError(f"Invalid response for generating Excalidraw diagram: {response}")
         except Exception:
             raise AssertionError(f"Invalid response for generating Excalidraw diagram: {raw_response}")
-        if not response or not isinstance(response, List) or not isinstance(response[0], Dict):
+        if not response or not isinstance(response["elements"], List) or not isinstance(response["elements"][0], Dict):
             # TODO Some additional validation here that it's a valid Excalidraw diagram
             raise AssertionError(f"Invalid response for improving diagram description: {response}")
 
@@ -1770,6 +1782,7 @@ Manage your automations [here](/automations).
 class MessageProcessor:
     def __init__(self):
         self.references = {}
+        self.usage = {}
         self.raw_response = ""
 
     def convert_message_chunk_to_json(self, raw_chunk: str) -> Dict[str, Any]:
@@ -1793,6 +1806,8 @@ class MessageProcessor:
         chunk_type = ChatEvent(chunk["type"])
         if chunk_type == ChatEvent.REFERENCES:
             self.references = chunk["data"]
+        elif chunk_type == ChatEvent.USAGE:
+            self.usage = chunk["data"]
         elif chunk_type == ChatEvent.MESSAGE:
             chunk_data = chunk["data"]
             if isinstance(chunk_data, dict):
@@ -1837,7 +1852,7 @@ async def read_chat_stream(response_iterator: AsyncGenerator[str, None]) -> Dict
     if buffer:
         processor.process_message_chunk(buffer)
 
-    return {"response": processor.raw_response, "references": processor.references}
+    return {"response": processor.raw_response, "references": processor.references, "usage": processor.usage}
 
 
 def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False):

khoj/routers/research.py

@@ -16,7 +16,7 @@ from khoj.processor.conversation.utils import (
     construct_tool_chat_history,
 )
 from khoj.processor.tools.online_search import read_webpages, search_online
-from khoj.processor.tools.run_code import run_code, truncate_code_context
+from khoj.processor.tools.run_code import run_code
 from khoj.routers.api import extract_references_and_questions
 from khoj.routers.helpers import (
     ChatEvent,
@@ -28,6 +28,7 @@ from khoj.utils.helpers import (
     function_calling_description_for_llm,
     is_none_or_empty,
     timer,
+    truncate_code_context,
 )
 from khoj.utils.rawconfig import LocationData
 

khoj/utils/cli.py

@@ -40,6 +40,8 @@ def cli(args=None):
         type=pathlib.Path,
         help="Path to UNIX socket for server. Use to run server behind reverse proxy. Default: /tmp/uvicorn.sock",
     )
+    parser.add_argument("--sslcert", type=str, help="Path to SSL certificate file")
+    parser.add_argument("--sslkey", type=str, help="Path to SSL key file")
     parser.add_argument("--version", "-V", action="store_true", help="Print the installed Khoj version and exit")
     parser.add_argument(
         "--disable-chat-on-gpu", action="store_true", default=False, help="Disable using GPU for the offline chat model"

khoj/utils/constants.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from typing import Dict
 
 app_root_directory = Path(__file__).parent.parent.parent
 web_directory = app_root_directory / "khoj/interface/web/"
@@ -31,3 +32,19 @@ default_config = {
         "image": {"encoder": "sentence-transformers/clip-ViT-B-32", "model_directory": "~/.khoj/search/image/"},
     },
 }
+
+model_to_cost: Dict[str, Dict[str, float]] = {
+    # OpenAI Pricing: https://openai.com/api/pricing/
+    "gpt-4o": {"input": 2.50, "output": 10.00},
+    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
+    "o1-preview": {"input": 15.0, "output": 60.00},
+    "o1-mini": {"input": 3.0, "output": 12.0},
+    # Gemini Pricing: https://ai.google.dev/pricing
+    "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
+    "gemini-1.5-flash-002": {"input": 0.075, "output": 0.30},
+    "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
+    "gemini-1.5-pro-002": {"input": 1.25, "output": 5.00},
+    # Anthropic Pricing: https://www.anthropic.com/pricing#anthropic-api_
+    "claude-3-5-sonnet-20241022": {"input": 3.0, "output": 15.0},
+    "claude-3-5-haiku-20241022": {"input": 1.0, "output": 5.0},
+}