khoj 1.33.3.dev32__py3-none-any.whl

khoj/processor/conversation/google/utils.py

@@ -0,0 +1,260 @@
+import logging
+import random
+from threading import Thread
+
+import google.generativeai as genai
+from google.generativeai.types.answer_types import FinishReason
+from google.generativeai.types.generation_types import StopCandidateException
+from google.generativeai.types.safety_types import (
+    HarmBlockThreshold,
+    HarmCategory,
+    HarmProbability,
+)
+from langchain.schema import ChatMessage
+from tenacity import (
+    before_sleep_log,
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    wait_random_exponential,
+)
+
+from khoj.processor.conversation.utils import (
+    ThreadedGenerator,
+    commit_conversation_trace,
+    get_image_from_url,
+)
+from khoj.utils import state
+from khoj.utils.helpers import (
+    get_chat_usage_metrics,
+    is_none_or_empty,
+    is_promptrace_enabled,
+)
+
+logger = logging.getLogger(__name__)
+
+
+MAX_OUTPUT_TOKENS_GEMINI = 8192
+
+
+@retry(
+    wait=wait_random_exponential(min=1, max=10),
+    stop=stop_after_attempt(2),
+    before_sleep=before_sleep_log(logger, logging.DEBUG),
+    reraise=True,
+)
+def gemini_completion_with_backoff(
+    messages, system_prompt, model_name, temperature=0, api_key=None, model_kwargs=None, tracer={}
+) -> str:
+    genai.configure(api_key=api_key)
+    model_kwargs = model_kwargs or dict()
+    model_kwargs["temperature"] = temperature
+    model_kwargs["max_output_tokens"] = MAX_OUTPUT_TOKENS_GEMINI
+    model = genai.GenerativeModel(
+        model_name,
+        generation_config=model_kwargs,
+        system_instruction=system_prompt,
+        safety_settings={
+            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+        },
+    )
+
+    formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
+
+    # Start chat session. All messages up to the last are considered to be part of the chat history
+    chat_session = model.start_chat(history=formatted_messages[0:-1])
+
+    try:
+        # Generate the response. The last message is considered to be the current prompt
+        response = chat_session.send_message(formatted_messages[-1]["parts"])
+        response_text = response.text
+    except StopCandidateException as e:
+        response = None
+        response_text, _ = handle_gemini_response(e.args)
+        # Respond with reason for stopping
+        logger.warning(
+            f"LLM Response Prevented for {model_name}: {response_text}.\n"
+            + f"Last Message by {messages[-1].role}: {messages[-1].content}"
+        )
+
+    # Aggregate cost of chat
+    input_tokens = response.usage_metadata.prompt_token_count if response else 0
+    output_tokens = response.usage_metadata.candidates_token_count if response else 0
+    tracer["usage"] = get_chat_usage_metrics(model_name, input_tokens, output_tokens, tracer.get("usage"))
+
+    # Save conversation trace
+    tracer["chat_model"] = model_name
+    tracer["temperature"] = temperature
+    if is_promptrace_enabled():
+        commit_conversation_trace(messages, response_text, tracer)
+
+    return response_text
+
+
+@retry(
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    stop=stop_after_attempt(2),
+    before_sleep=before_sleep_log(logger, logging.DEBUG),
+    reraise=True,
+)
+def gemini_chat_completion_with_backoff(
+    messages,
+    compiled_references,
+    online_results,
+    model_name,
+    temperature,
+    api_key,
+    system_prompt,
+    completion_func=None,
+    model_kwargs=None,
+    tracer: dict = {},
+):
+    g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
+    t = Thread(
+        target=gemini_llm_thread,
+        args=(g, messages, system_prompt, model_name, temperature, api_key, model_kwargs, tracer),
+    )
+    t.start()
+    return g
+
+
+def gemini_llm_thread(
+    g, messages, system_prompt, model_name, temperature, api_key, model_kwargs=None, tracer: dict = {}
+):
+    try:
+        genai.configure(api_key=api_key)
+        model_kwargs = model_kwargs or dict()
+        model_kwargs["temperature"] = temperature
+        model_kwargs["max_output_tokens"] = MAX_OUTPUT_TOKENS_GEMINI
+        model_kwargs["stop_sequences"] = ["Notes:\n["]
+        model = genai.GenerativeModel(
+            model_name,
+            generation_config=model_kwargs,
+            system_instruction=system_prompt,
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+            },
+        )
+
+        aggregated_response = ""
+        formatted_messages = [{"role": message.role, "parts": message.content} for message in messages]
+
+        # all messages up to the last are considered to be part of the chat history
+        chat_session = model.start_chat(history=formatted_messages[0:-1])
+        # the last message is considered to be the current prompt
+        for chunk in chat_session.send_message(formatted_messages[-1]["parts"], stream=True):
+            message, stopped = handle_gemini_response(chunk.candidates, chunk.prompt_feedback)
+            message = message or chunk.text
+            aggregated_response += message
+            g.send(message)
+            if stopped:
+                raise StopCandidateException(message)
+
+        # Calculate cost of chat
+        input_tokens = chunk.usage_metadata.prompt_token_count
+        output_tokens = chunk.usage_metadata.candidates_token_count
+        tracer["usage"] = get_chat_usage_metrics(model_name, input_tokens, output_tokens, tracer.get("usage"))
+
+        # Save conversation trace
+        tracer["chat_model"] = model_name
+        tracer["temperature"] = temperature
+        if is_promptrace_enabled():
+            commit_conversation_trace(messages, aggregated_response, tracer)
+    except StopCandidateException as e:
+        logger.warning(
+            f"LLM Response Prevented for {model_name}: {e.args[0]}.\n"
+            + f"Last Message by {messages[-1].role}: {messages[-1].content}"
+        )
+    except Exception as e:
+        logger.error(f"Error in gemini_llm_thread: {e}", exc_info=True)
+    finally:
+        g.close()
+
+
+def handle_gemini_response(candidates, prompt_feedback=None):
+    """Check if Gemini response was blocked and return an explanatory error message."""
+    # Check if the response was blocked due to safety concerns with the prompt
+    if len(candidates) == 0 and prompt_feedback:
+        message = f"\nI'd prefer to not respond to that due to **{prompt_feedback.block_reason.name}** issues with your query."
+        stopped = True
+    # Check if the response was blocked due to safety concerns with the generated content
+    elif candidates[0].finish_reason == FinishReason.SAFETY:
+        message = generate_safety_response(candidates[0].safety_ratings)
+        stopped = True
+    # Check if finish reason is empty, therefore generation is in progress
+    elif not candidates[0].finish_reason:
+        message = None
+        stopped = False
+    # Check if the response was stopped due to reaching maximum token limit or other reasons
+    elif candidates[0].finish_reason != FinishReason.STOP:
+        message = f"\nI can't talk further about that because of **{candidates[0].finish_reason.name} issue.**"
+        stopped = True
+    # Otherwise, the response is valid and can be used
+    else:
+        message = None
+        stopped = False
+    return message, stopped
+
+
+def generate_safety_response(safety_ratings):
+    """Generate a conversational response based on the safety ratings of the response."""
+    # Get the safety rating with the highest probability
+    max_safety_rating = sorted(safety_ratings, key=lambda x: x.probability, reverse=True)[0]
+    # Remove the "HARM_CATEGORY_" prefix and title case the category name
+    max_safety_category = " ".join(max_safety_rating.category.name.split("_")[2:]).title()
+    # Add a bit of variety to the discomfort level based on the safety rating probability
+    discomfort_level = {
+        HarmProbability.HARM_PROBABILITY_UNSPECIFIED: " ",
+        HarmProbability.LOW: "a bit ",
+        HarmProbability.MEDIUM: "moderately ",
+        HarmProbability.HIGH: random.choice(["very ", "quite ", "fairly "]),
+    }[max_safety_rating.probability]
+    # Generate a response using a random response template
+    safety_response_choice = random.choice(
+        [
+            "\nUmm, I'd rather not to respond to that. The conversation has some probability of going into **{category}** territory.",
+            "\nI'd prefer not to talk about **{category}** related topics. It makes me {discomfort_level}uncomfortable.",
+            "\nI feel {discomfort_level}squeamish talking about **{category}** related stuff! Can we talk about something less controversial?",
+            "\nThat sounds {discomfort_level}outside the [Overtone Window](https://en.wikipedia.org/wiki/Overton_window) of acceptable conversation. Should we stick to something less {category} related?",
+        ]
+    )
+    return safety_response_choice.format(
+        category=max_safety_category, probability=max_safety_rating.probability.name, discomfort_level=discomfort_level
+    )
+
+
+def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str = None) -> tuple[list[str], str]:
+    # Extract system message
+    system_prompt = system_prompt or ""
+    for message in messages.copy():
+        if message.role == "system":
+            system_prompt += message.content
+            messages.remove(message)
+    system_prompt = None if is_none_or_empty(system_prompt) else system_prompt
+
+    for message in messages:
+        # Convert message content to string list from chatml dictionary list
+        if isinstance(message.content, list):
+            # Convert image_urls to PIL.Image and place them at beginning of list (better for Gemini)
+            message.content = [
+                get_image_from_url(item["image_url"]["url"]).content
+                if item["type"] == "image_url"
+                else item.get("text", "")
+                for item in sorted(message.content, key=lambda x: 0 if x["type"] == "image_url" else 1)
+            ]
+        elif isinstance(message.content, str):
+            message.content = [message.content]
+
+        if message.role == "assistant":
+            message.role = "model"
+
+    if len(messages) == 1:
+        messages[0].role = "user"
+
+    return messages, system_prompt

khoj/processor/conversation/offline/chat_model.py

@@ -0,0 +1,308 @@
+import json
+import logging
+import os
+from datetime import datetime, timedelta
+from threading import Thread
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+import pyjson5
+from langchain.schema import ChatMessage
+from llama_cpp import Llama
+
+from khoj.database.models import Agent, ChatModel, KhojUser
+from khoj.processor.conversation import prompts
+from khoj.processor.conversation.offline.utils import download_model
+from khoj.processor.conversation.utils import (
+    ThreadedGenerator,
+    clean_json,
+    commit_conversation_trace,
+    generate_chatml_messages_with_context,
+    messages_to_print,
+)
+from khoj.utils import state
+from khoj.utils.constants import empty_escape_sequences
+from khoj.utils.helpers import (
+    ConversationCommand,
+    is_none_or_empty,
+    is_promptrace_enabled,
+    truncate_code_context,
+)
+from khoj.utils.rawconfig import FileAttachment, LocationData
+from khoj.utils.yaml import yaml_dump
+
+logger = logging.getLogger(__name__)
+
+
+def extract_questions_offline(
+    text: str,
+    model: str = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
+    loaded_model: Union[Any, None] = None,
+    conversation_log={},
+    use_history: bool = True,
+    should_extract_questions: bool = True,
+    location_data: LocationData = None,
+    user: KhojUser = None,
+    max_prompt_size: int = None,
+    temperature: float = 0.7,
+    personality_context: Optional[str] = None,
+    query_files: str = None,
+    tracer: dict = {},
+) -> List[str]:
+    """
+    Infer search queries to retrieve relevant notes to answer user query
+    """
+    all_questions = text.split("? ")
+    all_questions = [q + "?" for q in all_questions[:-1]] + [all_questions[-1]]
+
+    if not should_extract_questions:
+        return all_questions
+
+    assert loaded_model is None or isinstance(loaded_model, Llama), "loaded_model must be of type Llama, if configured"
+    offline_chat_model = loaded_model or download_model(model, max_tokens=max_prompt_size)
+
+    location = f"{location_data}" if location_data else "Unknown"
+    username = prompts.user_name.format(name=user.get_full_name()) if user and user.get_full_name() else ""
+
+    # Extract Past User Message and Inferred Questions from Conversation Log
+    chat_history = ""
+
+    if use_history:
+        for chat in conversation_log.get("chat", [])[-4:]:
+            if chat["by"] == "khoj":
+                chat_history += f"Q: {chat['intent']['query']}\n"
+                chat_history += f"Khoj: {chat['message']}\n\n"
+
+    # Get dates relative to today for prompt creation
+    today = datetime.today()
+    yesterday = (today - timedelta(days=1)).strftime("%Y-%m-%d")
+    last_year = today.year - 1
+    example_questions = prompts.extract_questions_offline.format(
+        query=text,
+        chat_history=chat_history,
+        current_date=today.strftime("%Y-%m-%d"),
+        day_of_week=today.strftime("%A"),
+        current_month=today.strftime("%Y-%m"),
+        yesterday_date=yesterday,
+        last_year=last_year,
+        this_year=today.year,
+        location=location,
+        username=username,
+        personality_context=personality_context,
+    )
+
+    messages = generate_chatml_messages_with_context(
+        example_questions,
+        model_name=model,
+        loaded_model=offline_chat_model,
+        max_prompt_size=max_prompt_size,
+        model_type=ChatModel.ModelType.OFFLINE,
+        query_files=query_files,
+    )
+
+    state.chat_lock.acquire()
+    try:
+        response = send_message_to_model_offline(
+            messages,
+            loaded_model=offline_chat_model,
+            model_name=model,
+            max_prompt_size=max_prompt_size,
+            temperature=temperature,
+            response_type="json_object",
+            tracer=tracer,
+        )
+    finally:
+        state.chat_lock.release()
+
+    # Extract and clean the chat model's response
+    try:
+        response = clean_json(empty_escape_sequences)
+        response = pyjson5.loads(response)
+        questions = [q.strip() for q in response["queries"] if q.strip()]
+        questions = filter_questions(questions)
+    except:
+        logger.warning(f"Llama returned invalid JSON. Falling back to using user message as search query.\n{response}")
+        return all_questions
+    logger.debug(f"Questions extracted by {model}: {questions}")
+    return questions
+
+
+def filter_questions(questions: List[str]):
+    # Skip questions that seem to be apologizing for not being able to answer the question
+    hint_words = [
+        "sorry",
+        "apologize",
+        "unable",
+        "can't",
+        "cannot",
+        "don't know",
+        "don't understand",
+        "do not know",
+        "do not understand",
+    ]
+    filtered_questions = set()
+    for q in questions:
+        if not any([word in q.lower() for word in hint_words]) and not is_none_or_empty(q):
+            filtered_questions.add(q)
+
+    return list(filtered_questions)
+
+
+def converse_offline(
+    user_query,
+    references=[],
+    online_results={},
+    code_results={},
+    conversation_log={},
+    model_name: str = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
+    loaded_model: Union[Any, None] = None,
+    completion_func=None,
+    conversation_commands=[ConversationCommand.Default],
+    max_prompt_size=None,
+    tokenizer_name=None,
+    location_data: LocationData = None,
+    user_name: str = None,
+    agent: Agent = None,
+    query_files: str = None,
+    generated_files: List[FileAttachment] = None,
+    additional_context: List[str] = None,
+    generated_asset_results: Dict[str, Dict] = {},
+    tracer: dict = {},
+) -> Union[ThreadedGenerator, Iterator[str]]:
+    """
+    Converse with user using Llama
+    """
+    # Initialize Variables
+    assert loaded_model is None or isinstance(loaded_model, Llama), "loaded_model must be of type Llama, if configured"
+    offline_chat_model = loaded_model or download_model(model_name, max_tokens=max_prompt_size)
+    tracer["chat_model"] = model_name
+    current_date = datetime.now()
+
+    if agent and agent.personality:
+        system_prompt = prompts.custom_system_prompt_offline_chat.format(
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
+    else:
+        system_prompt = prompts.system_prompt_offline_chat.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
+
+    if location_data:
+        location_prompt = prompts.user_location.format(location=f"{location_data}")
+        system_prompt = f"{system_prompt}\n{location_prompt}"
+
+    if user_name:
+        user_name_prompt = prompts.user_name.format(name=user_name)
+        system_prompt = f"{system_prompt}\n{user_name_prompt}"
+
+    # Get Conversation Primer appropriate to Conversation Type
+    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(references):
+        return iter([prompts.no_notes_found.format()])
+    elif conversation_commands == [ConversationCommand.Online] and is_none_or_empty(online_results):
+        completion_func(chat_response=prompts.no_online_results_found.format())
+        return iter([prompts.no_online_results_found.format()])
+
+    context_message = ""
+    if not is_none_or_empty(references):
+        context_message = f"{prompts.notes_conversation_offline.format(references=yaml_dump(references))}\n\n"
+    if ConversationCommand.Online in conversation_commands or ConversationCommand.Webpage in conversation_commands:
+        simplified_online_results = online_results.copy()
+        for result in online_results:
+            if online_results[result].get("webpages"):
+                simplified_online_results[result] = online_results[result]["webpages"]
+
+        context_message += f"{prompts.online_search_conversation_offline.format(online_results=yaml_dump(simplified_online_results))}\n\n"
+    if ConversationCommand.Code in conversation_commands and not is_none_or_empty(code_results):
+        context_message += (
+            f"{prompts.code_executed_context.format(code_results=truncate_code_context(code_results))}\n\n"
+        )
+    context_message = context_message.strip()
+
+    # Setup Prompt with Primer or Conversation History
+    messages = generate_chatml_messages_with_context(
+        user_query,
+        system_prompt,
+        conversation_log,
+        context_message=context_message,
+        model_name=model_name,
+        loaded_model=offline_chat_model,
+        max_prompt_size=max_prompt_size,
+        tokenizer_name=tokenizer_name,
+        model_type=ChatModel.ModelType.OFFLINE,
+        query_files=query_files,
+        generated_files=generated_files,
+        generated_asset_results=generated_asset_results,
+        program_execution_context=additional_context,
+    )
+
+    logger.debug(f"Conversation Context for {model_name}: {messages_to_print(messages)}")
+
+    g = ThreadedGenerator(references, online_results, completion_func=completion_func)
+    t = Thread(target=llm_thread, args=(g, messages, offline_chat_model, max_prompt_size, tracer))
+    t.start()
+    return g
+
+
+def llm_thread(g, messages: List[ChatMessage], model: Any, max_prompt_size: int = None, tracer: dict = {}):
+    stop_phrases = ["<s>", "INST]", "Notes:"]
+    aggregated_response = ""
+
+    state.chat_lock.acquire()
+    try:
+        response_iterator = send_message_to_model_offline(
+            messages, loaded_model=model, stop=stop_phrases, max_prompt_size=max_prompt_size, streaming=True
+        )
+        for response in response_iterator:
+            response_delta = response["choices"][0]["delta"].get("content", "")
+            aggregated_response += response_delta
+            g.send(response_delta)
+
+        # Save conversation trace
+        if is_promptrace_enabled():
+            commit_conversation_trace(messages, aggregated_response, tracer)
+
+    finally:
+        state.chat_lock.release()
+        g.close()
+
+
+def send_message_to_model_offline(
+    messages: List[ChatMessage],
+    loaded_model=None,
+    model_name="bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
+    temperature: float = 0.2,
+    streaming=False,
+    stop=[],
+    max_prompt_size: int = None,
+    response_type: str = "text",
+    tracer: dict = {},
+):
+    assert loaded_model is None or isinstance(loaded_model, Llama), "loaded_model must be of type Llama, if configured"
+    offline_chat_model = loaded_model or download_model(model_name, max_tokens=max_prompt_size)
+    messages_dict = [{"role": message.role, "content": message.content} for message in messages]
+    seed = int(os.getenv("KHOJ_LLM_SEED")) if os.getenv("KHOJ_LLM_SEED") else None
+    response = offline_chat_model.create_chat_completion(
+        messages_dict,
+        stop=stop,
+        stream=streaming,
+        temperature=temperature,
+        response_format={"type": response_type},
+        seed=seed,
+    )
+
+    if streaming:
+        return response
+
+    response_text = response["choices"][0]["message"].get("content", "")
+
+    # Save conversation trace for non-streaming responses
+    # Streamed responses need to be saved by the calling function
+    tracer["chat_model"] = model_name
+    tracer["temperature"] = temperature
+    if is_promptrace_enabled():
+        commit_conversation_trace(messages, response_text, tracer)
+
+    return response_text

khoj/processor/conversation/offline/utils.py

@@ -0,0 +1,80 @@
+import glob
+import logging
+import math
+import os
+from typing import Any, Dict
+
+from huggingface_hub.constants import HF_HUB_CACHE
+
+from khoj.utils import state
+from khoj.utils.helpers import get_device_memory
+
+logger = logging.getLogger(__name__)
+
+
+def download_model(repo_id: str, filename: str = "*Q4_K_M.gguf", max_tokens: int = None):
+    # Initialize Model Parameters
+    # Use n_ctx=0 to get context size from the model
+    kwargs: Dict[str, Any] = {"n_threads": 4, "n_ctx": 0, "verbose": False}
+
+    # Decide whether to load model to GPU or CPU
+    device = "gpu" if state.chat_on_gpu and state.device != "cpu" else "cpu"
+    kwargs["n_gpu_layers"] = -1 if device == "gpu" else 0
+
+    # Add chat format if known
+    if "llama-3" in repo_id.lower():
+        kwargs["chat_format"] = "llama-3"
+    elif "gemma-2" in repo_id.lower():
+        kwargs["chat_format"] = "gemma"
+
+    # Check if the model is already downloaded
+    model_path = load_model_from_cache(repo_id, filename)
+    chat_model = None
+    try:
+        chat_model = load_model(model_path, repo_id, filename, kwargs)
+    except:
+        # Load model on CPU if GPU is not available
+        kwargs["n_gpu_layers"], device = 0, "cpu"
+        chat_model = load_model(model_path, repo_id, filename, kwargs)
+
+    # Now load the model with context size set based on:
+    # 1. context size supported by model and
+    # 2. configured size or machine (V)RAM
+    kwargs["n_ctx"] = infer_max_tokens(chat_model.n_ctx(), max_tokens)
+    chat_model = load_model(model_path, repo_id, filename, kwargs)
+
+    logger.debug(
+        f"{'Loaded' if model_path else 'Downloaded'} chat model to {device.upper()} with {kwargs['n_ctx']} token context window."
+    )
+    return chat_model
+
+
+def load_model(model_path: str, repo_id: str, filename: str = "*Q4_K_M.gguf", kwargs: dict = {}):
+    from llama_cpp.llama import Llama
+
+    if model_path:
+        return Llama(model_path, **kwargs)
+    else:
+        return Llama.from_pretrained(repo_id=repo_id, filename=filename, **kwargs)
+
+
+def load_model_from_cache(repo_id: str, filename: str, repo_type="models"):
+    # Construct the path to the model file in the cache directory
+    repo_org, repo_name = repo_id.split("/")
+    object_id = "--".join([repo_type, repo_org, repo_name])
+    model_path = os.path.sep.join([HF_HUB_CACHE, object_id, "snapshots", "**", filename])
+
+    # Check if the model file exists
+    paths = glob.glob(model_path)
+    if paths:
+        return paths[0]
+    else:
+        return None
+
+
+def infer_max_tokens(model_context_window: int, configured_max_tokens=None) -> int:
+    """Infer max prompt size based on device memory and max context window supported by the model"""
+    configured_max_tokens = math.inf if configured_max_tokens is None else configured_max_tokens
+    vram_based_n_ctx = int(get_device_memory() / 1e6)  # based on heuristic
+    configured_max_tokens = configured_max_tokens or math.inf  # do not use if set to None
+    return min(configured_max_tokens, vram_based_n_ctx, model_context_window)

khoj/processor/conversation/offline/whisper.py

@@ -0,0 +1,15 @@
+import whisper
+from asgiref.sync import sync_to_async
+
+from khoj.utils import state
+
+
+async def transcribe_audio_offline(audio_filename: str, model: str) -> str:
+    """
+    Transcribe audio file offline using Whisper
+    """
+    # Send the audio data to the Whisper API
+    if not state.whisper_model:
+        state.whisper_model = whisper.load_model(model)
+    response = await sync_to_async(state.whisper_model.transcribe)(audio_filename)
+    return response["text"]