khoj 1.26.5.dev32__py3-none-any.whl → 1.27.2.dev1__py3-none-any.whl

khoj/processor/conversation/anthropic/anthropic_chat.py

@@ -11,8 +11,12 @@ from khoj.processor.conversation import prompts
 from khoj.processor.conversation.anthropic.utils import (
     anthropic_chat_completion_with_backoff,
     anthropic_completion_with_backoff,
+    format_messages_for_anthropic,
+)
+from khoj.processor.conversation.utils import (
+    construct_structured_message,
+    generate_chatml_messages_with_context,
 )
-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
 
@@ -27,6 +31,8 @@ def extract_questions_anthropic(
     temperature=0.7,
     location_data: LocationData = None,
     user: KhojUser = None,
+    query_images: Optional[list[str]] = None,
+    vision_enabled: bool = False,
     personality_context: Optional[str] = None,
 ):
     """
@@ -68,6 +74,13 @@ def extract_questions_anthropic(
         text=text,
     )
 
+    prompt = construct_structured_message(
+        message=prompt,
+        images=query_images,
+        model_type=ChatModelOptions.ModelType.ANTHROPIC,
+        vision_enabled=vision_enabled,
+    )
+
     messages = [ChatMessage(content=prompt, role="user")]
 
     response = anthropic_completion_with_backoff(
@@ -101,17 +114,7 @@ def anthropic_send_message_to_model(messages, api_key, model):
     """
     Send message to model
     """
-    # Anthropic requires the first message to be a 'user' message, and the system prompt is not to be sent in the messages parameter
-    system_prompt = None
-
-    if len(messages) == 1:
-        messages[0].role = "user"
-    else:
-        system_prompt = ""
-        for message in messages.copy():
-            if message.role == "system":
-                system_prompt += message.content
-                messages.remove(message)
+    messages, system_prompt = format_messages_for_anthropic(messages)
 
     # Get Response from GPT. Don't use response_type because Anthropic doesn't support it.
     return anthropic_completion_with_backoff(
@@ -127,7 +130,7 @@ def converse_anthropic(
     user_query,
     online_results: Optional[Dict[str, Dict]] = None,
     conversation_log={},
-    model: Optional[str] = "claude-instant-1.2",
+    model: Optional[str] = "claude-3-5-sonnet-20241022",
     api_key: Optional[str] = None,
     completion_func=None,
     conversation_commands=[ConversationCommand.Default],
@@ -136,6 +139,8 @@ def converse_anthropic(
     location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
+    query_images: Optional[list[str]] = None,
+    vision_available: bool = False,
 ):
     """
     Converse with user using Anthropic's Claude
@@ -189,17 +194,12 @@ def converse_anthropic(
         model_name=model,
         max_prompt_size=max_prompt_size,
         tokenizer_name=tokenizer_name,
+        query_images=query_images,
+        vision_enabled=vision_available,
         model_type=ChatModelOptions.ModelType.ANTHROPIC,
     )
 
-    if len(messages) > 1:
-        if messages[0].role == "assistant":
-            messages = messages[1:]
-
-    for message in messages.copy():
-        if message.role == "system":
-            system_prompt += message.content
-            messages.remove(message)
+    messages, system_prompt = format_messages_for_anthropic(messages, system_prompt)
 
     truncated_messages = "\n".join({f"{message.content[:40]}..." for message in messages})
     logger.debug(f"Conversation Context for Claude: {truncated_messages}")

khoj/processor/conversation/anthropic/utils.py

@@ -3,6 +3,7 @@ from threading import Thread
 from typing import Dict, List
 
 import anthropic
