khoj 1.42.1.dev10__py3-none-any.whl → 1.42.2.dev16__py3-none-any.whl

khoj/routers/api_automation.py

@@ -0,0 +1,243 @@
+import json
+import logging
+import threading
+from typing import Optional
+
+import cron_descriptor
+import pytz
+from apscheduler.job import Job
+from apscheduler.triggers.cron import CronTrigger
+from fastapi import APIRouter, Request
+from fastapi.responses import Response
+from starlette.authentication import requires
+
+from khoj.database.adapters import AutomationAdapters, ConversationAdapters
+from khoj.database.models import KhojUser
+from khoj.processor.conversation.utils import clean_json
+from khoj.routers.helpers import schedule_automation, schedule_query
+from khoj.utils.helpers import is_none_or_empty
+
+# Initialize Router
+api_automation = APIRouter()
+logger = logging.getLogger(__name__)
+
+
+@api_automation.get("", response_class=Response)
+@requires(["authenticated"])
+def get_automations(request: Request) -> Response:
+    user: KhojUser = request.user.object
+
+    # Collate all automations created by user that are still active
+    automations_info = [automation_info for automation_info in AutomationAdapters.get_automations_metadata(user)]
+
+    # Return tasks information as a JSON response
+    return Response(content=json.dumps(automations_info), media_type="application/json", status_code=200)
+
+
+@api_automation.delete("", response_class=Response)
+@requires(["authenticated"])
+def delete_automation(request: Request, automation_id: str) -> Response:
+    user: KhojUser = request.user.object
+
+    try:
+        automation_info = AutomationAdapters.delete_automation(user, automation_id)
+    except ValueError:
+        return Response(status_code=204)
+
+    # Return deleted automation information as a JSON response
+    return Response(content=json.dumps(automation_info), media_type="application/json", status_code=200)
+
+
+@api_automation.post("", response_class=Response)
+@requires(["authenticated"])
+def post_automation(
+    request: Request,
+    q: str,
+    crontime: str,
+    subject: Optional[str] = None,
+    city: Optional[str] = None,
+    region: Optional[str] = None,
+    country: Optional[str] = None,
+    timezone: Optional[str] = None,
+) -> Response:
+    user: KhojUser = request.user.object
+
+    # Perform validation checks
+    if is_none_or_empty(q) or is_none_or_empty(crontime):
+        return Response(content="A query and crontime is required", status_code=400)
+    if not cron_descriptor.get_description(crontime):
+        return Response(content="Invalid crontime", status_code=400)
+
+    # Infer subject, query to run
+    _, query_to_run, generated_subject = schedule_query(q, chat_history=[], user=user)
+    subject = subject or generated_subject
+
+    # Normalize query parameters
+    # Add /automated_task prefix to query if not present
+    query_to_run = query_to_run.strip()
+    if not query_to_run.startswith("/automated_task"):
+        query_to_run = f"/automated_task {query_to_run}"
+
+    # Normalize crontime for AP Scheduler CronTrigger
+    crontime = crontime.strip()
+    if len(crontime.split(" ")) > 5:
+        # Truncate crontime to 5 fields
+        crontime = " ".join(crontime.split(" ")[:5])
+
+    # Convert crontime to standard unix crontime
+    crontime = crontime.replace("?", "*")
+
+    # Disallow minute level automation recurrence
+    minute_value = crontime.split(" ")[0]
+    if not minute_value.isdigit():
+        return Response(
+            content="Minute level recurrence is unsupported. Please create a less frequent schedule.",
+            status_code=400,
+        )
+
+    # Create new Conversation Session associated with this new task
+    title = f"Automation: {subject}"
+    conversation = ConversationAdapters.create_conversation_session(user, request.user.client_app, title=title)
+
+    # Schedule automation with query_to_run, timezone, subject directly provided by user
+    try:
+        # Use the query to run as the scheduling request if the scheduling request is unset
+        calling_url = request.url.replace(query=f"{request.url.query}")
+        automation = schedule_automation(
+            query_to_run, subject, crontime, timezone, q, user, calling_url, str(conversation.id)
+        )
+    except Exception as e:
+        logger.error(f"Error creating automation {q} for {user.email}: {e}", exc_info=True)
+        return Response(
+            content=f"Unable to create automation. Ensure the automation doesn't already exist.",
+            media_type="text/plain",
+            status_code=500,
+        )
+
+    # Collate info about the created user automation
+    automation_info = AutomationAdapters.get_automation_metadata(user, automation)
+
+    # Return information about the created automation as a JSON response
+    return Response(content=json.dumps(automation_info), media_type="application/json", status_code=200)
+
+
+@api_automation.post("/trigger", response_class=Response)
+@requires(["authenticated"])
+def trigger_manual_job(
+    request: Request,
+    automation_id: str,
+):
+    user: KhojUser = request.user.object
+
+    # Check, get automation to edit
+    try:
+        automation: Job = AutomationAdapters.get_automation(user, automation_id)
+    except ValueError as e:
+        logger.error(f"Error triggering automation {automation_id} for {user.email}: {e}", exc_info=True)
+        return Response(content="Invalid automation", status_code=403)
+
+    # Trigger the job without waiting for the result.
+    scheduled_chat_func = automation.func
+
+    # Run the function in a separate thread
+    thread = threading.Thread(target=scheduled_chat_func, args=automation.args, kwargs=automation.kwargs)
+    thread.start()
+
+    return Response(content="Automation triggered", status_code=200)
+
+
+@api_automation.put("", response_class=Response)
+@requires(["authenticated"])
+def edit_job(
+    request: Request,
+    automation_id: str,
+    q: Optional[str],
+    subject: Optional[str],
+    crontime: Optional[str],
+    city: Optional[str] = None,
+    region: Optional[str] = None,
+    country: Optional[str] = None,
+    timezone: Optional[str] = None,
+) -> Response:
+    user: KhojUser = request.user.object
+
+    # Perform validation checks
+    if is_none_or_empty(q) or is_none_or_empty(subject) or is_none_or_empty(crontime):
+        return Response(content="A query, subject and crontime is required", status_code=400)
+    if not cron_descriptor.get_description(crontime):
+        return Response(content="Invalid crontime", status_code=400)
+
+    # Check, get automation to edit
+    try:
+        automation: Job = AutomationAdapters.get_automation(user, automation_id)
+    except ValueError as e:
+        logger.error(f"Error editing automation {automation_id} for {user.email}: {e}", exc_info=True)
+        return Response(content="Invalid automation", status_code=403)
+
+    # Infer subject, query to run
+    _, query_to_run, _ = schedule_query(q, chat_history=[], user=user)
+    subject = subject
+
+    # Normalize query parameters
+    # Add /automated_task prefix to query if not present
+    query_to_run = query_to_run.strip()
+    if not query_to_run.startswith("/automated_task"):
+        query_to_run = f"/automated_task {query_to_run}"
+    # Normalize crontime for AP Scheduler CronTrigger
+    crontime = crontime.strip()
+    if len(crontime.split(" ")) > 5:
+        # Truncate crontime to 5 fields
+        crontime = " ".join(crontime.split(" ")[:5])
+    # Convert crontime to standard unix crontime
+    crontime = crontime.replace("?", "*")
+
+    # Disallow minute level automation recurrence
+    minute_value = crontime.split(" ")[0]
+    if not minute_value.isdigit():
+        return Response(
+            content="Recurrence of every X minutes is unsupported. Please create a less frequent schedule.",
+            status_code=400,
+        )
+
+    # Construct updated automation metadata
+    automation_metadata: dict[str, str] = json.loads(clean_json(automation.name))
+    automation_metadata["scheduling_request"] = q
+    automation_metadata["query_to_run"] = query_to_run
+    automation_metadata["subject"] = subject.strip()
+    automation_metadata["crontime"] = crontime
+    conversation_id = automation_metadata.get("conversation_id")
+
+    if not conversation_id:
+        title = f"Automation: {subject}"
+
+        # Create new Conversation Session associated with this new task
+        conversation = ConversationAdapters.create_conversation_session(user, request.user.client_app, title=title)
+
+        conversation_id = str(conversation.id)
+        automation_metadata["conversation_id"] = conversation_id
+
+    # Modify automation with updated query, subject
+    automation.modify(
+        name=json.dumps(automation_metadata),
+        kwargs={
+            "query_to_run": query_to_run,
+            "subject": subject,
+            "scheduling_request": q,
+            "user": user,
+            "calling_url": request.url,
+            "conversation_id": conversation_id,
+        },
+    )
+
+    # Reschedule automation if crontime updated
+    user_timezone = pytz.timezone(timezone)
+    trigger = CronTrigger.from_crontab(crontime, user_timezone)
+    if automation.trigger != trigger:
+        automation.reschedule(trigger=trigger)
+
+    # Collate info about the updated user automation
+    automation = AutomationAdapters.get_automation(user, automation.id)
+    automation_info = AutomationAdapters.get_automation_metadata(user, automation)
+
+    # Return modified automation information as a JSON response
+    return Response(content=json.dumps(automation_info), media_type="application/json", status_code=200)

