khoj 1.29.2.dev35__py3-none-any.whl → 1.30.2.dev22__py3-none-any.whl

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
@@ -49,8 +48,6 @@ except ImportError:
 
 model_to_prompt_size = {
     # OpenAI Models
-    "gpt-3.5-turbo": 12000,
-    "gpt-4-turbo-preview": 20000,
     "gpt-4o": 20000,
     "gpt-4o-mini": 20000,
     "o1-preview": 20000,
@@ -59,12 +56,15 @@ model_to_prompt_size = {
     "gemini-1.5-flash": 20000,
     "gemini-1.5-pro": 20000,
     # Anthropic Models
-    "claude-3-5-sonnet-20240620": 20000,
-    "claude-3-opus-20240229": 20000,
+    "claude-3-5-sonnet-20241022": 20000,
+    "claude-3-5-haiku-20241022": 20000,
     # Offline Models
-    "TheBloke/Mistral-7B-Instruct-v0.2-GGUF": 3500,
-    "NousResearch/Hermes-2-Pro-Mistral-7B-GGUF": 3500,
     "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": 20000,
+    "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF": 20000,
+    "bartowski/Llama-3.2-3B-Instruct-GGUF": 20000,
+    "bartowski/gemma-2-9b-it-GGUF": 6000,
+    "bartowski/gemma-2-2b-it-GGUF": 6000,
+    "Qwen/Qwen2.5-14B-Instruct-GGUF": 20000,
 }
 model_to_tokenizer: Dict[str, str] = {}
 
@@ -212,6 +212,8 @@ class ChatEvent(Enum):
     REFERENCES = "references"
     STATUS = "status"
     METADATA = "metadata"
+    USAGE = "usage"
+    END_RESPONSE = "end_response"
 
 
 def message_to_log(
@@ -290,7 +292,7 @@ def save_to_conversation_log(
         user_message=q,
     )
 
-    if in_debug_mode() or state.verbose > 1:
+    if os.getenv("PROMPTRACE_DIR"):
         merge_message_into_conversation_trace(q, chat_response, tracer)
 
     logger.info(
@@ -577,7 +579,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.
@@ -588,6 +590,11 @@ def commit_conversation_trace(
     except ImportError:
         return None
 
+    # Infer repository path from environment variable or provided path
+    repo_path = repo_path or 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)
@@ -600,9 +607,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)
@@ -740,6 +744,20 @@ Metadata
 
 def messages_to_print(messages: list[ChatMessage], max_length: int = 70) -> str:
     """
-    Format, truncate messages to print
+    Format and truncate messages to print, ensuring JSON serializable content
     """
-    return "\n".join([f"{json.dumps(message.content)[:max_length]}..." for message in messages])
+
+    def safe_serialize(content: Any) -> str:
+        try:
+            # Try JSON serialization
+            json.dumps(content)
+            return content
+        except (TypeError, json.JSONDecodeError):
+            # Handle non-serializable types
+            if hasattr(content, "format") and content.format == "WEBP":
+                return "[WebP Image]"
+            elif hasattr(content, "__dict__"):
+                return str(content.__dict__)
+            return str(content)
+
+    return "\n".join([f"{json.dumps(safe_serialize(message.content))[:max_length]}..." for message in messages])

khoj/routers/api.py

@@ -613,7 +613,7 @@ async def post_automation(
         )
 
     if not subject:
-        subject = await acreate_title_from_query(q)
+        subject = await acreate_title_from_query(q, user)
 
     title = f"Automation: {subject}"
 

khoj/routers/api_chat.py

@@ -46,7 +46,7 @@ from khoj.routers.helpers import (
     FeedbackData,
     acreate_title_from_history,
     agenerate_chat_response,
-    aget_relevant_tools_to_execute,
+    aget_data_sources_and_output_format,
     construct_automation_created_message,
     create_automation,
     gather_raw_query_files,
@@ -667,27 +667,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 +709,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 +723,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
@@ -752,7 +762,7 @@ async def chat(
         attached_file_context = gather_raw_query_files(query_files)
 
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
-            conversation_commands = await aget_relevant_tools_to_execute(
+            chosen_io = await aget_data_sources_and_output_format(
                 q,
                 meta_log,
                 is_automated_task,
@@ -762,6 +772,7 @@ async def chat(
                 query_files=attached_file_context,
                 tracer=tracer,
             )
+            conversation_commands = chosen_io.get("sources") + [chosen_io.get("output")]
 
             # If we're doing research, we don't want to do anything else
             if ConversationCommand.Research in conversation_commands:
@@ -776,7 +787,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
 
@@ -833,7 +844,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(
@@ -852,7 +863,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)(
@@ -879,7 +890,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.
@@ -894,7 +905,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
 
@@ -915,7 +926,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
 
@@ -962,7 +973,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
 
@@ -1104,7 +1115,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
 
@@ -1131,7 +1142,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
 
@@ -1165,7 +1176,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)(
@@ -1212,7 +1223,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
 
@@ -1251,6 +1262,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/helpers.py

@@ -336,7 +336,7 @@ async def acheck_if_safe_prompt(system_prompt: str, user: KhojUser = None, lax:
     return is_safe, reason
 
 
-async def aget_relevant_tools_to_execute(
+async def aget_data_sources_and_output_format(
     query: str,
     conversation_history: dict,
     is_task: bool,
@@ -345,33 +345,33 @@ async def aget_relevant_tools_to_execute(
     agent: Agent = None,
     query_files: str = None,
     tracer: dict = {},
-):
+) -> Dict[str, Any]:
     """
-    Given a query, determine which of the available tools the agent should use in order to answer appropriately.
+    Given a query, determine which of the available data sources and output modes the agent should use to answer appropriately.
     """
 
-    tool_options = dict()
-    tool_options_str = ""
+    source_options = dict()
+    source_options_str = ""
 
-    agent_tools = agent.input_tools if agent else []
+    agent_sources = agent.input_tools if agent else []
 
-    for tool, description in tool_descriptions_for_llm.items():
-        tool_options[tool.value] = description
-        if len(agent_tools) == 0 or tool.value in agent_tools:
-            tool_options_str += f'- "{tool.value}": "{description}"\n'
+    for source, description in tool_descriptions_for_llm.items():
+        source_options[source.value] = description
+        if len(agent_sources) == 0 or source.value in agent_sources:
+            source_options_str += f'- "{source.value}": "{description}"\n'
 
-    mode_options = dict()
-    mode_options_str = ""
+    output_options = dict()
+    output_options_str = ""
 
-    output_modes = agent.output_modes if agent else []
+    agent_outputs = agent.output_modes if agent else []
 
-    for mode, description in mode_descriptions_for_llm.items():
+    for output, description in mode_descriptions_for_llm.items():
         # Do not allow tasks to schedule another task
-        if is_task and mode == ConversationCommand.Automation:
+        if is_task and output == ConversationCommand.Automation:
             continue
-        mode_options[mode.value] = description
-        if len(output_modes) == 0 or mode.value in output_modes:
-            mode_options_str += f'- "{mode.value}": "{description}"\n'
+        output_options[output.value] = description
+        if len(agent_outputs) == 0 or output.value in agent_outputs:
+            output_options_str += f'- "{output.value}": "{description}"\n'
 
     chat_history = construct_chat_history(conversation_history)
 
@@ -384,8 +384,8 @@ async def aget_relevant_tools_to_execute(
 
     relevant_tools_prompt = prompts.pick_relevant_tools.format(
         query=query,
-        tools=tool_options_str,
-        outputs=mode_options_str,
+        sources=source_options_str,
+        outputs=output_options_str,
         chat_history=chat_history,
         personality_context=personality_context,
     )
@@ -402,45 +402,43 @@ async def aget_relevant_tools_to_execute(
     try:
         response = clean_json(response)
         response = json.loads(response)
-        input_tools = [q.strip() for q in response["source"] if q.strip()]
-        if not isinstance(input_tools, list) or not input_tools or len(input_tools) == 0:
-            logger.error(f"Invalid response for determining relevant tools: {input_tools}")
-            return tool_options
-
-        output_modes = [q.strip() for q in response["output"] if q.strip()]
-        if not isinstance(output_modes, list) or not output_modes or len(output_modes) == 0:
-            logger.error(f"Invalid response for determining relevant output modes: {output_modes}")
-            return mode_options
-
-        final_response = [] if not is_task else [ConversationCommand.AutomatedTask]
-        for llm_suggested_tool in input_tools:
-            # Add a double check to verify it's in the agent list, because the LLM sometimes gets confused by the tool options.
-            if llm_suggested_tool in tool_options.keys() and (
-                len(agent_tools) == 0 or llm_suggested_tool in agent_tools
-            ):
-                # Check whether the tool exists as a valid ConversationCommand
-                final_response.append(ConversationCommand(llm_suggested_tool))
 
-        for llm_suggested_output in output_modes:
+        selected_sources = [q.strip() for q in response.get("source", []) if q.strip()]
+        selected_output = response.get("output", "text").strip()  # Default to text output
+
+        if not isinstance(selected_sources, list) or not selected_sources or len(selected_sources) == 0:
+            raise ValueError(
+                f"Invalid response for determining relevant tools: {selected_sources}. Raw Response: {response}"
+            )
+
+        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 llm_suggested_output in mode_options.keys() and (
-                len(output_modes) == 0 or llm_suggested_output in output_modes
+            if (
+                selected_source in source_options.keys()
+                and isinstance(result["sources"], list)
+                and (len(agent_sources) == 0 or selected_source in agent_sources)
             ):
                 # Check whether the tool exists as a valid ConversationCommand
-                final_response.append(ConversationCommand(llm_suggested_output))
+                result["sources"].append(ConversationCommand(selected_source))
+
+        # Add a double check to verify it's in the agent list, because the LLM sometimes gets confused by the tool options.
+        if selected_output in output_options.keys() and (len(agent_outputs) == 0 or selected_output in agent_outputs):
+            # Check whether the tool exists as a valid ConversationCommand
+            result["output"] = ConversationCommand(selected_output)
 
-        if is_none_or_empty(final_response):
-            if len(agent_tools) == 0:
-                final_response = [ConversationCommand.Default, ConversationCommand.Text]
+        if is_none_or_empty(result):
+            if len(agent_sources) == 0:
+                result = {"sources": [ConversationCommand.Default], "output": ConversationCommand.Text}
             else:
-                final_response = [ConversationCommand.General, ConversationCommand.Text]
-    except Exception:
-        logger.error(f"Invalid response for determining relevant tools: {response}")
-        if len(agent_tools) == 0:
-            final_response = [ConversationCommand.Default, ConversationCommand.Text]
-        else:
-            final_response = agent_tools
-    return final_response
+                result = {"sources": [ConversationCommand.General], "output": ConversationCommand.Text}
+    except Exception as e:
+        logger.error(f"Invalid response for determining relevant tools: {response}. Error: {e}", exc_info=True)
+        sources = agent_sources if len(agent_sources) > 0 else [ConversationCommand.Default]
+        output = agent_outputs[0] if len(agent_outputs) > 0 else ConversationCommand.Text
+        result = {"sources": sources, "output": output}
+
+    return result
 
 
 async def infer_webpage_urls(
@@ -755,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(
@@ -824,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 ""
     )
@@ -840,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}")
 
@@ -1772,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]:
@@ -1795,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):
@@ -1839,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/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/"
@@ -10,9 +11,10 @@ telemetry_server = "https://khoj.beta.haletic.com/v1/telemetry"
 content_directory = "~/.khoj/content/"
 default_offline_chat_models = [
     "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
+    "bartowski/Llama-3.2-3B-Instruct-GGUF",
     "bartowski/gemma-2-9b-it-GGUF",
     "bartowski/gemma-2-2b-it-GGUF",
-    "bartowski/Phi-3.5-mini-instruct-GGUF",
+    "Qwen/Qwen2.5-14B-Instruct-GGUF",
 ]
 default_openai_chat_models = ["gpt-4o-mini", "gpt-4o"]
 default_gemini_chat_models = ["gemini-1.5-flash", "gemini-1.5-pro"]
@@ -30,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},
+}

khoj/utils/helpers.py

@@ -540,3 +540,27 @@ def get_country_code_from_timezone(tz: str) -> str:
 def get_country_name_from_timezone(tz: str) -> str:
     """Get country name from timezone"""
     return country_names.get(get_country_code_from_timezone(tz), "United States")
+
+
+def get_cost_of_chat_message(model_name: str, input_tokens: int = 0, output_tokens: int = 0, prev_cost: float = 0.0):
+    """
+    Calculate cost of chat message based on input and output tokens
+    """
+
+    # Calculate cost of input and output tokens. Costs are per million tokens
+    input_cost = constants.model_to_cost.get(model_name, {}).get("input", 0) * (input_tokens / 1e6)
+    output_cost = constants.model_to_cost.get(model_name, {}).get("output", 0) * (output_tokens / 1e6)
+
+    return input_cost + output_cost + prev_cost
+
+
+def get_chat_usage_metrics(model_name: str, input_tokens: int = 0, output_tokens: int = 0, usage: dict = {}):
+    """
+    Get usage metrics for chat message based on input and output tokens
+    """
+    prev_usage = usage or {"input_tokens": 0, "output_tokens": 0, "cost": 0.0}
+    return {
+        "input_tokens": prev_usage["input_tokens"] + input_tokens,
+        "output_tokens": prev_usage["output_tokens"] + output_tokens,
+        "cost": get_cost_of_chat_message(model_name, input_tokens, output_tokens, prev_cost=prev_usage["cost"]),
+    }

{khoj-1.29.2.dev35.dist-info → khoj-1.30.2.dev22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: khoj
-Version: 1.29.2.dev35
+Version: 1.30.2.dev22
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev