khoj 1.16.1.dev15__py3-none-any.whl

khoj/processor/conversation/anthropic/utils.py

@@ -0,0 +1,114 @@
+import logging
+from threading import Thread
+from typing import Dict, List
+
+import anthropic
+from tenacity import (
+    before_sleep_log,
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    wait_random_exponential,
+)
+
+from khoj.processor.conversation.utils import ThreadedGenerator
+
+logger = logging.getLogger(__name__)
+
+anthropic_clients: Dict[str, anthropic.Anthropic] = {}
+
+
+DEFAULT_MAX_TOKENS_ANTHROPIC = 3000
+
+
+@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 anthropic_completion_with_backoff(
+    messages, system_prompt, model_name, temperature=0, api_key=None, model_kwargs=None, max_tokens=None
+) -> str:
+    if api_key not in anthropic_clients:
+        client: anthropic.Anthropic = anthropic.Anthropic(api_key=api_key)
+        anthropic_clients[api_key] = client
+    else:
+        client = anthropic_clients[api_key]
+
+    formatted_messages = [{"role": message.role, "content": message.content} for message in messages]
+
+    aggregated_response = ""
+    max_tokens = max_tokens or DEFAULT_MAX_TOKENS_ANTHROPIC
+
+    model_kwargs = model_kwargs or dict()
+    if system_prompt:
+        model_kwargs["system"] = system_prompt
+
+    with client.messages.stream(
+        messages=formatted_messages,
+        model=model_name,  # type: ignore
+        temperature=temperature,
+        timeout=20,
+        max_tokens=max_tokens,
+        **(model_kwargs),
+    ) as stream:
+        for text in stream.text_stream:
+            aggregated_response += text
+
+    return aggregated_response
+
+
+@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 anthropic_chat_completion_with_backoff(
+    messages,
+    compiled_references,
+    online_results,
+    model_name,
+    temperature,
+    api_key,
+    system_prompt,
+    max_prompt_size=None,
+    completion_func=None,
+    model_kwargs=None,
+):
+    g = ThreadedGenerator(compiled_references, online_results, completion_func=completion_func)
+    t = Thread(
+        target=anthropic_llm_thread,
+        args=(g, messages, system_prompt, model_name, temperature, api_key, max_prompt_size, model_kwargs),
+    )
+    t.start()
+    return g
+
+
+def anthropic_llm_thread(
+    g, messages, system_prompt, model_name, temperature, api_key, max_prompt_size=None, model_kwargs=None
+):
+    if api_key not in anthropic_clients:
+        client: anthropic.Anthropic = anthropic.Anthropic(api_key=api_key)
+        anthropic_clients[api_key] = client
+    else:
+        client: anthropic.Anthropic = anthropic_clients[api_key]
+
+    formatted_messages: List[anthropic.types.MessageParam] = [
+        anthropic.types.MessageParam(role=message.role, content=message.content) for message in messages
+    ]
+
+    with client.messages.stream(
+        messages=formatted_messages,
+        model=model_name,  # type: ignore
+        temperature=temperature,
+        system=system_prompt,
+        timeout=20,
+        max_tokens=DEFAULT_MAX_TOKENS_ANTHROPIC,
+        **(model_kwargs or dict()),
+    ) as stream:
+        for text in stream.text_stream:
+            g.send(text)
+
+    g.close()

khoj/processor/conversation/offline/chat_model.py

@@ -0,0 +1,231 @@
+import json
+import logging
+from datetime import datetime, timedelta
+from threading import Thread
+from typing import Any, Iterator, List, Union
+
+from langchain.schema import ChatMessage
+from llama_cpp import Llama
+
+from khoj.database.models import Agent
+from khoj.processor.conversation import prompts
+from khoj.processor.conversation.offline.utils import download_model
+from khoj.processor.conversation.utils import (
+    ThreadedGenerator,
+    generate_chatml_messages_with_context,
+)
+from khoj.utils import state
+from khoj.utils.constants import empty_escape_sequences
+from khoj.utils.helpers import ConversationCommand, is_none_or_empty
+from khoj.utils.rawconfig import LocationData
+
+logger = logging.getLogger(__name__)
+
+
+def extract_questions_offline(
+    text: str,
+    model: str = "NousResearch/Hermes-2-Pro-Mistral-7B-GGUF",
+    loaded_model: Union[Any, None] = None,
+    conversation_log={},
+    use_history: bool = True,
+    should_extract_questions: bool = True,
+    location_data: LocationData = None,
+    max_prompt_size: int = None,
+) -> 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.city}, {location_data.region}, {location_data.country}" if location_data else "Unknown"
+
+    # 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" and "text-to-image" not in chat["intent"].get("type"):
+                chat_history += f"Q: {chat['intent']['query']}\n"
+                chat_history += f"Khoj: {chat['message']}\n\n"
+
+    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"),
+        yesterday_date=yesterday,
+        last_year=last_year,
+        this_year=today.year,
+        location=location,
+    )
+    messages = generate_chatml_messages_with_context(
+        example_questions, model_name=model, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
+    )
+
+    state.chat_lock.acquire()
+    try:
+        response = send_message_to_model_offline(
+            messages, loaded_model=offline_chat_model, max_prompt_size=max_prompt_size
+        )
+    finally:
+        state.chat_lock.release()
+
+    # Extract, Clean Message from GPT's Response
+    try:
+        # This will expect to be a list with a single string with a list of questions
+        questions_str = (
+            str(response)
+            .strip(empty_escape_sequences)
+            .replace("['", '["')
+            .replace("<s>", "")
+            .replace("</s>", "")
+            .replace("']", '"]')
+            .replace("', '", '", "')
+        )
+        questions: List[str] = json.loads(questions_str)
+        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"Extracted Questions by Llama: {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=[],
+    conversation_log={},
+    model: str = "NousResearch/Hermes-2-Pro-Mistral-7B-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,
+) -> 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, max_tokens=max_prompt_size)
+    compiled_references_message = "\n\n".join({f"{item['compiled']}" for item in references})
+
+    current_date = datetime.now().strftime("%Y-%m-%d")
+
+    if agent and agent.personality:
+        system_prompt = prompts.custom_system_prompt_offline_chat.format(
+            name=agent.name, bio=agent.personality, current_date=current_date
+        )
+    else:
+        system_prompt = prompts.system_prompt_offline_chat.format(current_date=current_date)
+
+    conversation_primer = prompts.query_prompt.format(query=user_query)
+
+    if location_data:
+        location = f"{location_data.city}, {location_data.region}, {location_data.country}"
+        location_prompt = prompts.user_location.format(location=location)
+        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(compiled_references_message):
+        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()])
+
+    if ConversationCommand.Online 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"]
+
+        conversation_primer = f"{prompts.online_search_conversation.format(online_results=str(simplified_online_results))}\n{conversation_primer}"
+    if not is_none_or_empty(compiled_references_message):
+        conversation_primer = f"{prompts.notes_conversation_offline.format(references=compiled_references_message)}\n{conversation_primer}"
+
+    # Setup Prompt with Primer or Conversation History
+    messages = generate_chatml_messages_with_context(
+        conversation_primer,
+        system_prompt,
+        conversation_log,
+        model_name=model,
+        loaded_model=offline_chat_model,
+        max_prompt_size=max_prompt_size,
+        tokenizer_name=tokenizer_name,
+    )
+
+    g = ThreadedGenerator(references, online_results, completion_func=completion_func)
+    t = Thread(target=llm_thread, args=(g, messages, offline_chat_model, max_prompt_size))
+    t.start()
+    return g
+
+
+def llm_thread(g, messages: List[ChatMessage], model: Any, max_prompt_size: int = None):
+    stop_phrases = ["<s>", "INST]", "Notes:"]
+
+    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:
+            g.send(response["choices"][0]["delta"].get("content", ""))
+    finally:
+        state.chat_lock.release()
+    g.close()
+
+
+def send_message_to_model_offline(
+    messages: List[ChatMessage],
+    loaded_model=None,
+    model="NousResearch/Hermes-2-Pro-Mistral-7B-GGUF",
+    streaming=False,
+    stop=[],
+    max_prompt_size: int = None,
+):
+    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)
+    messages_dict = [{"role": message.role, "content": message.content} for message in messages]
+    response = offline_chat_model.create_chat_completion(messages_dict, stop=stop, stream=streaming)
+    if streaming:
+        return response
+    else:
+        return response["choices"][0]["message"].get("content", "")

khoj/processor/conversation/offline/utils.py

@@ -0,0 +1,78 @@
+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"
+
+    # 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() / 2e6)  # 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"]

khoj/processor/conversation/openai/gpt.py

@@ -0,0 +1,187 @@
+import json
+import logging
+from datetime import datetime, timedelta
+from typing import Dict, Optional
+
+from langchain.schema import ChatMessage
+
+from khoj.database.models import Agent
+from khoj.processor.conversation import prompts
+from khoj.processor.conversation.openai.utils import (
+    chat_completion_with_backoff,
+    completion_with_backoff,
+)
+from khoj.processor.conversation.utils import generate_chatml_messages_with_context
+from khoj.utils.helpers import ConversationCommand, is_none_or_empty
+from khoj.utils.rawconfig import LocationData
+
+logger = logging.getLogger(__name__)
+
+
+def extract_questions(
+    text,
+    model: Optional[str] = "gpt-4-turbo-preview",
+    conversation_log={},
+    api_key=None,
+    api_base_url=None,
+    temperature=0,
+    max_tokens=100,
+    location_data: LocationData = None,
+):
+    """
+    Infer search queries to retrieve relevant notes to answer user query
+    """
+    location = f"{location_data.city}, {location_data.region}, {location_data.country}" if location_data else "Unknown"
+
+    # Extract Past User Message and Inferred Questions from Conversation Log
+    chat_history = "".join(
+        [
+            f'Q: {chat["intent"]["query"]}\nKhoj: {{"queries": {chat["intent"].get("inferred-queries") or list([chat["intent"]["query"]])}}}\nA: {chat["message"]}\n\n'
+            for chat in conversation_log.get("chat", [])[-4:]
+            if chat["by"] == "khoj" and "text-to-image" not in chat["intent"].get("type")
+        ]
+    )
+
+    # Get dates relative to today for prompt creation
+    today = datetime.today()
+    current_new_year = today.replace(month=1, day=1)
+    last_new_year = current_new_year.replace(year=today.year - 1)
+
+    prompt = prompts.extract_questions.format(
+        current_date=today.strftime("%Y-%m-%d"),
+        day_of_week=today.strftime("%A"),
+        last_new_year=last_new_year.strftime("%Y"),
+        last_new_year_date=last_new_year.strftime("%Y-%m-%d"),
+        current_new_year_date=current_new_year.strftime("%Y-%m-%d"),
+        bob_tom_age_difference={current_new_year.year - 1984 - 30},
+        bob_age={current_new_year.year - 1984},
+        chat_history=chat_history,
+        text=text,
+        yesterday_date=(today - timedelta(days=1)).strftime("%Y-%m-%d"),
+        location=location,
+    )
+    messages = [ChatMessage(content=prompt, role="user")]
+
+    # Get Response from GPT
+    response = completion_with_backoff(
+        messages=messages,
+        model=model,
+        temperature=temperature,
+        api_base_url=api_base_url,
+        model_kwargs={"response_format": {"type": "json_object"}},
+        openai_api_key=api_key,
+    )
+
+    # Extract, Clean Message from GPT's Response
+    try:
+        response = response.strip()
+        response = json.loads(response)
+        response = [q.strip() for q in response["queries"] if q.strip()]
+        if not isinstance(response, list) or not response:
+            logger.error(f"Invalid response for constructing subqueries: {response}")
+            return [text]
+        return response
+    except:
+        logger.warning(f"GPT returned invalid JSON. Falling back to using user message as search query.\n{response}")
+        questions = [text]
+
+    logger.debug(f"Extracted Questions by GPT: {questions}")
+    return questions
+
+
+def send_message_to_model(messages, api_key, model, response_type="text", api_base_url=None):
+    """
+    Send message to model
+    """
+
+    # Get Response from GPT
+    return completion_with_backoff(
+        messages=messages,
+        model=model,
+        openai_api_key=api_key,
+        api_base_url=api_base_url,
+        model_kwargs={"response_format": {"type": response_type}},
+    )
+
+
+def converse(
+    references,
+    user_query,
+    online_results: Optional[Dict[str, Dict]] = None,
+    conversation_log={},
+    model: str = "gpt-3.5-turbo",
+    api_key: Optional[str] = None,
+    api_base_url: Optional[str] = None,
+    temperature: float = 0.2,
+    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,
+):
+    """
+    Converse with user using OpenAI's ChatGPT
+    """
+    # Initialize Variables
+    current_date = datetime.now().strftime("%Y-%m-%d")
+    compiled_references = "\n\n".join({f"# {item['compiled']}" for item in references})
+
+    conversation_primer = prompts.query_prompt.format(query=user_query)
+
+    if agent and agent.personality:
+        system_prompt = prompts.custom_personality.format(
+            name=agent.name, bio=agent.personality, current_date=current_date
+        )
+    else:
+        system_prompt = prompts.personality.format(current_date=current_date)
+
+    if location_data:
+        location = f"{location_data.city}, {location_data.region}, {location_data.country}"
+        location_prompt = prompts.user_location.format(location=location)
+        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(compiled_references):
+        completion_func(chat_response=prompts.no_notes_found.format())
+        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()])
+
+    if not is_none_or_empty(online_results):
+        conversation_primer = (
+            f"{prompts.online_search_conversation.format(online_results=str(online_results))}\n{conversation_primer}"
+        )
+    if not is_none_or_empty(compiled_references):
+        conversation_primer = f"{prompts.notes_conversation.format(query=user_query, references=compiled_references)}\n\n{conversation_primer}"
+
+    # Setup Prompt with Primer or Conversation History
+    messages = generate_chatml_messages_with_context(
+        conversation_primer,
+        system_prompt,
+        conversation_log,
+        model_name=model,
+        max_prompt_size=max_prompt_size,
+        tokenizer_name=tokenizer_name,
+    )
+    truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})
+    logger.debug(f"Conversation Context for GPT: {truncated_messages}")
+
+    # Get Response from GPT
+    return chat_completion_with_backoff(
+        messages=messages,
+        compiled_references=references,
+        online_results=online_results,
+        model_name=model,
+        temperature=temperature,
+        openai_api_key=api_key,
+        api_base_url=api_base_url,
+        completion_func=completion_func,
+        model_kwargs={"stop": ["Notes:\n["]},
+    )