khoj 1.24.2.dev4__py3-none-any.whl → 1.24.2.dev6__py3-none-any.whl

khoj/processor/tools/online_search.py

@@ -10,7 +10,7 @@ import aiohttp
 from bs4 import BeautifulSoup
 from markdownify import markdownify
 
-from khoj.database.models import KhojUser
+from khoj.database.models import Agent, KhojUser
 from khoj.routers.helpers import (
     ChatEvent,
     extract_relevant_info,
@@ -57,16 +57,17 @@ async def search_online(
     send_status_func: Optional[Callable] = None,
     custom_filters: List[str] = [],
     uploaded_image_url: str = None,
+    agent: Agent = None,
 ):
     query += " ".join(custom_filters)
     if not is_internet_connected():
-        logger.warn("Cannot search online as not connected to internet")
+        logger.warning("Cannot search online as not connected to internet")
         yield {}
         return
 
     # Breakdown the query into subqueries to get the correct answer
     subqueries = await generate_online_subqueries(
-        query, conversation_history, location, user, uploaded_image_url=uploaded_image_url
+        query, conversation_history, location, user, uploaded_image_url=uploaded_image_url, agent=agent
     )
     response_dict = {}
 
@@ -101,7 +102,7 @@ async def search_online(
             async for event in send_status_func(f"**Reading web pages**: {webpage_links_str}"):
                 yield {ChatEvent.STATUS: event}
     tasks = [
-        read_webpage_and_extract_content(subquery, link, content, subscribed=subscribed)
+        read_webpage_and_extract_content(subquery, link, content, subscribed=subscribed, agent=agent)
         for link, subquery, content in webpages
     ]
     results = await asyncio.gather(*tasks)
@@ -143,6 +144,7 @@ async def read_webpages(
     subscribed: bool = False,
     send_status_func: Optional[Callable] = None,
     uploaded_image_url: str = None,
+    agent: Agent = None,
 ):
     "Infer web pages to read from the query and extract relevant information from them"
     logger.info(f"Inferring web pages to read")
@@ -156,7 +158,7 @@ async def read_webpages(
         webpage_links_str = "\n- " + "\n- ".join(list(urls))
         async for event in send_status_func(f"**Reading web pages**: {webpage_links_str}"):
             yield {ChatEvent.STATUS: event}
-    tasks = [read_webpage_and_extract_content(query, url, subscribed=subscribed) for url in urls]
+    tasks = [read_webpage_and_extract_content(query, url, subscribed=subscribed, agent=agent) for url in urls]
     results = await asyncio.gather(*tasks)
 
     response: Dict[str, Dict] = defaultdict(dict)
@@ -167,14 +169,14 @@ async def read_webpages(
 
 
 async def read_webpage_and_extract_content(
-    subquery: str, url: str, content: str = None, subscribed: bool = False
+    subquery: str, url: str, content: str = None, subscribed: bool = False, agent: Agent = None
 ) -> Tuple[str, Union[None, str], str]:
     try:
         if is_none_or_empty(content):
             with timer(f"Reading web page at '{url}' took", logger):
                 content = await read_webpage_with_olostep(url) if OLOSTEP_API_KEY else await read_webpage_with_jina(url)
         with timer(f"Extracting relevant information from web page at '{url}' took", logger):
-            extracted_info = await extract_relevant_info(subquery, content, subscribed=subscribed)
+            extracted_info = await extract_relevant_info(subquery, content, subscribed=subscribed, agent=agent)
         return subquery, extracted_info, url
     except Exception as e:
         logger.error(f"Failed to read web page at '{url}' with {e}")

khoj/routers/api.py

@@ -27,7 +27,13 @@ from khoj.database.adapters import (
     get_user_photo,
     get_user_search_model_or_default,
 )
-from khoj.database.models import ChatModelOptions, KhojUser, SpeechToTextModelOptions
+from khoj.database.models import (
+    Agent,
+    ChatModelOptions,
+    KhojUser,
+    SpeechToTextModelOptions,
+)
+from khoj.processor.conversation import prompts
 from khoj.processor.conversation.anthropic.anthropic_chat import (
     extract_questions_anthropic,
 )
@@ -106,6 +112,7 @@ async def execute_search(
     r: Optional[bool] = False,
     max_distance: Optional[Union[float, None]] = None,
     dedupe: Optional[bool] = True,
+    agent: Optional[Agent] = None,
 ):
     start_time = time.time()
 
@@ -157,6 +164,7 @@ async def execute_search(
                     t,
                     question_embedding=encoded_asymmetric_query,
                     max_distance=max_distance,
+                    agent=agent,
                 )
             ]
 
@@ -333,6 +341,7 @@ async def extract_references_and_questions(
     location_data: LocationData = None,
     send_status_func: Optional[Callable] = None,
     uploaded_image_url: Optional[str] = None,
+    agent: Agent = None,
 ):
     user = request.user.object if request.user.is_authenticated else None
 
@@ -348,9 +357,10 @@ async def extract_references_and_questions(
         return
 
     if not await sync_to_async(EntryAdapters.user_has_entries)(user=user):
-        logger.debug("No documents in knowledge base. Use a Khoj client to sync and chat with your docs.")
-        yield compiled_references, inferred_queries, q
-        return
+        if not await sync_to_async(EntryAdapters.agent_has_entries)(agent=agent):
+            logger.debug("No documents in knowledge base. Use a Khoj client to sync and chat with your docs.")
+            yield compiled_references, inferred_queries, q
+            return
 
     # Extract filter terms from user message
     defiltered_query = q
@@ -368,6 +378,8 @@ async def extract_references_and_questions(
     using_offline_chat = False
     logger.debug(f"Filters in query: {filters_in_query}")
 
+    personality_context = prompts.personality_context.format(personality=agent.personality) if agent else ""
+
     # Infer search queries from user message
     with timer("Extracting search queries took", logger):
         # If we've reached here, either the user has enabled offline chat or the openai model is enabled.
@@ -392,6 +404,7 @@ async def extract_references_and_questions(
                 location_data=location_data,
                 user=user,
                 max_prompt_size=conversation_config.max_prompt_size,
+                personality_context=personality_context,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.OPENAI:
             openai_chat_config = conversation_config.openai_config
@@ -408,6 +421,7 @@ async def extract_references_and_questions(
                 user=user,
                 uploaded_image_url=uploaded_image_url,
                 vision_enabled=vision_enabled,
+                personality_context=personality_context,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.ANTHROPIC:
             api_key = conversation_config.openai_config.api_key
@@ -419,6 +433,7 @@ async def extract_references_and_questions(
                 conversation_log=meta_log,
                 location_data=location_data,
                 user=user,
+                personality_context=personality_context,
             )
         elif conversation_config.model_type == ChatModelOptions.ModelType.GOOGLE:
             api_key = conversation_config.openai_config.api_key
@@ -431,6 +446,7 @@ async def extract_references_and_questions(
                 location_data=location_data,
                 max_tokens=conversation_config.max_prompt_size,
                 user=user,
+                personality_context=personality_context,
             )
 
     # Collate search results as context for GPT
@@ -452,6 +468,7 @@ async def extract_references_and_questions(
                     r=True,
                     max_distance=d,
                     dedupe=False,
+                    agent=agent,
                 )
             )
         search_results = text_search.deduplicated_search_responses(search_results)

khoj/routers/api_agents.py

@@ -1,13 +1,22 @@
 import json
 import logging
+from typing import Dict, List, Optional
 
+from asgiref.sync import sync_to_async
 from fastapi import APIRouter, Request
 from fastapi.requests import Request
 from fastapi.responses import Response
+from pydantic import BaseModel
+from starlette.authentication import requires
 
 from khoj.database.adapters import AgentAdapters
-from khoj.database.models import KhojUser
-from khoj.routers.helpers import CommonQueryParams
+from khoj.database.models import Agent, KhojUser
+from khoj.routers.helpers import CommonQueryParams, acheck_if_safe_prompt
+from khoj.utils.helpers import (
+    ConversationCommand,
+    command_descriptions_for_agent,
+    mode_descriptions_for_agent,
+)
 
 # Initialize Router
 logger = logging.getLogger(__name__)
@@ -16,6 +25,18 @@ logger = logging.getLogger(__name__)
 api_agents = APIRouter()
 
 
+class ModifyAgentBody(BaseModel):
+    name: str
+    persona: str
+    privacy_level: str
+    icon: str
+    color: str
+    chat_model: str
+    files: Optional[List[str]] = []
+    input_tools: Optional[List[str]] = []
+    output_modes: Optional[List[str]] = []
+
+
 @api_agents.get("", response_class=Response)
 async def all_agents(
     request: Request,
@@ -25,17 +46,22 @@ async def all_agents(
     agents = await AgentAdapters.aget_all_accessible_agents(user)
     agents_packet = list()
     for agent in agents:
+        files = agent.fileobject_set.all()
+        file_names = [file.file_name for file in files]
         agents_packet.append(
             {
                 "slug": agent.slug,
-                "avatar": agent.avatar,
                 "name": agent.name,
                 "persona": agent.personality,
-                "public": agent.public,
                 "creator": agent.creator.username if agent.creator else None,
                 "managed_by_admin": agent.managed_by_admin,
                 "color": agent.style_color,
                 "icon": agent.style_icon,
+                "privacy_level": agent.privacy_level,
+                "chat_model": agent.chat_model.chat_model,
+                "files": file_names,
+                "input_tools": agent.input_tools,
+                "output_modes": agent.output_modes,
             }
         )
 
@@ -43,3 +69,197 @@ async def all_agents(
     agents_packet.sort(key=lambda x: x["name"])
     agents_packet.sort(key=lambda x: x["slug"] == "khoj", reverse=True)
     return Response(content=json.dumps(agents_packet), media_type="application/json", status_code=200)
+
+
+@api_agents.get("/options", response_class=Response)
+async def get_agent_configuration_options(
+    request: Request,
+    common: CommonQueryParams,
+) -> Response:
+    agent_input_tools = [key for key, _ in Agent.InputToolOptions.choices]
+    agent_output_modes = [key for key, _ in Agent.OutputModeOptions.choices]
+
+    agent_input_tool_with_descriptions: Dict[str, str] = {}
+    for key in agent_input_tools:
+        conversation_command = ConversationCommand(key)
+        agent_input_tool_with_descriptions[key] = command_descriptions_for_agent[conversation_command]
+
+    agent_output_modes_with_descriptions: Dict[str, str] = {}
+    for key in agent_output_modes:
+        conversation_command = ConversationCommand(key)
+        agent_output_modes_with_descriptions[key] = mode_descriptions_for_agent[conversation_command]
+
+    return Response(
+        content=json.dumps(
+            {
+                "input_tools": agent_input_tool_with_descriptions,
+                "output_modes": agent_output_modes_with_descriptions,
+            }
+        ),
+        media_type="application/json",
+        status_code=200,
+    )
+
+
+@api_agents.get("/{agent_slug}", response_class=Response)
+async def get_agent(
+    request: Request,
+    common: CommonQueryParams,
+    agent_slug: str,
+) -> Response:
+    user: KhojUser = request.user.object if request.user.is_authenticated else None
+    agent = await AgentAdapters.aget_readonly_agent_by_slug(agent_slug, user)
+
+    if not agent:
+        return Response(
+            content=json.dumps({"error": f"Agent with name {agent_slug} not found."}),
+            media_type="application/json",
+            status_code=404,
+        )
+
+    files = agent.fileobject_set.all()
+    file_names = [file.file_name for file in files]
+    agents_packet = {
+        "slug": agent.slug,
+        "name": agent.name,
+        "persona": agent.personality,
+        "creator": agent.creator.username if agent.creator else None,
+        "managed_by_admin": agent.managed_by_admin,
+        "color": agent.style_color,
+        "icon": agent.style_icon,
+        "privacy_level": agent.privacy_level,
+        "chat_model": agent.chat_model.chat_model,
+        "files": file_names,
+        "input_tools": agent.input_tools,
+        "output_modes": agent.output_modes,
+    }
+
+    return Response(content=json.dumps(agents_packet), media_type="application/json", status_code=200)
+
+
+@api_agents.delete("/{agent_slug}", response_class=Response)
+@requires(["authenticated"])
+async def delete_agent(
+    request: Request,
+    common: CommonQueryParams,
+    agent_slug: str,
+) -> Response:
+    user: KhojUser = request.user.object
+
+    agent = await AgentAdapters.aget_agent_by_slug(agent_slug, user)
+
+    if not agent:
+        return Response(
+            content=json.dumps({"error": f"Agent with name {agent_slug} not found."}),
+            media_type="application/json",
+            status_code=404,
+        )
+
+    await AgentAdapters.adelete_agent_by_slug(agent_slug, user)
+
+    return Response(content=json.dumps({"message": "Agent deleted."}), media_type="application/json", status_code=200)
+
+
+@api_agents.post("", response_class=Response)
+@requires(["authenticated"])
+async def create_agent(
+    request: Request,
+    common: CommonQueryParams,
+    body: ModifyAgentBody,
+) -> Response:
+    user: KhojUser = request.user.object
+
+    is_safe_prompt, reason = await acheck_if_safe_prompt(body.persona)
+    if not is_safe_prompt:
+        return Response(
+            content=json.dumps({"error": f"{reason}"}),
+            media_type="application/json",
+            status_code=400,
+        )
+
+    agent = await AgentAdapters.aupdate_agent(
+        user,
+        body.name,
+        body.persona,
+        body.privacy_level,
+        body.icon,
+        body.color,
+        body.chat_model,
+        body.files,
+        body.input_tools,
+        body.output_modes,
+    )
+
+    agents_packet = {
+        "slug": agent.slug,
+        "name": agent.name,
+        "persona": agent.personality,
+        "creator": agent.creator.username if agent.creator else None,
+        "managed_by_admin": agent.managed_by_admin,
+        "color": agent.style_color,
+        "icon": agent.style_icon,
+        "privacy_level": agent.privacy_level,
+        "chat_model": agent.chat_model.chat_model,
+        "files": body.files,
+        "input_tools": agent.input_tools,
+        "output_modes": agent.output_modes,
+    }
+
+    return Response(content=json.dumps(agents_packet), media_type="application/json", status_code=200)
+
+
+@api_agents.patch("", response_class=Response)
+@requires(["authenticated"])
+async def update_agent(
+    request: Request,
+    common: CommonQueryParams,
+    body: ModifyAgentBody,
+) -> Response:
+    user: KhojUser = request.user.object
+
+    is_safe_prompt, reason = await acheck_if_safe_prompt(body.persona)
+    if not is_safe_prompt:
+        return Response(
+            content=json.dumps({"error": f"{reason}"}),
+            media_type="application/json",
+            status_code=400,
+        )
+
+    selected_agent = await AgentAdapters.aget_agent_by_name(body.name, user)
+
+    if not selected_agent:
+        return Response(
+            content=json.dumps({"error": f"Agent with name {body.name} not found."}),
+            media_type="application/json",
+            status_code=404,
+        )
+
+    agent = await AgentAdapters.aupdate_agent(
+        user,
+        body.name,
+        body.persona,
+        body.privacy_level,
+        body.icon,
+        body.color,
+        body.chat_model,
+        body.files,
+        body.input_tools,
+        body.output_modes,
+    )
+
+    agents_packet = {
+        "slug": agent.slug,
+        "name": agent.name,
+        "persona": agent.personality,
+        "creator": agent.creator.username if agent.creator else None,
+        "managed_by_admin": agent.managed_by_admin,
+        "color": agent.style_color,
+        "icon": agent.style_icon,
+        "privacy_level": agent.privacy_level,
+        "chat_model": agent.chat_model.chat_model,
+        "files": body.files,
+        "input_tools": agent.input_tools,
+        "output_modes": agent.output_modes,
+    }
+
+    return Response(content=json.dumps(agents_packet), media_type="application/json", status_code=200)

khoj/routers/api_chat.py

@@ -17,13 +17,14 @@ from starlette.authentication import has_required_scope, requires
 
 from khoj.app.settings import ALLOWED_HOSTS
 from khoj.database.adapters import (
+    AgentAdapters,
     ConversationAdapters,
     EntryAdapters,
     FileObjectAdapters,
     PublicConversationAdapters,
     aget_user_name,
 )
-from khoj.database.models import KhojUser
+from khoj.database.models import Agent, KhojUser
 from khoj.processor.conversation.prompts import help_message, no_entries_found
 from khoj.processor.conversation.utils import save_to_conversation_log
 from khoj.processor.image.generate import text_to_image
@@ -211,7 +212,6 @@ def chat_history(
         agent_metadata = {
             "slug": conversation.agent.slug,
             "name": conversation.agent.name,
-            "avatar": conversation.agent.avatar,
             "isCreator": conversation.agent.creator == user,
             "color": conversation.agent.style_color,
             "icon": conversation.agent.style_icon,
@@ -268,7 +268,6 @@ def get_shared_chat(
         agent_metadata = {
             "slug": conversation.agent.slug,
             "name": conversation.agent.name,
-            "avatar": conversation.agent.avatar,
             "isCreator": conversation.agent.creator == user,
             "color": conversation.agent.style_color,
             "icon": conversation.agent.style_icon,
@@ -418,7 +417,7 @@ def chat_sessions(
         conversations = conversations[:8]
 
     sessions = conversations.values_list(
-        "id", "slug", "title", "agent__slug", "agent__name", "agent__avatar", "created_at", "updated_at"
+        "id", "slug", "title", "agent__slug", "agent__name", "created_at", "updated_at"
     )
 
     session_values = [
@@ -426,9 +425,8 @@ def chat_sessions(
             "conversation_id": str(session[0]),
             "slug": session[2] or session[1],
             "agent_name": session[4],
-            "agent_avatar": session[5],
-            "created": session[6].strftime("%Y-%m-%d %H:%M:%S"),
-            "updated": session[7].strftime("%Y-%m-%d %H:%M:%S"),
+            "created": session[5].strftime("%Y-%m-%d %H:%M:%S"),
+            "updated": session[6].strftime("%Y-%m-%d %H:%M:%S"),
         }
         for session in sessions
     ]
@@ -590,7 +588,7 @@ async def chat(
             nonlocal connection_alive, ttft
             if not connection_alive or await request.is_disconnected():
                 connection_alive = False
-                logger.warn(f"User {user} disconnected from {common.client} client")
+                logger.warning(f"User {user} disconnected from {common.client} client")
                 return
             try:
                 if event_type == ChatEvent.END_LLM_RESPONSE:
@@ -658,6 +656,11 @@ async def chat(
             return
         conversation_id = conversation.id
 
+        agent: Agent | None = None
+        default_agent = await AgentAdapters.aget_default_agent()
+        if conversation.agent and conversation.agent != default_agent:
+            agent = conversation.agent
+
         await is_ready_to_chat(user)
 
         user_name = await aget_user_name(user)
@@ -677,7 +680,12 @@ async def chat(
 
         if conversation_commands == [ConversationCommand.Default] or is_automated_task:
             conversation_commands = await aget_relevant_information_sources(
-                q, meta_log, is_automated_task, subscribed=subscribed, uploaded_image_url=uploaded_image_url
+                q,
+                meta_log,
+                is_automated_task,
+                subscribed=subscribed,
+                uploaded_image_url=uploaded_image_url,
+                agent=agent,
             )
             conversation_commands_str = ", ".join([cmd.value for cmd in conversation_commands])
             async for result in send_event(
@@ -685,7 +693,7 @@ async def chat(
             ):
                 yield result
 
-            mode = await aget_relevant_output_modes(q, meta_log, is_automated_task, uploaded_image_url)
+            mode = await aget_relevant_output_modes(q, meta_log, is_automated_task, uploaded_image_url, agent)
             async for result in send_event(ChatEvent.STATUS, f"**Decided Response Mode:** {mode.value}"):
                 yield result
             if mode not in conversation_commands:
@@ -734,7 +742,7 @@ async def chat(
                         yield result
 
                     response = await extract_relevant_summary(
-                        q, contextual_data, subscribed=subscribed, uploaded_image_url=uploaded_image_url
+                        q, contextual_data, subscribed=subscribed, uploaded_image_url=uploaded_image_url, agent=agent
                     )
                     response_log = str(response)
                     async for result in send_llm_response(response_log):
@@ -816,6 +824,7 @@ async def chat(
             location,
             partial(send_event, ChatEvent.STATUS),
             uploaded_image_url=uploaded_image_url,
+            agent=agent,
         ):
             if isinstance(result, dict) and ChatEvent.STATUS in result:
                 yield result[ChatEvent.STATUS]
@@ -853,6 +862,7 @@ async def chat(
                     partial(send_event, ChatEvent.STATUS),
                     custom_filters,
                     uploaded_image_url=uploaded_image_url,
+                    agent=agent,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
@@ -876,6 +886,7 @@ async def chat(
                     subscribed,
                     partial(send_event, ChatEvent.STATUS),
                     uploaded_image_url=uploaded_image_url,
+                    agent=agent,
                 ):
                     if isinstance(result, dict) and ChatEvent.STATUS in result:
                         yield result[ChatEvent.STATUS]
@@ -922,6 +933,7 @@ async def chat(
                 subscribed=subscribed,
                 send_status_func=partial(send_event, ChatEvent.STATUS),
                 uploaded_image_url=uploaded_image_url,
+                agent=agent,
             ):
                 if isinstance(result, dict) and ChatEvent.STATUS in result:
                     yield result[ChatEvent.STATUS]
@@ -1132,6 +1144,7 @@ async def get_chat(
                 yield result
             return
         conversation_id = conversation.id
+        agent = conversation.agent if conversation.agent else None
 
         await is_ready_to_chat(user)