khoj/routers/api_chat.py

@@ -40,7 +40,6 @@ from khoj.processor.tools.online_search import (
     search_online,
 )
 from khoj.processor.tools.run_code import run_code
-from khoj.routers.api import extract_references_and_questions
 from khoj.routers.email import send_query_feedback
 from khoj.routers.helpers import (
     ApiImageRateLimiter,
@@ -63,6 +62,7 @@ from khoj.routers.helpers import (
     is_query_empty,
     is_ready_to_chat,
     read_chat_stream,
+    search_documents,
     update_telemetry_state,
     validate_chat_model,
 )
@@ -752,7 +752,7 @@ async def chat(
                                 q,
                                 chat_response="",
                                 user=user,
-                                meta_log=meta_log,
+                                chat_history=chat_history,
                                 compiled_references=compiled_references,
                                 online_results=online_results,
                                 code_results=code_results,
@@ -918,7 +918,7 @@ async def chat(
         if city or region or country or country_code:
             location = LocationData(city=city, region=region, country=country, country_code=country_code)
         user_message_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        meta_log = conversation.conversation_log
+        chat_history = conversation.messages
 
         # If interrupt flag is set, wait for the previous turn to be saved before proceeding
         if interrupt_flag:
@@ -964,14 +964,14 @@ async def chat(
             operator_results = [OperatorRun(**iter_dict) for iter_dict in last_message.operatorContext or []]
             train_of_thought = [thought.model_dump() for thought in last_message.trainOfThought or []]
             # Drop the interrupted message from conversation history
-            meta_log["chat"].pop()
+            chat_history.pop()
             logger.info(f"Loaded interrupted partial context from conversation {conversation_id}.")
 
         if conversation_commands == [ConversationCommand.Default]:
             try:
                 chosen_io = await aget_data_sources_and_output_format(
                     q,
-                    meta_log,
+                    chat_history,
                     is_automated_task,
                     user=user,
                     query_images=uploaded_images,
@@ -1011,7 +1011,7 @@ async def chat(
                 user=user,
                 query=defiltered_query,
                 conversation_id=conversation_id,
-                conversation_history=meta_log,
+                conversation_history=conversation.messages,
                 previous_iterations=list(research_results),
                 query_images=uploaded_images,
                 agent=agent,
@@ -1055,115 +1055,13 @@ async def chat(
             if state.verbose > 1:
                 logger.debug(f'Researched Results: {"".join(r.summarizedResult for r in research_results)}')
 
-        used_slash_summarize = conversation_commands == [ConversationCommand.Summarize]
-        # Skip trying to summarize if
-        if (
-            # summarization intent was inferred
-            ConversationCommand.Summarize in conversation_commands
-            # and not triggered via slash command
-            and not used_slash_summarize
-            # but we can't actually summarize
-            and len(file_filters) == 0
-        ):
-            conversation_commands.remove(ConversationCommand.Summarize)
-        elif ConversationCommand.Summarize in conversation_commands:
-            response_log = ""
-            agent_has_entries = await EntryAdapters.aagent_has_entries(agent)
-            if len(file_filters) == 0 and not agent_has_entries:
-                response_log = "No files selected for summarization. Please add files using the section on the left."
-                async for result in send_llm_response(response_log, tracer.get("usage")):
-                    yield result
-            else:
-                async for response in generate_summary_from_files(
-                    q=q,
-                    user=user,
-                    file_filters=file_filters,
-                    meta_log=meta_log,
-                    query_images=uploaded_images,
-                    agent=agent,
-                    send_status_func=partial(send_event, ChatEvent.STATUS),
-                    query_files=attached_file_context,
-                    tracer=tracer,
-                ):
-                    if isinstance(response, dict) and ChatEvent.STATUS in response:
-                        yield response[ChatEvent.STATUS]
-                    else:
-                        if isinstance(response, str):
-                            response_log = response
-                            async for result in send_llm_response(response, tracer.get("usage")):
-                                yield result
-
-            summarized_document = FileAttachment(
-                name="Summarized Document",
-                content=response_log,
-                type="text/plain",
-                size=len(response_log.encode("utf-8")),
-            )
-
-            async for result in send_event(ChatEvent.GENERATED_ASSETS, {"files": [summarized_document.model_dump()]}):
-                yield result
-
-            generated_files.append(summarized_document)
-
-        custom_filters = []
-        if conversation_commands == [ConversationCommand.Help]:
-            if not q:
-                chat_model = await ConversationAdapters.aget_user_chat_model(user)
-                if chat_model == None:
-                    chat_model = await ConversationAdapters.aget_default_chat_model(user)
-                model_type = chat_model.model_type
-                formatted_help = help_message.format(model=model_type, version=state.khoj_version, device=get_device())
-                async for result in send_llm_response(formatted_help, tracer.get("usage")):
-                    yield result
-                return
-            # Adding specification to search online specifically on khoj.dev pages.
-            custom_filters.append("site:khoj.dev")
-            conversation_commands.append(ConversationCommand.Online)
-
-        if ConversationCommand.Automation in conversation_commands:
-            try:
-                automation, crontime, query_to_run, subject = await create_automation(
-                    q, timezone, user, request.url, meta_log, tracer=tracer
-                )
-            except Exception as e:
-                logger.error(f"Error scheduling task {q} for {user.email}: {e}")
-                error_message = f"Unable to create automation. Ensure the automation doesn't already exist."
-                async for result in send_llm_response(error_message, tracer.get("usage")):
-                    yield result
-                return
-
-            llm_response = construct_automation_created_message(automation, crontime, query_to_run, subject)
-            # Trigger task to save conversation to DB
-            asyncio.create_task(
-                save_to_conversation_log(
-                    q,
-                    llm_response,
-                    user,
-                    meta_log,
-                    user_message_time,
-                    intent_type="automation",
-                    client_application=request.user.client_app,
-                    conversation_id=conversation_id,
-                    inferred_queries=[query_to_run],
-                    automation_id=automation.id,
-                    query_images=uploaded_images,
-                    train_of_thought=train_of_thought,
-                    raw_query_files=raw_query_files,
-                    tracer=tracer,
-                )
-            )
-            # Send LLM Response
-            async for result in send_llm_response(llm_response, tracer.get("usage")):
-                yield result
-            return
-
         # Gather Context
         ## Extract Document References
         if not ConversationCommand.Research in conversation_commands:
             try:
-                async for result in extract_references_and_questions(
+                async for result in search_documents(
                     user,
-                    meta_log,
+                    chat_history,
                     q,
                     (n or 7),
                     d,
@@ -1212,11 +1110,11 @@ async def chat(
             try:
                 async for result in search_online(
                     defiltered_query,
-                    meta_log,
+                    chat_history,
                     location,
                     user,
                     partial(send_event, ChatEvent.STATUS),
-                    custom_filters,
+                    custom_filters=[],
                     max_online_searches=3,
                     query_images=uploaded_images,
                     query_files=attached_file_context,
@@ -1240,7 +1138,7 @@ async def chat(
             try:
                 async for result in read_webpages(
                     defiltered_query,
-                    meta_log,
+                    chat_history,
                     location,
                     user,
                     partial(send_event, ChatEvent.STATUS),
@@ -1281,7 +1179,7 @@ async def chat(
                 context = f"# Iteration 1:\n#---\nNotes:\n{compiled_references}\n\nOnline Results:{online_results}"
                 async for result in run_code(
                     defiltered_query,
-                    meta_log,
+                    chat_history,
                     context,
                     location,
                     user,
@@ -1306,7 +1204,7 @@ async def chat(
                 async for result in operate_environment(
                     defiltered_query,
                     user,
-                    meta_log,
+                    chat_history,
                     location,
                     list(operator_results)[-1] if operator_results else None,
                     query_images=uploaded_images,
@@ -1356,7 +1254,7 @@ async def chat(
             async for result in text_to_image(
                 defiltered_query,
                 user,
-                meta_log,
+                chat_history,
                 location_data=location,
                 references=compiled_references,
                 online_results=online_results,
@@ -1400,7 +1298,7 @@ async def chat(
 
             async for result in generate_mermaidjs_diagram(
                 q=defiltered_query,
-                conversation_history=meta_log,
+                chat_history=chat_history,
                 location_data=location,
                 note_references=compiled_references,
                 online_results=online_results,
@@ -1456,40 +1354,36 @@ async def chat(
 
         llm_response, chat_metadata = await agenerate_chat_response(
             defiltered_query,
-            meta_log,
+            chat_history,
             conversation,
             compiled_references,
             online_results,
             code_results,
             operator_results,
             research_results,
-            inferred_queries,
-            conversation_commands,
             user,
-            request.user.client_app,
             location,
             user_name,
             uploaded_images,
-            train_of_thought,
             attached_file_context,
-            raw_query_files,
-            generated_images,
             generated_files,
-            generated_mermaidjs_diagram,
             program_execution_context,
             generated_asset_results,
             is_subscribed,
             tracer,
         )
 
+        full_response = ""
         async for item in llm_response:
-            # Should not happen with async generator, end is signaled by loop exit. Skip.
-            if item is None:
+            # Should not happen with async generator. Skip.
+            if item is None or not isinstance(item, ResponseWithThought):
+                logger.warning(f"Unexpected item type in LLM response: {type(item)}. Skipping.")
                 continue
             if cancellation_event.is_set():
                 break
-            message = item.response if isinstance(item, ResponseWithThought) else item
-            if isinstance(item, ResponseWithThought) and item.thought:
+            message = item.response
+            full_response += message if message else ""
+            if item.thought:
                 async for result in send_event(ChatEvent.THOUGHT, item.thought):
                     yield result
                 continue
@@ -1506,6 +1400,31 @@ async def chat(
                     logger.warning(f"Error during streaming. Stopping send: {e}")
                 break
 
+        # Save conversation once finish streaming
+        asyncio.create_task(
+            save_to_conversation_log(
+                q,
+                chat_response=full_response,
+                user=user,
+                chat_history=chat_history,
+                compiled_references=compiled_references,
+                online_results=online_results,
+                code_results=code_results,
+                operator_results=operator_results,
+                research_results=research_results,
+                inferred_queries=inferred_queries,
+                client_application=request.user.client_app,
+                conversation_id=str(conversation.id),
+                query_images=uploaded_images,
+                train_of_thought=train_of_thought,
+                raw_query_files=raw_query_files,
+                generated_images=generated_images,
+                raw_generated_files=generated_files,
+                generated_mermaidjs_diagram=generated_mermaidjs_diagram,
+                tracer=tracer,
+            )
+        )
+
         # Signal end of LLM response after the loop finishes
         if not cancellation_event.is_set():
             async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""):