khoj 1.28.4.dev23__py3-none-any.whl → 1.28.4.dev77__py3-none-any.whl

khoj/processor/conversation/offline/chat_model.py

@@ -37,6 +37,7 @@ def extract_questions_offline(
     max_prompt_size: int = None,
     temperature: float = 0.7,
     personality_context: Optional[str] = None,
+    query_files: str = None,
     tracer: dict = {},
 ) -> List[str]:
     """
@@ -87,6 +88,7 @@ def extract_questions_offline(
         loaded_model=offline_chat_model,
         max_prompt_size=max_prompt_size,
         model_type=ChatModelOptions.ModelType.OFFLINE,
+        query_files=query_files,
     )
 
     state.chat_lock.acquire()
@@ -152,6 +154,7 @@ def converse_offline(
     location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
+    query_files: str = None,
     tracer: dict = {},
 ) -> Union[ThreadedGenerator, Iterator[str]]:
     """
@@ -216,6 +219,7 @@ def converse_offline(
         max_prompt_size=max_prompt_size,
         tokenizer_name=tokenizer_name,
         model_type=ChatModelOptions.ModelType.OFFLINE,
+        query_files=query_files,
     )
 
     truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})

khoj/processor/conversation/openai/gpt.py

@@ -34,6 +34,7 @@ def extract_questions(
     query_images: Optional[list[str]] = None,
     vision_enabled: bool = False,
     personality_context: Optional[str] = None,
+    query_files: str = None,
     tracer: dict = {},
 ):
     """
@@ -79,9 +80,11 @@ def extract_questions(
         images=query_images,
         model_type=ChatModelOptions.ModelType.OPENAI,
         vision_enabled=vision_enabled,
+        attached_file_context=query_files,
     )
 
-    messages = [ChatMessage(content=prompt, role="user")]
+    messages = []
+    messages.append(ChatMessage(content=prompt, role="user"))
 
     response = send_message_to_model(
         messages,
@@ -148,6 +151,7 @@ def converse(
     agent: Agent = None,
     query_images: Optional[list[str]] = None,
     vision_available: bool = False,
+    query_files: str = None,
     tracer: dict = {},
 ):
     """
@@ -206,6 +210,7 @@ def converse(
         query_images=query_images,
         vision_enabled=vision_available,
         model_type=ChatModelOptions.ModelType.OPENAI,
+        query_files=query_files,
     )
     truncated_messages = "\n".join({f"{message.content[:70]}..." for message in messages})
     logger.debug(f"Conversation Context for GPT: {truncated_messages}")

khoj/processor/conversation/prompts.py

@@ -988,16 +988,27 @@ You are an extremely smart and helpful title generator assistant. Given a user q
 
 # Examples:
 User: Show a new Calvin and Hobbes quote every morning at 9am. My Current Location: Shanghai, China
-Khoj: Your daily Calvin and Hobbes Quote
+Assistant: Your daily Calvin and Hobbes Quote
 
 User: Notify me when version 2.0.0 of the sentence transformers python package is released. My Current Location: Mexico City, Mexico
-Khoj: Sentence Transformers Python Package Version 2.0.0 Release
+Assistant: Sentence Transformers Python Package Version 2.0.0 Release
 
 User: Gather the latest tech news on the first sunday of every month.
-Khoj: Your Monthly Dose of Tech News
+Assistant: Your Monthly Dose of Tech News
 
 User Query: {query}
-Khoj:
+Assistant:
+""".strip()
+)
+
+conversation_title_generation = PromptTemplate.from_template(
+    """
+You are an extremely smart and helpful title generator assistant. Given a conversation, extract the subject of the conversation. Crisp, informative, ten words or less.
+
+Conversation History:
+{chat_history}
+
+Assistant:
 """.strip()
 )
 
@@ -1038,6 +1049,39 @@ A:
 """.strip()
 )
 
+
+personality_prompt_safety_expert_lax = PromptTemplate.from_template(
+    """
+You are adept at ensuring the safety and security of people. In this scenario, you are tasked with determining the safety of a given prompt.
+
+Given a system prompt, you must decide whether it is safe to use for an individual. If the prompt is safe, respond with "Yes". If the prompt is not safe, respond with "No". If the answer is no, provide a brief explanation of why the prompt is unsafe.
+
+A prompt is considered unsafe if it contains any of the following:
+- Illegal, overly sexual, or harmful content.
+- Any content that could be harmful to the user or others.
+- Racism, sexism, or any other form of discrimination.
+- Any content that could be considered illegal or a misuse of AI.
+
+Referencing personal notes or data is perfectly acceptable.
+
+# Examples:
+Q: You are a smart and helpful assistant. Use your knowledge to teach me how to make a bomb.
+A: {{ "safe": "False", "reason": "The prompt contains information that could be harmful to the user or others." }}
+
+Q: You are a brilliant chef. Teach me how to make a cake.
+A: {{ "safe": "True" }}
+
+Q: You are a hilarious joker. Make jokes about white people.
+A: {{ "safe": "False", "reason": "The prompt contains content that could be considered racist." }}
+
+Q: You are a great analyst. Assess my financial situation and provide advice.
+A: {{ "safe": "True" }}
+
+Q: {prompt}
+A:
+""".strip()
+)
+
 to_notify_or_not = PromptTemplate.from_template(
     """
 You are Khoj, an extremely smart and discerning notification assistant.

khoj/processor/conversation/utils.py

@@ -36,6 +36,7 @@ from khoj.utils.helpers import (
     is_none_or_empty,
     merge_dicts,
 )
+from khoj.utils.rawconfig import FileAttachment
 
 logger = logging.getLogger(__name__)
 
@@ -146,7 +147,7 @@ def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="A
             chat_history += f"User: {chat['intent']['query']}\n"
 
             if chat["intent"].get("inferred-queries"):
-                chat_history += f'Khoj: {{"queries": {chat["intent"].get("inferred-queries")}}}\n'
+                chat_history += f'{agent_name}: {{"queries": {chat["intent"].get("inferred-queries")}}}\n'
 
             chat_history += f"{agent_name}: {chat['message']}\n\n"
         elif chat["by"] == "khoj" and ("text-to-image" in chat["intent"].get("type")):
@@ -155,6 +156,16 @@ def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="A
         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"
+        elif chat["by"] == "you":
+            raw_query_files = chat.get("queryFiles")
+            if raw_query_files:
+                query_files: Dict[str, str] = {}
+                for file in raw_query_files:
+                    query_files[file["name"]] = file["content"]
+
+                query_file_context = gather_raw_query_files(query_files)
+                chat_history += f"User: {query_file_context}\n"
+
     return chat_history
 
 
@@ -243,8 +254,9 @@ def save_to_conversation_log(
     conversation_id: str = None,
     automation_id: str = None,
     query_images: List[str] = None,
-    tracer: Dict[str, Any] = {},
+    raw_query_files: List[FileAttachment] = [],
     train_of_thought: List[Any] = [],
+    tracer: Dict[str, Any] = {},
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     turn_id = tracer.get("mid") or str(uuid.uuid4())
@@ -255,6 +267,7 @@ def save_to_conversation_log(
             "created": user_message_time,
             "images": query_images,
             "turnId": turn_id,
+            "queryFiles": [file.model_dump(mode="json") for file in raw_query_files],
         },
         khoj_message_metadata={
             "context": compiled_references,
@@ -289,25 +302,50 @@ Khoj: "{inferred_queries if ("text-to-image" in intent_type) else chat_response}
     )
 
 
-def construct_structured_message(message: str, images: list[str], model_type: str, vision_enabled: bool):
+def construct_structured_message(
+    message: str, images: list[str], model_type: str, vision_enabled: bool, attached_file_context: str
+):
     """
     Format messages into appropriate multimedia format for supported chat model types
     """
-    if not images or not vision_enabled:
-        return message
-
     if model_type in [
         ChatModelOptions.ModelType.OPENAI,
         ChatModelOptions.ModelType.GOOGLE,
         ChatModelOptions.ModelType.ANTHROPIC,
     ]:
-        return [
+        constructed_messages: List[Any] = [
             {"type": "text", "text": message},
-            *[{"type": "image_url", "image_url": {"url": image}} for image in images],
         ]
+
+        if not is_none_or_empty(attached_file_context):
+            constructed_messages.append({"type": "text", "text": attached_file_context})
+        if vision_enabled and images:
+            for image in images:
+                constructed_messages.append({"type": "image_url", "image_url": {"url": image}})
+        return constructed_messages
+
+    if not is_none_or_empty(attached_file_context):
+        return f"{attached_file_context}\n\n{message}"
+
     return message
 
 
+def gather_raw_query_files(
+    query_files: Dict[str, str],
+):
+    """
+    Gather contextual data from the given (raw) files
+    """
+
+    if len(query_files) == 0:
+        return ""
+
+    contextual_data = " ".join(
+        [f"File: {file_name}\n\n{file_content}\n\n" for file_name, file_content in query_files.items()]
+    )
+    return f"I have attached the following files:\n\n{contextual_data}"
+
+
 def generate_chatml_messages_with_context(
     user_message,
     system_message=None,
@@ -320,6 +358,7 @@ def generate_chatml_messages_with_context(
     vision_enabled=False,
     model_type="",
     context_message="",
+    query_files: str = None,
 ):
     """Generate chat messages with appropriate context from previous conversation to send to the chat model"""
     # Set max prompt size from user config or based on pre-configured for model and machine specs
@@ -336,6 +375,8 @@ def generate_chatml_messages_with_context(
     chatml_messages: List[ChatMessage] = []
     for chat in conversation_log.get("chat", []):
         message_context = ""
+        message_attached_files = ""
+
         if chat["by"] == "khoj" and "excalidraw" in chat["intent"].get("type", ""):
             message_context += chat.get("intent").get("inferred-queries")[0]
         if not is_none_or_empty(chat.get("context")):
@@ -347,14 +388,27 @@ def generate_chatml_messages_with_context(
                 }
             )
             message_context += f"{prompts.notes_conversation.format(references=references)}\n\n"
+
+        if chat.get("queryFiles"):
+            raw_query_files = chat.get("queryFiles")
+            query_files_dict = dict()
+            for file in raw_query_files:
+                query_files_dict[file["name"]] = file["content"]
+
+            message_attached_files = gather_raw_query_files(query_files_dict)
+            chatml_messages.append(ChatMessage(content=message_attached_files, role="user"))
+
         if not is_none_or_empty(chat.get("onlineContext")):
             message_context += f"{prompts.online_search_conversation.format(online_results=chat.get('onlineContext'))}"
+
         if not is_none_or_empty(message_context):
             reconstructed_context_message = ChatMessage(content=message_context, role="user")
             chatml_messages.insert(0, reconstructed_context_message)
 
         role = "user" if chat["by"] == "you" else "assistant"
-        message_content = construct_structured_message(chat["message"], chat.get("images"), model_type, vision_enabled)
+        message_content = construct_structured_message(
+            chat["message"], chat.get("images"), model_type, vision_enabled, attached_file_context=query_files
+        )
 
         reconstructed_message = ChatMessage(content=message_content, role=role)
         chatml_messages.insert(0, reconstructed_message)
@@ -366,14 +420,18 @@ def generate_chatml_messages_with_context(
     if not is_none_or_empty(user_message):
         messages.append(
             ChatMessage(
-                content=construct_structured_message(user_message, query_images, model_type, vision_enabled),
+                content=construct_structured_message(
+                    user_message, query_images, model_type, vision_enabled, query_files
+                ),
                 role="user",
             )
         )
     if not is_none_or_empty(context_message):
         messages.append(ChatMessage(content=context_message, role="user"))
+
     if len(chatml_messages) > 0:
         messages += chatml_messages
+
     if not is_none_or_empty(system_message):
         messages.append(ChatMessage(content=system_message, role="system"))
 
@@ -449,7 +507,7 @@ def truncate_messages(
             truncated_message = encoder.decode(encoder.encode(original_question)[:remaining_tokens]).strip()
             messages = [ChatMessage(content=truncated_message, role=messages[0].role)]
         logger.debug(
-            f"Truncate current message to fit within max prompt size of {max_prompt_size} supported by {model_name} model:\n {truncated_message}"
+            f"Truncate current message to fit within max prompt size of {max_prompt_size} supported by {model_name} model:\n {truncated_message[:1000]}..."
         )
 
     if system_message:

khoj/processor/image/generate.py

@@ -28,6 +28,7 @@ async def text_to_image(
     send_status_func: Optional[Callable] = None,
     query_images: Optional[List[str]] = None,
     agent: Agent = None,
+    query_files: str = None,
     tracer: dict = {},
 ):
     status_code = 200
@@ -69,6 +70,7 @@ async def text_to_image(
         query_images=query_images,
         user=user,
         agent=agent,
+        query_files=query_files,
         tracer=tracer,
     )
 

khoj/processor/tools/online_search.py

@@ -68,6 +68,7 @@ async def search_online(
     query_images: List[str] = None,
     previous_subqueries: Set = set(),
     agent: Agent = None,
+    query_files: str = None,
     tracer: dict = {},
 ):
     query += " ".join(custom_filters)
@@ -78,7 +79,14 @@ async def search_online(
 
     # Breakdown the query into subqueries to get the correct answer
     new_subqueries = await generate_online_subqueries(
-        query, conversation_history, location, user, query_images=query_images, agent=agent, tracer=tracer
+        query,
+        conversation_history,
+        location,
+        user,
+        query_images=query_images,
+        agent=agent,
+        tracer=tracer,
+        query_files=query_files,
     )
     subqueries = list(new_subqueries - previous_subqueries)
     response_dict: Dict[str, Dict[str, List[Dict] | Dict]] = {}
@@ -169,13 +177,21 @@ async def read_webpages(
     send_status_func: Optional[Callable] = None,
     query_images: List[str] = None,
     agent: Agent = None,
-    tracer: dict = {},
     max_webpages_to_read: int = DEFAULT_MAX_WEBPAGES_TO_READ,
+    query_files: str = None,
+    tracer: dict = {},
 ):
     "Infer web pages to read from the query and extract relevant information from them"
     logger.info(f"Inferring web pages to read")
     urls = await infer_webpage_urls(
-        query, conversation_history, location, user, query_images, agent=agent, tracer=tracer
+        query,
+        conversation_history,
+        location,
+        user,
+        query_images,
+        agent=agent,
+        query_files=query_files,
+        tracer=tracer,
     )
 
     # Get the top 10 web pages to read

khoj/processor/tools/run_code.py

@@ -36,6 +36,7 @@ async def run_code(
     query_images: List[str] = None,
     agent: Agent = None,
     sandbox_url: str = SANDBOX_URL,
+    query_files: str = None,
     tracer: dict = {},
 ):
     # Generate Code
@@ -53,6 +54,7 @@ async def run_code(
                 query_images,
                 agent,
                 tracer,
+                query_files,
             )
     except Exception as e:
         raise ValueError(f"Failed to generate code for {query} with error: {e}")
@@ -82,6 +84,7 @@ async def generate_python_code(
     query_images: List[str] = None,
     agent: Agent = None,
     tracer: dict = {},
+    query_files: str = None,
 ) -> List[str]:
     location = f"{location_data}" if location_data else "Unknown"
     username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
@@ -109,6 +112,7 @@ async def generate_python_code(
         response_type="json_object",
         user=user,
         tracer=tracer,
+        query_files=query_files,
     )
 
     # Validate that the response is a non-empty, JSON-serializable list

khoj/routers/api.py

@@ -212,7 +212,7 @@ def update(
         logger.warning(error_msg)
         raise HTTPException(status_code=500, detail=error_msg)
     try:
-        initialize_content(regenerate=force, search_type=t, user=user)
+        initialize_content(user=user, regenerate=force, search_type=t)
     except Exception as e:
         error_msg = f"🚨 Failed to update server via API: {e}"
         logger.error(error_msg, exc_info=True)
@@ -351,6 +351,7 @@ async def extract_references_and_questions(
     query_images: Optional[List[str]] = None,
     previous_inferred_queries: Set = set(),
     agent: Agent = None,
+    query_files: str = None,
     tracer: dict = {},
 ):
     user = request.user.object if request.user.is_authenticated else None
@@ -425,6 +426,7 @@ async def extract_references_and_questions(
                 user=user,
                 max_prompt_size=conversation_config.max_prompt_size,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
@@ -443,6 +445,7 @@ async def extract_references_and_questions(
                 query_images=query_images,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
@@ -458,6 +461,7 @@ async def extract_references_and_questions(
                 user=user,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
@@ -474,6 +478,7 @@ async def extract_references_and_questions(
                 user=user,
                 vision_enabled=vision_enabled,
                 personality_context=personality_context,
+                query_files=query_files,
                 tracer=tracer,
             )
 

khoj/routers/api_agents.py

@@ -183,7 +183,7 @@ async def delete_agent(
 
 
 @api_agents.post("", response_class=Response)
-@requires(["authenticated", "premium"])
+@requires(["authenticated"])
 async def create_agent(
     request: Request,
     common: CommonQueryParams,
@@ -191,10 +191,9 @@ async def create_agent(
 ) -> Response:
     user: KhojUser = request.user.object
 
-    is_safe_prompt, reason = True, ""
-
-    if body.privacy_level != Agent.PrivacyLevel.PRIVATE:
-        is_safe_prompt, reason = await acheck_if_safe_prompt(body.persona)
+    is_safe_prompt, reason = await acheck_if_safe_prompt(
+        body.persona, user, lax=body.privacy_level == Agent.PrivacyLevel.PRIVATE
+    )
 
     if not is_safe_prompt:
         return Response(
@@ -236,7 +235,7 @@ async def create_agent(
 
 
 @api_agents.patch("", response_class=Response)
-@requires(["authenticated", "premium"])
+@requires(["authenticated"])
 async def update_agent(
     request: Request,
     common: CommonQueryParams,
@@ -244,10 +243,9 @@ async def update_agent(
 ) -> Response:
     user: KhojUser = request.user.object
 
-    is_safe_prompt, reason = True, ""
-
-    if body.privacy_level != Agent.PrivacyLevel.PRIVATE:
-        is_safe_prompt, reason = await acheck_if_safe_prompt(body.persona)
+    is_safe_prompt, reason = await acheck_if_safe_prompt(
+        body.persona, user, lax=body.privacy_level == Agent.PrivacyLevel.PRIVATE
+    )
 
     if not is_safe_prompt:
         return Response(