khoj 1.33.3.dev32__py3-none-any.whl

khoj/processor/embeddings.py

@@ -0,0 +1,122 @@
+import logging
+from typing import List
+
+import requests
+import tqdm
+from sentence_transformers import CrossEncoder, SentenceTransformer
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+from torch import nn
+
+from khoj.utils.helpers import fix_json_dict, get_device, merge_dicts, timer
+from khoj.utils.rawconfig import SearchResponse
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingsModel:
+    def __init__(
+        self,
+        model_name: str = "thenlper/gte-small",
+        embeddings_inference_endpoint: str = None,
+        embeddings_inference_endpoint_api_key: str = None,
+        query_encode_kwargs: dict = {},
+        docs_encode_kwargs: dict = {},
+        model_kwargs: dict = {},
+    ):
+        default_query_encode_kwargs = {"show_progress_bar": False, "normalize_embeddings": True}
+        default_docs_encode_kwargs = {"show_progress_bar": True, "normalize_embeddings": True}
+        self.query_encode_kwargs = merge_dicts(fix_json_dict(query_encode_kwargs), default_query_encode_kwargs)
+        self.docs_encode_kwargs = merge_dicts(fix_json_dict(docs_encode_kwargs), default_docs_encode_kwargs)
+        self.model_kwargs = merge_dicts(fix_json_dict(model_kwargs), {"device": get_device()})
+        self.model_name = model_name
+        self.inference_endpoint = embeddings_inference_endpoint
+        self.api_key = embeddings_inference_endpoint_api_key
+        with timer(f"Loaded embedding model {self.model_name}", logger):
+            self.embeddings_model = SentenceTransformer(self.model_name, **self.model_kwargs)
+
+    def inference_server_enabled(self) -> bool:
+        return self.api_key is not None and self.inference_endpoint is not None
+
+    def embed_query(self, query):
+        if self.inference_server_enabled():
+            return self.embed_with_api([query])[0]
+        return self.embeddings_model.encode([query], **self.query_encode_kwargs)[0]
+
+    @retry(
+        retry=retry_if_exception_type(requests.exceptions.HTTPError),
+        wait=wait_random_exponential(multiplier=1, max=10),
+        stop=stop_after_attempt(5),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+    )
+    def embed_with_api(self, docs):
+        payload = {"inputs": docs}
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        try:
+            response = requests.post(self.inference_endpoint, json=payload, headers=headers)
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            logger.error(
+                f" Error while calling inference endpoint {self.inference_endpoint} with error {e}, response {response.json()} ",
+                exc_info=True,
+            )
+            raise e
+        return response.json()["embeddings"]
+
+    def embed_documents(self, docs):
+        if self.inference_server_enabled():
+            if "huggingface" not in self.inference_endpoint:
+                logger.warning(
+                    f"Unsupported inference endpoint: {self.inference_endpoint}. Only HuggingFace supported. Generating embeddings on device instead."
+                )
+                return self.embeddings_model.encode(docs, **self.docs_encode_kwargs).tolist()
+            # break up the docs payload in chunks of 1000 to avoid hitting rate limits
+            embeddings = []
+            with tqdm.tqdm(total=len(docs)) as pbar:
+                for i in range(0, len(docs), 1000):
+                    docs_to_embed = docs[i : i + 1000]
+                    generated_embeddings = self.embed_with_api(docs_to_embed)
+                    embeddings += generated_embeddings
+                    pbar.update(1000)
+            return embeddings
+        return self.embeddings_model.encode(docs, **self.docs_encode_kwargs).tolist() if docs else []
+
+
+class CrossEncoderModel:
+    def __init__(
+        self,
+        model_name: str = "mixedbread-ai/mxbai-rerank-xsmall-v1",
+        cross_encoder_inference_endpoint: str = None,
+        cross_encoder_inference_endpoint_api_key: str = None,
+        model_kwargs: dict = {},
+    ):
+        self.model_name = model_name
+        self.inference_endpoint = cross_encoder_inference_endpoint
+        self.api_key = cross_encoder_inference_endpoint_api_key
+        self.model_kwargs = merge_dicts(model_kwargs, {"device": get_device()})
+        with timer(f"Loaded cross-encoder model {self.model_name}", logger):
+            self.cross_encoder_model = CrossEncoder(model_name=self.model_name, **self.model_kwargs)
+
+    def inference_server_enabled(self) -> bool:
+        return self.api_key is not None and self.inference_endpoint is not None
+
+    def predict(self, query, hits: List[SearchResponse], key: str = "compiled"):
+        if self.inference_server_enabled() and "huggingface" in self.inference_endpoint:
+            target_url = f"{self.inference_endpoint}"
+            payload = {"inputs": {"query": query, "passages": [hit.additional[key] for hit in hits]}}
+            headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
+            response = requests.post(target_url, json=payload, headers=headers)
+            response.raise_for_status()
+            return response.json()["scores"]
+
+        cross_inp = [[query, hit.additional[key]] for hit in hits]
+        cross_scores = self.cross_encoder_model.predict(cross_inp, activation_fct=nn.Sigmoid())
+        return cross_scores