+from langchain.schema import ChatMessage
 from tenacity import (
     before_sleep_log,
     retry,
@@ -11,7 +12,8 @@ from tenacity import (
     wait_random_exponential,
 )
 
-from khoj.processor.conversation.utils import ThreadedGenerator
+from khoj.processor.conversation.utils import ThreadedGenerator, get_image_from_url
+from khoj.utils.helpers import is_none_or_empty
 
 logger = logging.getLogger(__name__)
 
@@ -115,3 +117,51 @@ def anthropic_llm_thread(
         logger.error(f"Error in anthropic_llm_thread: {e}", exc_info=True)
     finally:
         g.close()
+
+
+def format_messages_for_anthropic(messages: list[ChatMessage], system_prompt=None):
+    """
+    Format messages for Anthropic
+    """
+    # Extract system prompt
+    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
+
+    # Anthropic requires the first message to be a 'user' message
+    if len(messages) == 1:
+        messages[0].role = "user"
+    elif len(messages) > 1 and messages[0].role == "assistant":
+        messages = messages[1:]
+
+    # Convert image urls to base64 encoded images in Anthropic message format
+    for message in messages:
+        if isinstance(message.content, list):
+            content = []
+            # Sort the content. Anthropic models prefer that text comes after images.
+            message.content.sort(key=lambda x: 0 if x["type"] == "image_url" else 1)
+            for idx, part in enumerate(message.content):
+                if part["type"] == "text":
+                    content.append({"type": "text", "text": part["text"]})
+                elif part["type"] == "image_url":
+                    image = get_image_from_url(part["image_url"]["url"], type="b64")
+                    # Prefix each image with text block enumerating the image number
+                    # This helps the model reference the image in its response. Recommended by Anthropic
+                    content.extend(
+                        [
+                            {
+                                "type": "text",
+                                "text": f"Image {idx + 1}:",
+                            },
+                            {
+                                "type": "image",
+                                "source": {"type": "base64", "media_type": image.type, "data": image.content},
+                            },
+                        ]
+                    )
+            message.content = content
+
+    return messages, system_prompt

khoj/processor/conversation/google/utils.py

@@ -1,11 +1,8 @@
 import logging
 import random
-from io import BytesIO
 from threading import Thread
 
 import google.generativeai as genai
-import PIL.Image
-import requests
 from google.generativeai.types.answer_types import FinishReason
 from google.generativeai.types.generation_types import StopCandidateException
 from google.generativeai.types.safety_types import (
@@ -22,7 +19,7 @@ from tenacity import (
     wait_random_exponential,
 )
 
-from khoj.processor.conversation.utils import ThreadedGenerator
+from khoj.processor.conversation.utils import ThreadedGenerator, get_image_from_url
 from khoj.utils.helpers import is_none_or_empty
 
 logger = logging.getLogger(__name__)
@@ -207,7 +204,7 @@ def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str =
         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"]) if item["type"] == "image_url" else item["text"]
+                get_image_from_url(item["image_url"]["url"]).content if item["type"] == "image_url" else item["text"]
                 for item in sorted(message.content, key=lambda x: 0 if x["type"] == "image_url" else 1)
             ]
         elif isinstance(message.content, str):
@@ -220,13 +217,3 @@ def format_messages_for_gemini(messages: list[ChatMessage], system_prompt: str =
         messages[0].role = "user"
 
     return messages, system_prompt
-
-
-def get_image_from_url(image_url: str) -> PIL.Image:
-    try:
-        response = requests.get(image_url)
-        response.raise_for_status()  # Check if the request was successful
-        return PIL.Image.open(BytesIO(response.content))
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Failed to get image from URL {image_url}: {e}")
-        return None

khoj/processor/conversation/prompts.py

@@ -619,7 +619,7 @@ AI: It's currently 28°C and partly cloudy in Bali.
 Q: Share a painting using the weather for Bali every morning.
 Khoj: {{"output": "automation"}}
 
-Now it's your turn to pick the mode you would like to use to answer the user's question. Provide your response as a JSON.
+Now it's your turn to pick the mode you would like to use to answer the user's question. Provide your response as a JSON. Do not say anything else.
 
 Chat History:
 {chat_history}

khoj/processor/conversation/utils.py

@@ -1,10 +1,16 @@
+import base64
 import logging
 import math
+import mimetypes
 import queue
+from dataclasses import dataclass
 from datetime import datetime
+from io import BytesIO
 from time import perf_counter
 from typing import Any, Dict, List, Optional
 
+import PIL.Image
+import requests
 import tiktoken
 from langchain.schema import ChatMessage
 from llama_cpp.llama import Llama
@@ -152,7 +158,11 @@ def construct_structured_message(message: str, images: list[str], model_type: st
     if not images or not vision_enabled:
         return message
 
-    if model_type in [ChatModelOptions.ModelType.OPENAI, ChatModelOptions.ModelType.GOOGLE]:
+    if model_type in [
+        ChatModelOptions.ModelType.OPENAI,
+        ChatModelOptions.ModelType.GOOGLE,
+        ChatModelOptions.ModelType.ANTHROPIC,
+    ]:
         return [
             {"type": "text", "text": message},
             *[{"type": "image_url", "image_url": {"url": image}} for image in images],
@@ -306,3 +316,31 @@ def reciprocal_conversation_to_chatml(message_pair):
 def remove_json_codeblock(response: str):
     """Remove any markdown json codeblock formatting if present. Useful for non schema enforceable models"""
     return response.removeprefix("```json").removesuffix("```")
+
+
+@dataclass
+class ImageWithType:
+    content: Any
+    type: str
+
+
+def get_image_from_url(image_url: str, type="pil"):
+    try:
+        response = requests.get(image_url)
+        response.raise_for_status()  # Check if the request was successful
+
+        # Get content type from response or infer from URL
+        content_type = response.headers.get("content-type") or mimetypes.guess_type(image_url)[0] or "image/webp"
+
+        # Convert image to desired format
+        if type == "b64":
+            image_data = base64.b64encode(response.content).decode("utf-8")
+        elif type == "pil":
+            image_data = PIL.Image.open(BytesIO(response.content))
+        else:
+            raise ValueError(f"Invalid image type: {type}")
+
+        return ImageWithType(content=image_data, type=content_type)
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to get image from URL {image_url}: {e}")
+        return ImageWithType(content=None, type=None)

khoj/processor/image/generate.py

@@ -204,9 +204,10 @@ def generate_image_with_replicate(
 
     # 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}")
+        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"]

khoj/routers/api.py

@@ -25,8 +25,9 @@ from khoj.database.adapters import (
     AutomationAdapters,
     ConversationAdapters,
     EntryAdapters,
+    get_default_search_model,
+    get_user_default_search_model,
     get_user_photo,
-    get_user_search_model_or_default,
 )
 from khoj.database.models import (
     Agent,
@@ -149,7 +150,7 @@ async def execute_search(
     encoded_asymmetric_query = None
     if t != SearchType.Image:
         with timer("Encoding query took", logger=logger):
-            search_model = await sync_to_async(get_user_search_model_or_default)(user)
+            search_model = await sync_to_async(get_user_default_search_model)(user)
             encoded_asymmetric_query = state.embeddings_model[search_model.name].embed_query(defiltered_query)
 
     with concurrent.futures.ThreadPoolExecutor() as executor:
@@ -447,11 +448,13 @@ async def extract_references_and_questions(
             chat_model = conversation_config.chat_model
             inferred_queries = extract_questions_anthropic(
                 defiltered_query,
+                query_images=query_images,
                 model=chat_model,
                 api_key=api_key,
                 conversation_log=meta_log,
                 location_data=location_data,
                 user=user,
+                vision_enabled=vision_enabled,
                 personality_context=personality_context,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:

khoj/routers/api_model.py

@@ -94,39 +94,6 @@ async def update_voice_model(
     return Response(status_code=202, content=json.dumps({"status": "ok"}))
 
 
-@api_model.post("/search", status_code=200)
-@requires(["authenticated"])
-async def update_search_model(
-    request: Request,
-    id: str,
-    client: Optional[str] = None,
-):
-    user = request.user.object
-
-    prev_config = await adapters.aget_user_search_model(user)
-    new_config = await adapters.aset_user_search_model(user, int(id))
-
-    if prev_config and int(id) != prev_config.id and new_config:
-        await EntryAdapters.adelete_all_entries(user)
-
-    if not prev_config:
-        # If the use was just using the default config, delete all the entries and set the new config.
-        await EntryAdapters.adelete_all_entries(user)
-
-    if new_config is None:
-        return {"status": "error", "message": "Model not found"}
-    else:
-        update_telemetry_state(
-            request=request,
-            telemetry_type="api",
-            api="set_search_model",
-            client=client,
-            metadata={"search_model": new_config.setting.name},
-        )
-
-    return {"status": "ok"}
-
-
 @api_model.post("/paint", status_code=200)
 @requires(["authenticated"])
 async def update_paint_model(

khoj/routers/{subscription.py → api_subscription.py}

@@ -1,12 +1,14 @@
+import json
 import logging
 import os
 from datetime import datetime, timezone
 
 from asgiref.sync import sync_to_async
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Request, Response
 from starlette.authentication import requires
 
 from khoj.database import adapters
+from khoj.database.models import KhojUser, Subscription
 from khoj.routers.helpers import update_telemetry_state
 from khoj.utils import state
 
@@ -73,7 +75,7 @@ async def subscribe(request: Request):
     elif event_type in {"customer.subscription.deleted"}:
         # Reset the user to trial state
         user, is_new = await adapters.set_user_subscription(
-            customer_email, is_recurring=False, renewal_date=False, type="trial"
+            customer_email, is_recurring=False, renewal_date=False, type=Subscription.Type.TRIAL
         )
         success = user is not None
 
@@ -116,3 +118,19 @@ async def update_subscription(request: Request, email: str, operation: str):
         return {"success": False, "message": "No subscription found that is set to cancel"}
 
     return {"success": False, "message": "Invalid operation"}
+
+
+@subscription_router.post("/trial", response_class=Response)
+@requires(["authenticated"])
+async def start_trial(request: Request) -> Response:
+    user: KhojUser = request.user.object
+
+    # Start a trial for the user
+    updated_subscription = await adapters.astart_trial_subscription(user)
+
+    # Return trial status as a JSON response
+    return Response(
+        content=json.dumps({"trial_enabled": updated_subscription is not None}),
+        media_type="application/json",
+        status_code=200,
+    )

khoj/routers/helpers.py

@@ -38,6 +38,7 @@ from starlette.requests import URL
 
 from khoj.database import adapters
 from khoj.database.adapters import (
+    LENGTH_OF_FREE_TRIAL,
     AgentAdapters,
     AutomationAdapters,
     ConversationAdapters,
@@ -683,10 +684,7 @@ async def generate_better_diagram_description(
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
     )
 
-    if location_data:
-        location_prompt = prompts.user_location.format(location=f"{location_data}")
-    else:
-        location_prompt = "Unknown"
+    location = f"{location_data}" if location_data else "Unknown"
 
     user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
 
@@ -704,7 +702,7 @@ async def generate_better_diagram_description(
     improve_diagram_description_prompt = prompts.improve_diagram_description_prompt.format(
         query=q,
         chat_history=chat_history,
-        location=location_prompt,
+        location=location,
         current_date=today_date,
         references=user_references,
         online_results=simplified_online_results,
@@ -769,10 +767,7 @@ async def generate_better_image_prompt(
     )
     model_type = model_type or TextToImageModelConfig.ModelType.OPENAI
 
-    if location_data:
-        location_prompt = prompts.user_location.format(location=f"{location_data}")
-    else:
-        location_prompt = "Unknown"
+    location = f"{location_data}" if location_data else "Unknown"
 
     user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
 
@@ -789,7 +784,7 @@ async def generate_better_image_prompt(
         image_prompt = prompts.image_generation_improve_prompt_dalle.format(
             query=q,
             chat_history=conversation_history,
-            location=location_prompt,
+            location=location,
             current_date=today_date,
             references=user_references,
             online_results=simplified_online_results,
@@ -799,7 +794,7 @@ async def generate_better_image_prompt(
         image_prompt = prompts.image_generation_improve_prompt_sd.format(
             query=q,
             chat_history=conversation_history,
-            location=location_prompt,
+            location=location,
             current_date=today_date,
             references=user_references,
             online_results=simplified_online_results,
@@ -825,10 +820,13 @@ async def send_message_to_model_wrapper(
     conversation_config: ChatModelOptions = await ConversationAdapters.aget_default_conversation_config(user)
     vision_available = conversation_config.vision_enabled
     if not vision_available and query_images:
+        logger.warning(f"Vision is not enabled for default model: {conversation_config.chat_model}.")
         vision_enabled_config = await ConversationAdapters.aget_vision_enabled_config()
         if vision_enabled_config:
             conversation_config = vision_enabled_config
             vision_available = True
+    if vision_available and query_images:
+        logger.info(f"Using {conversation_config.chat_model} model to understand {len(query_images)} images.")
 
     subscribed = await ais_user_subscribed(user)
     chat_model = conversation_config.chat_model
@@ -1109,8 +1107,9 @@ def generate_chat_response(
             chat_response = converse_anthropic(
                 compiled_references,
                 q,
-                online_results,
-                meta_log,
+                query_images=query_images,
+                online_results=online_results,
+                conversation_log=meta_log,
                 model=conversation_config.chat_model,
                 api_key=api_key,
                 completion_func=partial_completion,
@@ -1120,6 +1119,7 @@ def generate_chat_response(
                 location_data=location_data,
                 user_name=user_name,
                 agent=agent,
+                vision_available=vision_available,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
             api_key = conversation_config.openai_config.api_key
@@ -1673,10 +1673,16 @@ def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False)
 
     user_subscription_state = get_user_subscription_state(user.email)
     user_subscription = adapters.get_user_subscription(user.email)
+
     subscription_renewal_date = (
         user_subscription.renewal_date.strftime("%d %b %Y")
         if user_subscription and user_subscription.renewal_date
-        else (user_subscription.created_at + timedelta(days=7)).strftime("%d %b %Y")
+        else None
+    )
+    subscription_enabled_trial_at = (
+        user_subscription.enabled_trial_at.strftime("%d %b %Y")
+        if user_subscription and user_subscription.enabled_trial_at
+        else None
     )
     given_name = get_user_name(user)
 
@@ -1699,13 +1705,6 @@ def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False)
     for chat_model in chat_models:
         chat_model_options.append({"name": chat_model.chat_model, "id": chat_model.id})
 
-    search_model_options = adapters.get_or_create_search_models().all()
-    all_search_model_options = list()
-    for search_model_option in search_model_options:
-        all_search_model_options.append({"name": search_model_option.name, "id": search_model_option.id})
-
-    current_search_model_option = adapters.get_user_search_model_or_default(user)
-
     selected_paint_model_config = ConversationAdapters.get_user_text_to_image_model_config(user)
     paint_model_options = ConversationAdapters.get_text_to_image_model_options().all()
     all_paint_model_options = list()
@@ -1738,8 +1737,6 @@ def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False)
         "has_documents": has_documents,
         "notion_token": notion_token,
         # user model settings
-        "search_model_options": all_search_model_options,
-        "selected_search_model_config": current_search_model_option.id,
         "chat_model_options": chat_model_options,
         "selected_chat_model_config": selected_chat_model_config.id if selected_chat_model_config else None,
         "paint_model_options": all_paint_model_options,
@@ -1749,6 +1746,7 @@ def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False)
         # user billing info
         "subscription_state": user_subscription_state,
         "subscription_renewal_date": subscription_renewal_date,
+        "subscription_enabled_trial_at": subscription_enabled_trial_at,
         # server settings
         "khoj_cloud_subscription_url": os.getenv("KHOJ_CLOUD_SUBSCRIPTION_URL"),
         "billing_enabled": state.billing_enabled,
@@ -1757,6 +1755,7 @@ def get_user_config(user: KhojUser, request: Request, is_detailed: bool = False)
         "khoj_version": state.khoj_version,
         "anonymous_mode": state.anonymous_mode,
         "notion_oauth_url": notion_oauth_url,
+        "length_of_free_trial": LENGTH_OF_FREE_TRIAL,
     }
 
 

khoj/search_type/text_search.py

@@ -8,7 +8,11 @@ import torch
 from asgiref.sync import sync_to_async
 from sentence_transformers import util
 
-from khoj.database.adapters import EntryAdapters, get_user_search_model_or_default
+from khoj.database.adapters import (
+    EntryAdapters,
+    get_default_search_model,
+    get_user_default_search_model,
+)
 from khoj.database.models import Agent
 from khoj.database.models import Entry as DbEntry
 from khoj.database.models import KhojUser
@@ -110,7 +114,7 @@ async def query(
     file_type = search_type_to_embeddings_type[type.value]
 
     query = raw_query
-    search_model = await sync_to_async(get_user_search_model_or_default)(user)
+    search_model = await sync_to_async(get_user_default_search_model)(user)
     if not max_distance:
         if search_model.bi_encoder_confidence_threshold:
             max_distance = search_model.bi_encoder_confidence_threshold

khoj/utils/helpers.py

@@ -352,7 +352,7 @@ tool_descriptions_for_llm = {
 }
 
 mode_descriptions_for_llm = {
-    ConversationCommand.Image: "Use this if the user is requesting you to create a new picture based on their description.",
+    ConversationCommand.Image: "Use this if you are confident the user is requesting you to create a new picture based on their description.",
     ConversationCommand.Automation: "Use this if you are confident the user is requesting a response at a scheduled date, time and frequency",
     ConversationCommand.Text: "Use this if a normal text response would be sufficient for accurately responding to the query.",
     ConversationCommand.Diagram: "Use this if the user is requesting a visual representation that requires primitives like lines, rectangles, and text.",

{khoj-1.26.5.dev32.dist-info → khoj-1.27.2.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: khoj
-Version: 1.26.5.dev32
+Version: 1.27.2.dev1
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev