khoj 1.26.2__py3-none-any.whl → 1.26.5.dev34__py3-none-any.whl

khoj/routers/helpers.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import hashlib
 import json
 import logging
@@ -14,6 +15,7 @@ from typing import (
     Annotated,
     Any,
     AsyncGenerator,
+    Callable,
     Dict,
     Iterator,
     List,
@@ -21,7 +23,7 @@ from typing import (
     Tuple,
     Union,
 )
-from urllib.parse import parse_qs, quote, urljoin, urlparse
+from urllib.parse import parse_qs, quote, unquote, urljoin, urlparse
 
 import cron_descriptor
 import pytz
@@ -30,11 +32,13 @@ from apscheduler.job import Job
 from apscheduler.triggers.cron import CronTrigger
 from asgiref.sync import sync_to_async
 from fastapi import Depends, Header, HTTPException, Request, UploadFile
+from pydantic import BaseModel
 from starlette.authentication import has_required_scope
 from starlette.requests import URL
 
 from khoj.database import adapters
 from khoj.database.adapters import (
+    LENGTH_OF_FREE_TRIAL,
     AgentAdapters,
     AutomationAdapters,
     ConversationAdapters,
@@ -215,6 +219,9 @@ def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="A
         elif chat["by"] == "khoj" and ("text-to-image" in chat["intent"].get("type")):
             chat_history += f"User: {chat['intent']['query']}\n"
             chat_history += f"{agent_name}: [generated image redacted for space]\n"
+        elif chat["by"] == "khoj" and ("excalidraw" in chat["intent"].get("type")):
+            chat_history += f"User: {chat['intent']['query']}\n"
+            chat_history += f"{agent_name}: {chat['intent']['inferred-queries'][0]}\n"
     return chat_history
 
 
@@ -235,6 +242,8 @@ def get_conversation_command(query: str, any_references: bool = False) -> Conver
         return ConversationCommand.AutomatedTask
     elif query.startswith("/summarize"):
         return ConversationCommand.Summarize
+    elif query.startswith("/diagram"):
+        return ConversationCommand.Diagram
     # If no relevant notes found for the given query
     elif not any_references:
         return ConversationCommand.General
@@ -290,7 +299,7 @@ async def aget_relevant_information_sources(
     conversation_history: dict,
     is_task: bool,
     user: KhojUser,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     agent: Agent = None,
 ):
     """
@@ -309,8 +318,8 @@ async def aget_relevant_information_sources(
 
     chat_history = construct_chat_history(conversation_history)
 
-    if uploaded_image_url:
-        query = f"[placeholder for user attached image]\n{query}"
+    if query_images:
+        query = f"[placeholder for {len(query_images)} user attached images]\n{query}"
 
     personality_context = (
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
@@ -367,7 +376,7 @@ async def aget_relevant_output_modes(
     conversation_history: dict,
     is_task: bool = False,
     user: KhojUser = None,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     agent: Agent = None,
 ):
     """
@@ -389,8 +398,8 @@ async def aget_relevant_output_modes(
 
     chat_history = construct_chat_history(conversation_history)
 
-    if uploaded_image_url:
-        query = f"[placeholder for user attached image]\n{query}"
+    if query_images:
+        query = f"[placeholder for {len(query_images)} user attached images]\n{query}"
 
     personality_context = (
         prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
@@ -433,7 +442,7 @@ async def infer_webpage_urls(
     conversation_history: dict,
     location_data: LocationData,
     user: KhojUser,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     agent: Agent = None,
 ) -> List[str]:
     """
@@ -459,7 +468,7 @@ async def infer_webpage_urls(
 
     with timer("Chat actor: Infer webpage urls to read", logger):
         response = await send_message_to_model_wrapper(
-            online_queries_prompt, uploaded_image_url=uploaded_image_url, response_type="json_object", user=user
+            online_queries_prompt, query_images=query_images, response_type="json_object", user=user
         )
 
     # Validate that the response is a non-empty, JSON-serializable list of URLs
@@ -479,7 +488,7 @@ async def generate_online_subqueries(
     conversation_history: dict,
     location_data: LocationData,
     user: KhojUser,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     agent: Agent = None,
 ) -> List[str]:
     """
@@ -505,7 +514,7 @@ async def generate_online_subqueries(
 
     with timer("Chat actor: Generate online search subqueries", logger):
         response = await send_message_to_model_wrapper(
-            online_queries_prompt, uploaded_image_url=uploaded_image_url, response_type="json_object", user=user
+            online_queries_prompt, query_images=query_images, response_type="json_object", user=user
         )
 
     # Validate that the response is a non-empty, JSON-serializable list
@@ -524,7 +533,7 @@ async def generate_online_subqueries(
 
 
 async def schedule_query(
-    q: str, conversation_history: dict, user: KhojUser, uploaded_image_url: str = None
+    q: str, conversation_history: dict, user: KhojUser, query_images: List[str] = None
 ) -> Tuple[str, ...]:
     """
     Schedule the date, time to run the query. Assume the server timezone is UTC.
@@ -537,7 +546,7 @@ async def schedule_query(
     )
 
     raw_response = await send_message_to_model_wrapper(
-        crontime_prompt, uploaded_image_url=uploaded_image_url, response_type="json_object", user=user
+        crontime_prompt, query_images=query_images, response_type="json_object", user=user
     )
 
     # Validate that the response is a non-empty, JSON-serializable list
@@ -583,7 +592,7 @@ async def extract_relevant_summary(
     q: str,
     corpus: str,
     conversation_history: dict,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
     user: KhojUser = None,
     agent: Agent = None,
 ) -> Union[str, None]:
@@ -612,11 +621,134 @@ async def extract_relevant_summary(
             extract_relevant_information,
             prompts.system_prompt_extract_relevant_summary,
             user=user,
-            uploaded_image_url=uploaded_image_url,
+            query_images=query_images,
         )
     return response.strip()
 
 
+async def generate_excalidraw_diagram(
+    q: str,
+    conversation_history: Dict[str, Any],
+    location_data: LocationData,
+    note_references: List[Dict[str, Any]],
+    online_results: Optional[dict] = None,
+    query_images: List[str] = None,
+    user: KhojUser = None,
+    agent: Agent = None,
+    send_status_func: Optional[Callable] = None,
+):
+    if send_status_func:
+        async for event in send_status_func("**Enhancing the Diagramming Prompt**"):
+            yield {ChatEvent.STATUS: event}
+
+    better_diagram_description_prompt = await generate_better_diagram_description(
+        q=q,
+        conversation_history=conversation_history,
+        location_data=location_data,
+        note_references=note_references,
+        online_results=online_results,
+        query_images=query_images,
+        user=user,
+        agent=agent,
+    )
+
+    if send_status_func:
+        async for event in send_status_func(f"**Diagram to Create:**:\n{better_diagram_description_prompt}"):
+            yield {ChatEvent.STATUS: event}
+
+    excalidraw_diagram_description = await generate_excalidraw_diagram_from_description(
+        q=better_diagram_description_prompt,
+        user=user,
+        agent=agent,
+    )
+
+    yield better_diagram_description_prompt, excalidraw_diagram_description
+
+
+async def generate_better_diagram_description(
+    q: str,
+    conversation_history: Dict[str, Any],
+    location_data: LocationData,
+    note_references: List[Dict[str, Any]],
+    online_results: Optional[dict] = None,
+    query_images: List[str] = None,
+    user: KhojUser = None,
+    agent: Agent = None,
+) -> str:
+    """
+    Generate a diagram description from the given query and context
+    """
+
+    today_date = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d, %A")
+    personality_context = (
+        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"
+
+    user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
+
+    chat_history = construct_chat_history(conversation_history)
+
+    simplified_online_results = {}
+
+    if online_results:
+        for result in online_results:
+            if online_results[result].get("answerBox"):
+                simplified_online_results[result] = online_results[result]["answerBox"]
+            elif online_results[result].get("webpages"):
+                simplified_online_results[result] = online_results[result]["webpages"]
+
+    improve_diagram_description_prompt = prompts.improve_diagram_description_prompt.format(
+        query=q,
+        chat_history=chat_history,
+        location=location_prompt,
+        current_date=today_date,
+        references=user_references,
+        online_results=simplified_online_results,
+        personality_context=personality_context,
+    )
+
+    with timer("Chat actor: Generate better diagram description", logger):
+        response = await send_message_to_model_wrapper(
+            improve_diagram_description_prompt, query_images=query_images, user=user
+        )
+        response = response.strip()
+        if response.startswith(('"', "'")) and response.endswith(('"', "'")):
+            response = response[1:-1]
+
+    return response
+
+
+async def generate_excalidraw_diagram_from_description(
+    q: str,
+    user: KhojUser = None,
+    agent: Agent = None,
+) -> str:
+    personality_context = (
+        prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
+    )
+
+    excalidraw_diagram_generation = prompts.excalidraw_diagram_generation_prompt.format(
+        personality_context=personality_context,
+        query=q,
+    )
+
+    with timer("Chat actor: Generate excalidraw diagram", logger):
+        raw_response = await send_message_to_model_wrapper(message=excalidraw_diagram_generation, user=user)
+        raw_response = raw_response.strip()
+        raw_response = remove_json_codeblock(raw_response)
+        response: Dict[str, str] = json.loads(raw_response)
+        if not response or not isinstance(response, List) or not isinstance(response[0], Dict):
+            # TODO Some additional validation here that it's a valid Excalidraw diagram
+            raise AssertionError(f"Invalid response for improving diagram description: {response}")
+
+    return response
+
+
 async def generate_better_image_prompt(
     q: str,
     conversation_history: str,
@@ -624,7 +756,7 @@ async def generate_better_image_prompt(
     note_references: List[Dict[str, Any]],
     online_results: Optional[dict] = None,
     model_type: Optional[str] = None,
-    uploaded_image_url: Optional[str] = None,
+    query_images: Optional[List[str]] = None,
     user: KhojUser = None,
     agent: Agent = None,
 ) -> str:
@@ -676,7 +808,7 @@ async def generate_better_image_prompt(
         )
 
     with timer("Chat actor: Generate contextual image prompt", logger):
-        response = await send_message_to_model_wrapper(image_prompt, uploaded_image_url=uploaded_image_url, user=user)
+        response = await send_message_to_model_wrapper(image_prompt, query_images=query_images, user=user)
         response = response.strip()
         if response.startswith(('"', "'")) and response.endswith(('"', "'")):
             response = response[1:-1]
@@ -689,11 +821,11 @@ async def send_message_to_model_wrapper(
     system_message: str = "",
     response_type: str = "text",
     user: KhojUser = None,
-    uploaded_image_url: str = None,
+    query_images: List[str] = None,
 ):
     conversation_config: ChatModelOptions = await ConversationAdapters.aget_default_conversation_config(user)
     vision_available = conversation_config.vision_enabled
-    if not vision_available and uploaded_image_url:
+    if not vision_available and query_images:
         vision_enabled_config = await ConversationAdapters.aget_vision_enabled_config()
         if vision_enabled_config:
             conversation_config = vision_enabled_config
@@ -746,7 +878,7 @@ async def send_message_to_model_wrapper(
             max_prompt_size=max_tokens,
             tokenizer_name=tokenizer,
             vision_enabled=vision_available,
-            uploaded_image_url=uploaded_image_url,
+            query_images=query_images,
             model_type=conversation_config.model_type,
         )
 
@@ -766,7 +898,7 @@ async def send_message_to_model_wrapper(
             max_prompt_size=max_tokens,
             tokenizer_name=tokenizer,
             vision_enabled=vision_available,
-            uploaded_image_url=uploaded_image_url,
+            query_images=query_images,
             model_type=conversation_config.model_type,
         )
 
@@ -784,7 +916,8 @@ async def send_message_to_model_wrapper(
             max_prompt_size=max_tokens,
             tokenizer_name=tokenizer,
             vision_enabled=vision_available,
-            uploaded_image_url=uploaded_image_url,
+            query_images=query_images,
+            model_type=conversation_config.model_type,
         )
 
         return gemini_send_message_to_model(
@@ -875,6 +1008,7 @@ def send_message_to_model_wrapper_sync(
             model_name=chat_model,
             max_prompt_size=max_tokens,
             vision_enabled=vision_available,
+            model_type=conversation_config.model_type,
         )
 
         return gemini_send_message_to_model(
@@ -900,7 +1034,7 @@ def generate_chat_response(
     conversation_id: str = None,
     location_data: LocationData = None,
     user_name: Optional[str] = None,
-    uploaded_image_url: Optional[str] = None,
+    query_images: Optional[List[str]] = None,
 ) -> Tuple[Union[ThreadedGenerator, Iterator[str]], Dict[str, str]]:
     # Initialize Variables
     chat_response = None
@@ -919,12 +1053,12 @@ def generate_chat_response(
             inferred_queries=inferred_queries,
             client_application=client_application,
             conversation_id=conversation_id,
-            uploaded_image_url=uploaded_image_url,
+            query_images=query_images,
         )
 
         conversation_config = ConversationAdapters.get_valid_conversation_config(user, conversation)
         vision_available = conversation_config.vision_enabled
-        if not vision_available and uploaded_image_url:
+        if not vision_available and query_images:
             vision_enabled_config = ConversationAdapters.get_vision_enabled_config()
             if vision_enabled_config:
                 conversation_config = vision_enabled_config
@@ -955,7 +1089,7 @@ def generate_chat_response(
             chat_response = converse(
                 compiled_references,
                 q,
-                image_url=uploaded_image_url,
+                query_images=query_images,
                 online_results=online_results,
                 conversation_log=meta_log,
                 model=chat_model,
@@ -993,8 +1127,9 @@ def generate_chat_response(
             chat_response = converse_gemini(
                 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,
@@ -1004,6 +1139,7 @@ def generate_chat_response(
                 location_data=location_data,
                 user_name=user_name,
                 agent=agent,
+                vision_available=vision_available,
             )
 
         metadata.update({"chat_model": conversation_config.chat_model})
@@ -1015,6 +1151,22 @@ def generate_chat_response(
     return chat_response, metadata
 
 
+class ChatRequestBody(BaseModel):
+    q: str
+    n: Optional[int] = 7
+    d: Optional[float] = None
+    stream: Optional[bool] = False
+    title: Optional[str] = None
+    conversation_id: Optional[str] = None
+    city: Optional[str] = None
+    region: Optional[str] = None
+    country: Optional[str] = None
+    country_code: Optional[str] = None
+    timezone: Optional[str] = None
+    images: Optional[list[str]] = None
+    create_new: Optional[bool] = False
+
+
 class ApiUserRateLimiter:
     def __init__(self, requests: int, subscribed_requests: int, window: int, slug: str):
         self.requests = requests
@@ -1060,13 +1212,58 @@ class ApiUserRateLimiter:
             )
             raise HTTPException(
                 status_code=429,
-                detail="We're glad you're enjoying Khoj! You've exceeded your usage limit for today. Come back tomorrow or subscribe to increase your usage limit via [your settings](https://app.khoj.dev/settings).",
+                detail="I'm glad you're enjoying interacting with me! But you've exceeded your usage limit for today. Come back tomorrow or subscribe to increase your usage limit via [your settings](https://app.khoj.dev/settings).",
             )
 
         # Add the current request to the cache
         UserRequests.objects.create(user=user, slug=self.slug)
 
 
+class ApiImageRateLimiter:
+    def __init__(self, max_images: int = 10, max_combined_size_mb: float = 10):
+        self.max_images = max_images
+        self.max_combined_size_mb = max_combined_size_mb
+
+    def __call__(self, request: Request, body: ChatRequestBody):
+        if state.billing_enabled is False:
+            return
+
+        # Rate limiting is disabled if user unauthenticated.
+        # Other systems handle authentication
+        if not request.user.is_authenticated:
+            return
+
+        if not body.images:
+            return
+
+        # Check number of images
+        if len(body.images) > self.max_images:
+            raise HTTPException(
+                status_code=429,
+                detail=f"Those are way too many images for me! I can handle up to {self.max_images} images per message.",
+            )
+
+        # Check total size of images
+        total_size_mb = 0.0
+        for image in body.images:
+            # Unquote the image in case it's URL encoded
+            image = unquote(image)
+            # Assuming the image is a base64 encoded string
+            # Remove the data:image/jpeg;base64, part if present
+            if "," in image:
+                image = image.split(",", 1)[1]
+
+            # Decode base64 to get the actual size
+            image_bytes = base64.b64decode(image)
+            total_size_mb += len(image_bytes) / (1024 * 1024)  # Convert bytes to MB
+
+        if total_size_mb > self.max_combined_size_mb:
+            raise HTTPException(
+                status_code=429,
+                detail=f"Those images are way too large for me! I can handle up to {self.max_combined_size_mb}MB of images per message.",
+            )
+
+
 class ConversationCommandRateLimiter:
     def __init__(self, trial_rate_limit: int, subscribed_rate_limit: int, slug: str):
         self.slug = slug
@@ -1477,10 +1674,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)
 
@@ -1553,6 +1756,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,
@@ -1561,6 +1765,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/routers/web_client.py

@@ -51,17 +51,6 @@ def chat_page(request: Request):
     )
 
 
-@web_client.get("/experimental", response_class=FileResponse)
-@requires(["authenticated"], redirect="login_page")
-def experimental_page(request: Request):
-    return templates.TemplateResponse(
-        "index.html",
-        context={
-            "request": request,
-        },
-    )
-
-
 @web_client.get("/factchecker", response_class=FileResponse)
 def fact_checker_page(request: Request):
     return templates.TemplateResponse(

khoj/utils/helpers.py

@@ -318,6 +318,7 @@ class ConversationCommand(str, Enum):
     Automation = "automation"
     AutomatedTask = "automated_task"
     Summarize = "summarize"
+    Diagram = "diagram"
 
 
 command_descriptions = {
@@ -326,10 +327,11 @@ command_descriptions = {
     ConversationCommand.Default: "The default command when no command specified. It intelligently auto-switches between general and notes mode.",
     ConversationCommand.Online: "Search for information on the internet.",
     ConversationCommand.Webpage: "Get information from webpage suggested by you.",
-    ConversationCommand.Image: "Generate images by describing your imagination in words.",
+    ConversationCommand.Image: "Generate illustrative, creative images by describing your imagination in words.",
     ConversationCommand.Automation: "Automatically run your query at a specified time or interval.",
     ConversationCommand.Help: "Get help with how to use or setup Khoj from the documentation",
     ConversationCommand.Summarize: "Get help with a question pertaining to an entire document.",
+    ConversationCommand.Diagram: "Draw a flowchart, diagram, or any other visual representation best expressed with primitives like lines, rectangles, and text.",
 }
 
 command_descriptions_for_agent = {
@@ -350,15 +352,17 @@ tool_descriptions_for_llm = {
 }
 
 mode_descriptions_for_llm = {
-    ConversationCommand.Image: "Use this if the user is requesting you to generate a picture based on their description.",
+    ConversationCommand.Image: "Use this if 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 the other response modes don't seem to fit the query.",
+    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.",
 }
 
 mode_descriptions_for_agent = {
     ConversationCommand.Image: "Agent can generate image in response.",
     ConversationCommand.Automation: "Agent can schedule a task to run at a scheduled date, time and frequency in response.",
     ConversationCommand.Text: "Agent can generate text in response.",
+    ConversationCommand.Diagram: "Agent can generate a visual representation that requires primitives like lines, rectangles, and text.",
 }
 
 

{khoj-1.26.2.dist-info → khoj-1.26.5.dev34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: khoj
-Version: 1.26.2
+Version: 1.26.5.dev34
 Summary: Your Second Brain
 Project-URL: Homepage, https://khoj.dev
 Project-URL: Documentation, https://docs.khoj.dev
@@ -56,7 +56,7 @@ Requires-Dist: pillow~=10.0.0
 Requires-Dist: psutil>=5.8.0
 Requires-Dist: psycopg2-binary==2.9.9
 Requires-Dist: pydantic[email]>=2.0.0
-Requires-Dist: pymupdf>=1.23.5
+Requires-Dist: pymupdf==1.24.11
 Requires-Dist: python-multipart>=0.0.7
 Requires-Dist: pytz~=2024.1
 Requires-Dist: pyyaml~=6.0