khoj/processor/image/generate.py

@@ -0,0 +1,215 @@
+import base64
+import io
+import logging
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+import openai
+import requests
+
+from khoj.database.adapters import ConversationAdapters
+from khoj.database.models import Agent, KhojUser, TextToImageModelConfig
+from khoj.routers.helpers import ChatEvent, generate_better_image_prompt
+from khoj.routers.storage import upload_image
+from khoj.utils import state
+from khoj.utils.helpers import convert_image_to_webp, timer
+from khoj.utils.rawconfig import LocationData
+
+logger = logging.getLogger(__name__)
+
+
+async def text_to_image(
+    message: str,
+    user: KhojUser,
+    conversation_log: dict,
+    location_data: LocationData,
+    references: List[Dict[str, Any]],
+    online_results: Dict[str, Any],
+    send_status_func: Optional[Callable] = None,
+    query_images: Optional[List[str]] = None,
+    agent: Agent = None,
+    query_files: str = None,
+    tracer: dict = {},
+):
+    status_code = 200
+    image = None
+    image_url = None
+
+    text_to_image_config = await ConversationAdapters.aget_user_text_to_image_model(user)
+    if not text_to_image_config:
+        # If the user has not configured a text to image model, return an unsupported on server error
+        status_code = 501
+        message = "Failed to generate image. Setup image generation on the server."
+        yield image_url or image, status_code, message
+        return
+
+    text2image_model = text_to_image_config.model_name
+    chat_history = ""
+    for chat in conversation_log.get("chat", [])[-4:]:
+        if chat["by"] == "khoj" and chat["intent"].get("type") in ["remember", "reminder"]:
+            chat_history += f"Q: {chat['intent']['query']}\n"
+            chat_history += f"A: {chat['message']}\n"
+        elif chat["by"] == "khoj" and chat.get("images"):
+            chat_history += f"Q: {chat['intent']['query']}\n"
+            chat_history += f"A: Improved Prompt: {chat['intent']['inferred-queries'][0]}\n"
+
+    if send_status_func:
+        async for event in send_status_func("**Enhancing the Painting Prompt**"):
+            yield {ChatEvent.STATUS: event}
+
+    # Generate a better image prompt
+    # Use the user's message, chat history, and other context
+    image_prompt = await generate_better_image_prompt(
+        message,
+        chat_history,
+        location_data=location_data,
+        note_references=references,
+        online_results=online_results,
+        model_type=text_to_image_config.model_type,
+        query_images=query_images,
+        user=user,
+        agent=agent,
+        query_files=query_files,
+        tracer=tracer,
+    )
+
+    if send_status_func:
+        async for event in send_status_func(f"**Painting to Imagine**:\n{image_prompt}"):
+            yield {ChatEvent.STATUS: event}
+
+    # Generate image using the configured model and API
+    with timer(f"Generate image with {text_to_image_config.model_type}", logger):
+        try:
+            if text_to_image_config.model_type == TextToImageModelConfig.ModelType.OPENAI:
+                webp_image_bytes = generate_image_with_openai(image_prompt, text_to_image_config, text2image_model)
+            elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.STABILITYAI:
+                webp_image_bytes = generate_image_with_stability(image_prompt, text_to_image_config, text2image_model)
+            elif text_to_image_config.model_type == TextToImageModelConfig.ModelType.REPLICATE:
+                webp_image_bytes = generate_image_with_replicate(image_prompt, text_to_image_config, text2image_model)
+        except openai.OpenAIError or openai.BadRequestError or openai.APIConnectionError as e:
+            if "content_policy_violation" in e.message:
+                logger.error(f"Image Generation blocked by OpenAI: {e}")
+                status_code = e.status_code  # type: ignore
+                message = f"Image generation blocked by OpenAI due to policy violation"  # type: ignore
+                yield image_url or image, status_code, message
+                return
+            else:
+                logger.error(f"Image Generation failed with {e}", exc_info=True)
+                message = f"Image generation failed using OpenAI"  # type: ignore
+                status_code = e.status_code  # type: ignore
+                yield image_url or image, status_code, message
+                return
+        except requests.RequestException as e:
+            logger.error(f"Image Generation failed with {e}", exc_info=True)
+            message = f"Image generation using {text2image_model} via {text_to_image_config.model_type} failed due to a network error."
+            status_code = 502
+            yield image_url or image, status_code, message
+            return
+
+    # Decide how to store the generated image
+    with timer("Upload image to S3", logger):
+        image_url = upload_image(webp_image_bytes, user.uuid)
+
+    if not image_url:
+        image = base64.b64encode(webp_image_bytes).decode("utf-8")
+
+    yield image_url or image, status_code, image_prompt
+
+
+def generate_image_with_openai(
+    improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str
+):
+    "Generate image using OpenAI API"
+
+    # Get the API key from the user's configuration
+    if text_to_image_config.api_key:
+        api_key = text_to_image_config.api_key
+    elif text_to_image_config.ai_model_api:
+        api_key = text_to_image_config.ai_model_api.api_key
+    elif state.openai_client:
+        api_key = state.openai_client.api_key
+    auth_header = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+
+    # Generate image using OpenAI API
+    OPENAI_IMAGE_GEN_STYLE = "vivid"
+    response = state.openai_client.images.generate(
+        prompt=improved_image_prompt,
+        model=text2image_model,
+        style=OPENAI_IMAGE_GEN_STYLE,
+        response_format="b64_json",
+        extra_headers=auth_header,
+    )
+
+    # Extract the base64 image from the response
+    image = response.data[0].b64_json
+    # Decode base64 png and convert it to webp for faster loading
+    return convert_image_to_webp(base64.b64decode(image))
+
+
+def generate_image_with_stability(
+    improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str
+):
+    "Generate image using Stability AI"
+
+    # Call Stability AI API to generate image
+    response = requests.post(
+        f"https://api.stability.ai/v2beta/stable-image/generate/sd3",
+        headers={"authorization": f"Bearer {text_to_image_config.api_key}", "accept": "image/*"},
+        files={"none": ""},
+        data={
+            "prompt": improved_image_prompt,
+            "model": text2image_model,
+            "mode": "text-to-image",
+            "output_format": "png",
+            "aspect_ratio": "1:1",
+        },
+    )
+    # Convert png to webp for faster loading
+    return convert_image_to_webp(response.content)
+
+
+def generate_image_with_replicate(
+    improved_image_prompt: str, text_to_image_config: TextToImageModelConfig, text2image_model: str
+):
+    "Generate image using Replicate API"
+
+    # Create image generation task on Replicate
+    replicate_create_prediction_url = f"https://api.replicate.com/v1/models/{text2image_model}/predictions"
+    headers = {
+        "Authorization": f"Bearer {text_to_image_config.api_key}",
+        "Content-Type": "application/json",
+    }
+    json = {
+        "input": {
+            "prompt": improved_image_prompt,
+            "num_outputs": 1,
+            "aspect_ratio": "1:1",
+            "output_format": "webp",
+            "output_quality": 100,
+        }
+    }
+    create_prediction = requests.post(replicate_create_prediction_url, headers=headers, json=json).json()
+
+    # Get status of image generation task
+    get_prediction_url = create_prediction["urls"]["get"]
+    get_prediction = requests.get(get_prediction_url, headers=headers).json()
+    status = get_prediction["status"]
+    retry_count = 1
+
+    # Poll the image generation task for completion status
+    while status not in ["succeeded", "failed", "canceled"] and retry_count < 20:
+        time.sleep(2)
+        get_prediction = requests.get(get_prediction_url, headers=headers).json()
+        status = get_prediction["status"]
+        retry_count += 1
+
+    # Raise exception if the image generation task fails
+    if status != "succeeded":
+        error = get_prediction.get("error")
+        if retry_count >= 10:
+            raise requests.RequestException("Image generation timed out")
+        raise requests.RequestException(f"Image generation failed with status: {status}, message: {error}")
+
+    # Get the generated image
+    image_url = get_prediction["output"][0] if isinstance(get_prediction["output"], list) else get_prediction["output"]
+    return io.BytesIO(requests.get(image_url).content).getvalue()

