khoj 2.0.0b12.dev5__py3-none-any.whl → 2.0.0b13.dev5__py3-none-any.whl

khoj/processor/conversation/offline/chat_model.py

@@ -1,224 +0,0 @@
-import asyncio
-import logging
-import os
-from datetime import datetime
-from threading import Thread
-from time import perf_counter
-from typing import Any, AsyncGenerator, Dict, List, Union
-
-from langchain_core.messages.chat import ChatMessage
-from llama_cpp import Llama
-
-from khoj.database.models import Agent, ChatMessageModel, ChatModel
-from khoj.processor.conversation import prompts
-from khoj.processor.conversation.offline.utils import download_model
-from khoj.processor.conversation.utils import (
-    ResponseWithThought,
-    commit_conversation_trace,
-    generate_chatml_messages_with_context,
-    messages_to_print,
-)
-from khoj.utils import state
-from khoj.utils.helpers import (
-    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__)
-
-
-async def converse_offline(
-    # Query
-    user_query: str,
-    # Context
-    references: list[dict] = [],
-    online_results={},
-    code_results={},
-    query_files: str = None,
-    generated_files: List[FileAttachment] = None,
-    additional_context: List[str] = None,
-    generated_asset_results: Dict[str, Dict] = {},
-    location_data: LocationData = None,
-    user_name: str = None,
-    chat_history: list[ChatMessageModel] = [],
-    # Model
-    model_name: str = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
-    loaded_model: Union[Any, None] = None,
-    max_prompt_size=None,
-    tokenizer_name=None,
-    agent: Agent = None,
-    tracer: dict = {},
-) -> AsyncGenerator[ResponseWithThought, None]:
-    """
-    Converse with user using Llama (Async Version)
-    """
-    # 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
-    context_message = ""
-    if not is_none_or_empty(references):
-        context_message = f"{prompts.notes_conversation_offline.format(references=yaml_dump(references))}\n\n"
-    if not is_none_or_empty(online_results):
-        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 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,
-        chat_history,
-        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)}")
-
-    # Use asyncio.Queue and a thread to bridge sync iterator
-    queue: asyncio.Queue[ResponseWithThought] = asyncio.Queue()
-    stop_phrases = ["<s>", "INST]", "Notes:"]
-
-    def _sync_llm_thread():
-        """Synchronous function to run in a separate thread."""
-        aggregated_response = ""
-        start_time = perf_counter()
-        state.chat_lock.acquire()
-        try:
-            response_iterator = send_message_to_model_offline(
-                messages,
-                loaded_model=offline_chat_model,
-                stop=stop_phrases,
-                max_prompt_size=max_prompt_size,
-                streaming=True,
-                tracer=tracer,
-            )
-            for response in response_iterator:
-                response_delta: str = response["choices"][0]["delta"].get("content", "")
-                # Log the time taken to start response
-                if aggregated_response == "" and response_delta != "":
-                    logger.info(f"First response took: {perf_counter() - start_time:.3f} seconds")
-                # Handle response chunk
-                aggregated_response += response_delta
-                # Put chunk into the asyncio queue (non-blocking)
-                try:
-                    queue.put_nowait(ResponseWithThought(text=response_delta))
-                except asyncio.QueueFull:
-                    # Should not happen with default queue size unless consumer is very slow
-                    logger.warning("Asyncio queue full during offline LLM streaming.")
-                    # Potentially block here or handle differently if needed
-                    asyncio.run(queue.put(ResponseWithThought(text=response_delta)))
-
-            # Log the time taken to stream the entire response
-            logger.info(f"Chat streaming took: {perf_counter() - start_time:.3f} seconds")
-
-            # Save conversation trace
-            tracer["chat_model"] = model_name
-            if is_promptrace_enabled():
-                commit_conversation_trace(messages, aggregated_response, tracer)
-
-        except Exception as e:
-            logger.error(f"Error in offline LLM thread: {e}", exc_info=True)
-        finally:
-            state.chat_lock.release()
-            # Signal end of stream
-            queue.put_nowait(None)
-
-    # Start the synchronous thread
-    thread = Thread(target=_sync_llm_thread)
-    thread.start()
-
-    # Asynchronously consume from the queue
-    while True:
-        chunk = await queue.get()
-        if chunk is None:  # End of stream signal
-            queue.task_done()
-            break
-        yield chunk
-        queue.task_done()
-
-    # Wait for the thread to finish (optional, ensures cleanup)
-    loop = asyncio.get_running_loop()
-    await loop.run_in_executor(None, thread.join)
-
-
-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: str = 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 ResponseWithThought(text=response_text)

khoj/processor/conversation/offline/utils.py

@@ -1,80 +0,0 @@
-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

@@ -1,15 +0,0 @@
-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/utils/fs_syncer.py

@@ -1,252 +0,0 @@
-import glob
-import logging
-import os
-from pathlib import Path
-from typing import Optional
-
-from bs4 import BeautifulSoup
-from magika import Magika
-
-from khoj.database.models import (
-    KhojUser,
-    LocalMarkdownConfig,
-    LocalOrgConfig,
-    LocalPdfConfig,
-    LocalPlaintextConfig,
-)
-from khoj.utils.config import SearchType
-from khoj.utils.helpers import get_absolute_path, is_none_or_empty
-from khoj.utils.rawconfig import TextContentConfig
-
-logger = logging.getLogger(__name__)
-magika = Magika()
-
-
-def collect_files(user: KhojUser, search_type: Optional[SearchType] = SearchType.All) -> dict:
-    files: dict[str, dict] = {"docx": {}, "image": {}}
-
-    if search_type == SearchType.All or search_type == SearchType.Org:
-        org_config = LocalOrgConfig.objects.filter(user=user).first()
-        files["org"] = get_org_files(construct_config_from_db(org_config)) if org_config else {}
-    if search_type == SearchType.All or search_type == SearchType.Markdown:
-        markdown_config = LocalMarkdownConfig.objects.filter(user=user).first()
-        files["markdown"] = get_markdown_files(construct_config_from_db(markdown_config)) if markdown_config else {}
-    if search_type == SearchType.All or search_type == SearchType.Plaintext:
-        plaintext_config = LocalPlaintextConfig.objects.filter(user=user).first()
-        files["plaintext"] = get_plaintext_files(construct_config_from_db(plaintext_config)) if plaintext_config else {}
-    if search_type == SearchType.All or search_type == SearchType.Pdf:
-        pdf_config = LocalPdfConfig.objects.filter(user=user).first()
-        files["pdf"] = get_pdf_files(construct_config_from_db(pdf_config)) if pdf_config else {}
-    files["image"] = {}
-    files["docx"] = {}
-    return files
-
-
-def construct_config_from_db(db_config) -> TextContentConfig:
-    return TextContentConfig(
-        input_files=db_config.input_files,
-        input_filter=db_config.input_filter,
-        index_heading_entries=db_config.index_heading_entries,
-    )
-
-
-def get_plaintext_files(config: TextContentConfig) -> dict[str, str]:
-    def is_plaintextfile(file: str):
-        "Check if file is plaintext file"
-        # Check if file path exists
-        content_group = magika.identify_path(Path(file)).output.group
-        # Use file extension to decide plaintext if file content is not identifiable
-        valid_text_file_extensions = ("txt", "md", "markdown", "org" "mbox", "rst", "html", "htm", "xml")
-        return file.endswith(valid_text_file_extensions) or content_group in ["text", "code"]
-
-    def extract_html_content(html_content: str):
-        "Extract content from HTML"
-        soup = BeautifulSoup(html_content, "html.parser")
-        return soup.get_text(strip=True, separator="\n")
-
-    # Extract required fields from config
-    input_files, input_filters = (
-        config.input_files,
-        config.input_filter,
-    )
-
-    # Input Validation
-    if is_none_or_empty(input_files) and is_none_or_empty(input_filters):
-        logger.debug("At least one of input-files or input-file-filter is required to be specified")
-        return {}
-
-    # Get all plain text files to process
-    absolute_plaintext_files, filtered_plaintext_files = set(), set()
-    if input_files:
-        absolute_plaintext_files = {get_absolute_path(jsonl_file) for jsonl_file in input_files}
-    if input_filters:
-        filtered_plaintext_files = {
-            filtered_file
-            for plaintext_file_filter in input_filters
-            for filtered_file in glob.glob(get_absolute_path(plaintext_file_filter), recursive=True)
-            if os.path.isfile(filtered_file)
-        }
-
-    all_target_files = sorted(absolute_plaintext_files | filtered_plaintext_files)
-
-    files_with_no_plaintext_extensions = {
-        target_files for target_files in all_target_files if not is_plaintextfile(target_files)
-    }
-    if any(files_with_no_plaintext_extensions):
-        logger.warning(f"Skipping unsupported files from plaintext indexing: {files_with_no_plaintext_extensions}")
-        all_target_files = list(set(all_target_files) - files_with_no_plaintext_extensions)
-
-    logger.debug(f"Processing files: {all_target_files}")
-
-    filename_to_content_map = {}
-    for file in all_target_files:
-        with open(file, "r", encoding="utf8") as f:
-            try:
-                plaintext_content = f.read()
-                if file.endswith(("html", "htm", "xml")):
-                    plaintext_content = extract_html_content(plaintext_content)
-                filename_to_content_map[file] = plaintext_content
-            except Exception as e:
-                logger.warning(f"Unable to read file: {file} as plaintext. Skipping file.")
-                logger.warning(e, exc_info=True)
-
-    return filename_to_content_map
-
-
-def get_org_files(config: TextContentConfig):
-    # Extract required fields from config
-    org_files, org_file_filters = (
-        config.input_files,
-        config.input_filter,
-    )
-
-    # Input Validation
-    if is_none_or_empty(org_files) and is_none_or_empty(org_file_filters):
-        logger.debug("At least one of org-files or org-file-filter is required to be specified")
-        return {}
-
-    # Get Org files to process
-    absolute_org_files, filtered_org_files = set(), set()
-    if org_files:
-        absolute_org_files = {get_absolute_path(org_file) for org_file in org_files}
-    if org_file_filters:
-        filtered_org_files = {
-            filtered_file
-            for org_file_filter in org_file_filters
-            for filtered_file in glob.glob(get_absolute_path(org_file_filter), recursive=True)
-            if os.path.isfile(filtered_file)
-        }
-
-    all_org_files = sorted(absolute_org_files | filtered_org_files)
-
-    files_with_non_org_extensions = {org_file for org_file in all_org_files if not org_file.endswith(".org")}
-    if any(files_with_non_org_extensions):
-        logger.warning(f"There maybe non org-mode files in the input set: {files_with_non_org_extensions}")
-
-    logger.debug(f"Processing files: {all_org_files}")
-
-    filename_to_content_map = {}
-    for file in all_org_files:
-        with open(file, "r", encoding="utf8") as f:
-            try:
-                filename_to_content_map[file] = f.read()
-            except Exception as e:
-                logger.warning(f"Unable to read file: {file} as org. Skipping file.")
-                logger.warning(e, exc_info=True)
-
-    return filename_to_content_map
-
-
-def get_markdown_files(config: TextContentConfig):
-    # Extract required fields from config
-    markdown_files, markdown_file_filters = (
-        config.input_files,
-        config.input_filter,
-    )
-
-    # Input Validation
-    if is_none_or_empty(markdown_files) and is_none_or_empty(markdown_file_filters):
-        logger.debug("At least one of markdown-files or markdown-file-filter is required to be specified")
-        return {}
-
-    # Get markdown files to process
-    absolute_markdown_files, filtered_markdown_files = set(), set()
-    if markdown_files:
-        absolute_markdown_files = {get_absolute_path(markdown_file) for markdown_file in markdown_files}
-
-    if markdown_file_filters:
-        filtered_markdown_files = {
-            filtered_file
-            for markdown_file_filter in markdown_file_filters
-            for filtered_file in glob.glob(get_absolute_path(markdown_file_filter), recursive=True)
-            if os.path.isfile(filtered_file)
-        }
-
-    all_markdown_files = sorted(absolute_markdown_files | filtered_markdown_files)
-
-    files_with_non_markdown_extensions = {
-        md_file for md_file in all_markdown_files if not md_file.endswith(".md") and not md_file.endswith(".markdown")
-    }
-
-    if any(files_with_non_markdown_extensions):
-        logger.warning(
-            f"[Warning] There maybe non markdown-mode files in the input set: {files_with_non_markdown_extensions}"
-        )
-
-    logger.debug(f"Processing files: {all_markdown_files}")
-
-    filename_to_content_map = {}
-    for file in all_markdown_files:
-        with open(file, "r", encoding="utf8") as f:
-            try:
-                filename_to_content_map[file] = f.read()
-            except Exception as e:
-                logger.warning(f"Unable to read file: {file} as markdown. Skipping file.")
-                logger.warning(e, exc_info=True)
-
-    return filename_to_content_map
-
-
-def get_pdf_files(config: TextContentConfig):
-    # Extract required fields from config
-    pdf_files, pdf_file_filters = (
-        config.input_files,
-        config.input_filter,
-    )
-
-    # Input Validation
-    if is_none_or_empty(pdf_files) and is_none_or_empty(pdf_file_filters):
-        logger.debug("At least one of pdf-files or pdf-file-filter is required to be specified")
-        return {}
-
-    # Get PDF files to process
-    absolute_pdf_files, filtered_pdf_files = set(), set()
-    if pdf_files:
-        absolute_pdf_files = {get_absolute_path(pdf_file) for pdf_file in pdf_files}
-    if pdf_file_filters:
-        filtered_pdf_files = {
-            filtered_file
-            for pdf_file_filter in pdf_file_filters
-            for filtered_file in glob.glob(get_absolute_path(pdf_file_filter), recursive=True)
-            if os.path.isfile(filtered_file)
-        }
-
-    all_pdf_files = sorted(absolute_pdf_files | filtered_pdf_files)
-
-    files_with_non_pdf_extensions = {pdf_file for pdf_file in all_pdf_files if not pdf_file.endswith(".pdf")}
-
-    if any(files_with_non_pdf_extensions):
-        logger.warning(f"[Warning] There maybe non pdf-mode files in the input set: {files_with_non_pdf_extensions}")
-
-    logger.debug(f"Processing files: {all_pdf_files}")
-
-    filename_to_content_map = {}
-    for file in all_pdf_files:
-        with open(file, "rb") as f:
-            try:
-                filename_to_content_map[file] = f.read()
-            except Exception as e:
-                logger.warning(f"Unable to read file: {file} as PDF. Skipping file.")
-                logger.warning(e, exc_info=True)
-
-    return filename_to_content_map