khoj/processor/speech/text_to_speech.py

@@ -0,0 +1,51 @@
+import json  # Used for working with JSON data
+import os
+
+import requests  # Used for making HTTP requests
+from bs4 import BeautifulSoup
+from markdown_it import MarkdownIt
+
+# Define constants for the script
+CHUNK_SIZE = 1024  # Size of chunks to read/write at a time
+ELEVEN_LABS_API_KEY = os.getenv("ELEVEN_LABS_API_KEY", None)  # Your API key for authentication
+VOICE_ID = "RPEIZnKMqlQiZyZd1Dae"  # ID of the voice model to use. MALE - Christopher - friendly guy next door.
+ELEVEN_API_URL = "https://api.elevenlabs.io/v1/text-to-speech"  # Base URL for the Text-to-Speech API
+
+markdown_renderer = MarkdownIt()
+
+
+def is_eleven_labs_enabled():
+    return ELEVEN_LABS_API_KEY is not None
+
+
+def generate_text_to_speech(
+    text_to_speak: str,
+    voice_id: str = VOICE_ID,
+):
+    if not is_eleven_labs_enabled():
+        return "Eleven Labs API key is not set"
+
+    # Convert the incoming text from markdown format to plain text
+    html = markdown_renderer.render(text_to_speak)
+    text = "".join(BeautifulSoup(html, features="lxml").findAll(text=True))
+
+    # Construct the URL for the Text-to-Speech API request
+    tts_url = f"{ELEVEN_API_URL}/{voice_id}/stream"
+
+    # Set up headers for the API request, including the API key for authentication
+    headers = {"Accept": "application/json", "xi-api-key": ELEVEN_LABS_API_KEY}
+
+    # Set up the data payload for the API request, including the text and voice settings
+    data = {
+        "text": text,
+        # "model_id": "eleven_multilingual_v2",
+        "voice_settings": {"stability": 0.5, "similarity_boost": 0.8, "style": 0.0, "use_speaker_boost": True},
+    }
+
+    # Make the POST request to the TTS API with headers and data, enabling streaming response
+    response = requests.post(tts_url, headers=headers, json=data, stream=True)
+
+    if response.ok:
+        return response
+    else:
+        raise Exception(f"Failed to generate text-to-speech: {response.text}")