diff --git a/src/interface/web/app/automations/page.tsx b/src/interface/web/app/automations/page.tsx index e4c9abc94..201a173d5 100644 --- a/src/interface/web/app/automations/page.tsx +++ b/src/interface/web/app/automations/page.tsx @@ -21,6 +21,7 @@ interface AutomationsData { schedule: string; crontime: string; next: string; + chat_model_id?: number | null; } import cronstrue from "cronstrue"; @@ -49,6 +50,7 @@ import { useSearchParams } from "next/navigation"; import Link from "next/link"; import { Popover, PopoverTrigger, PopoverContent } from "@/components/ui/popover"; import { + Brain, CalendarCheck, CalendarDot, CalendarDots, @@ -62,7 +64,7 @@ import { Plus, Trash, } from "@phosphor-icons/react"; -import { useAuthenticatedData, UserProfile } from "../common/auth"; +import { ModelOptions, useAuthenticatedData, useChatModelOptions, UserProfile } from "../common/auth"; import LoginPrompt from "../components/loginPrompt/loginPrompt"; import { useToast } from "@/components/ui/use-toast"; import { ToastAction } from "@/components/ui/toast"; @@ -271,6 +273,7 @@ interface AutomationsCardProps { setShowLoginPrompt: (showLoginPrompt: boolean) => void; authenticatedData: UserProfile | null; setToastMessage: (toastMessage: string) => void; + chatModelOptions?: ModelOptions[]; } function AutomationsCard(props: AutomationsCardProps) { @@ -335,6 +338,7 @@ function AutomationsCard(props: AutomationsCardProps) { automation={updatedAutomationData || automation} ipLocationData={props.locationData} setToastMessage={props.setToastMessage} + chatModelOptions={props.chatModelOptions} /> )} -
+
@@ -399,6 +403,21 @@ function AutomationsCard(props: AutomationsCardProps) { {intervalString}
+ {(() => { + const modelId = + updatedAutomationData?.chat_model_id ?? automation.chat_model_id; + const modelName = props.chatModelOptions?.find( + (m) => m.id === modelId, + )?.name; + return modelName ? ( +
+ +
+ {modelName} +
+
+ ) : null; + })()}
{props.suggestedCard && props.setNewAutomationData && ( )} @@ -428,6 +448,7 @@ interface SharedAutomationCardProps { authenticatedData: UserProfile | null; isMobileWidth: boolean; setToastMessage: (toastMessage: string) => void; + chatModelOptions?: ModelOptions[]; } function SharedAutomationCard(props: SharedAutomationCardProps) { @@ -465,6 +486,7 @@ function SharedAutomationCard(props: SharedAutomationCardProps) { automation={automation} ipLocationData={props.locationData} setToastMessage={props.setToastMessage} + chatModelOptions={props.chatModelOptions} /> ) : null; } @@ -476,6 +498,7 @@ const EditAutomationSchema = z.object({ dayOfMonth: z.optional(z.string()), timeRecurrence: z.string({ required_error: "Time Recurrence is required" }), schedulingRequest: z.string({ required_error: "Query to Run is required" }), + chatModelId: z.optional(z.number().nullable()), }); interface EditCardProps { @@ -488,6 +511,7 @@ interface EditCardProps { setShowLoginPrompt: (showLoginPrompt: boolean) => void; authenticatedData: UserProfile | null; setToastMessage: (toastMessage: string) => void; + chatModelOptions?: ModelOptions[]; } function EditCard(props: EditCardProps) { @@ -504,6 +528,7 @@ function EditCard(props: EditCardProps) { : "12:00 PM", dayOfMonth: automation?.crontime ? getDayOfMonthFromCron(automation.crontime) : "1", schedulingRequest: automation?.scheduling_request, + chatModelId: automation?.chat_model_id ?? undefined, }, }); @@ -532,6 +557,8 @@ function EditCard(props: EditCardProps) { updateQueryUrl += `&country=${encodeURIComponent(props.locationData.country)}`; if (props.locationData && props.locationData.timezone) updateQueryUrl += `&timezone=${encodeURIComponent(props.locationData.timezone)}`; + if (values.chatModelId != null) + updateQueryUrl += `&chat_model_id=${encodeURIComponent(values.chatModelId)}`; let method = props.createNew ? "POST" : "PUT"; @@ -547,6 +574,7 @@ function EditCard(props: EditCardProps) { schedule: cronToHumanReadableString(data.crontime), crontime: data.crontime, next: data.next, + chat_model_id: data.chat_model_id, }); }) .catch((error) => { @@ -603,6 +631,7 @@ function EditCard(props: EditCardProps) { create={props.createNew} isLoggedIn={props.isLoggedIn} setShowLoginPrompt={props.setShowLoginPrompt} + chatModelOptions={props.chatModelOptions} /> ); } @@ -615,6 +644,7 @@ interface AutomationModificationFormProps { setShowLoginPrompt: (showLoginPrompt: boolean) => void; authenticatedData: UserProfile | null; locationData: LocationData | null; + chatModelOptions?: ModelOptions[]; } function AutomationModificationForm(props: AutomationModificationFormProps) { @@ -860,6 +890,51 @@ function AutomationModificationForm(props: AutomationModificationFormProps) { )} /> + {props.chatModelOptions && props.chatModelOptions.length > 0 && ( + ( + + Model + + Which AI model should this automation use? + + + + + )} + /> + )}
{props.isLoggedIn ? ( isSaving ? ( @@ -925,6 +1000,7 @@ interface AutomationComponentWrapperProps { ipLocationData: LocationData | null | undefined; automation?: AutomationsData; setToastMessage: (toastMessage: string) => void; + chatModelOptions?: ModelOptions[]; } function AutomationComponentWrapper(props: AutomationComponentWrapperProps) { @@ -953,6 +1029,7 @@ function AutomationComponentWrapper(props: AutomationComponentWrapperProps) { setUpdatedAutomationData={props.setNewAutomationData} locationData={props.ipLocationData} setToastMessage={props.setToastMessage} + chatModelOptions={props.chatModelOptions} /> @@ -981,6 +1058,7 @@ function AutomationComponentWrapper(props: AutomationComponentWrapperProps) { setUpdatedAutomationData={props.setNewAutomationData} locationData={props.ipLocationData} setToastMessage={props.setToastMessage} + chatModelOptions={props.chatModelOptions} /> @@ -1001,6 +1079,8 @@ export default function Automations() { revalidateOnFocus: false, }); + const { models: chatModelOptions } = useChatModelOptions(); + const [isCreating, setIsCreating] = useState(false); const [newAutomationData, setNewAutomationData] = useState(null); const [allNewAutomations, setAllNewAutomations] = useState([]); @@ -1126,6 +1206,7 @@ export default function Automations() { isCreating={isCreating} ipLocationData={locationData} setToastMessage={setToastMessage} + chatModelOptions={chatModelOptions} /> ) : (
@@ -1193,6 +1277,7 @@ export default function Automations() { setShowLoginPrompt={setShowLoginPrompt} suggestedCard={true} setToastMessage={setToastMessage} + chatModelOptions={chatModelOptions} /> ))} diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py index f9a0ac1b1..3e23e3b23 100644 --- a/src/khoj/database/adapters/__init__.py +++ b/src/khoj/database/adapters/__init__.py @@ -1716,11 +1716,22 @@ async def aget_conversation_starters(user: KhojUser, max_results=3): return random.sample(all_questions, max_results) @staticmethod - async def aget_valid_chat_model(user: KhojUser, conversation: Conversation, is_subscribed: bool): + async def aget_valid_chat_model( + user: KhojUser, conversation: Conversation, is_subscribed: bool, chat_model_id: int = None + ): """ For paid users: Prefer any custom agent chat model > user default chat model > server default chat model. For free users: Prefer conversation specific agent's chat model > user default chat model > server default chat model. + An explicit chat_model_id override (e.g. from automations) takes highest priority. """ + if chat_model_id: + try: + chat_model = await ChatModel.objects.select_related("ai_model_api").aget(pk=chat_model_id) + if chat_model and chat_model.ai_model_api: + return chat_model + except ChatModel.DoesNotExist: + pass + agent: Agent = conversation.agent if await AgentAdapters.aget_default_agent() != conversation.agent else None if agent and agent.chat_model and (agent.is_hidden or is_subscribed): chat_model = await ChatModel.objects.select_related("ai_model_api").aget( @@ -2209,6 +2220,7 @@ def get_automation_metadata(user: KhojUser, automation: Job): "schedule": schedule, "crontime": crontime, "next": automation.next_run_time.strftime("%Y-%m-%d %I:%M %p %Z"), + "chat_model_id": automation_metadata.get("chat_model_id"), } @staticmethod diff --git a/src/khoj/routers/api_automation.py b/src/khoj/routers/api_automation.py index 8630a3246..47d7f6b2a 100644 --- a/src/khoj/routers/api_automation.py +++ b/src/khoj/routers/api_automation.py @@ -12,7 +12,7 @@ from starlette.authentication import requires from khoj.database.adapters import AutomationAdapters, ConversationAdapters -from khoj.database.models import KhojUser +from khoj.database.models import ChatModel, 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 @@ -59,6 +59,7 @@ def post_automation( region: Optional[str] = None, country: Optional[str] = None, timezone: Optional[str] = None, + chat_model_id: Optional[int] = None, ) -> Response: user: KhojUser = request.user.object @@ -95,6 +96,11 @@ def post_automation( status_code=400, ) + # Validate chat_model_id if provided + if chat_model_id is not None: + if not ChatModel.objects.filter(id=chat_model_id).exists(): + return Response(content="Invalid chat model", 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) @@ -104,7 +110,8 @@ def post_automation( # Use the query to run as the scheduling request if the scheduling request is unset calling_url = str(request.url.replace(query=f"{request.url.query}")) automation = schedule_automation( - query_to_run, subject, crontime, timezone, q, user, calling_url, str(conversation.id) + query_to_run, subject, crontime, timezone, q, user, calling_url, str(conversation.id), + chat_model_id=chat_model_id, ) except Exception as e: logger.error(f"Error creating automation {q} for {user.email}: {e}", exc_info=True) @@ -158,6 +165,7 @@ def edit_job( region: Optional[str] = None, country: Optional[str] = None, timezone: Optional[str] = None, + chat_model_id: Optional[int] = None, ) -> Response: user: KhojUser = request.user.object @@ -199,12 +207,18 @@ def edit_job( status_code=400, ) + # Validate chat_model_id if provided + if chat_model_id is not None: + if not ChatModel.objects.filter(id=chat_model_id).exists(): + return Response(content="Invalid chat model", 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 + automation_metadata["chat_model_id"] = chat_model_id conversation_id = automation_metadata.get("conversation_id") if not conversation_id: @@ -226,6 +240,7 @@ def edit_job( "user": user, "calling_url": str(request.url), "conversation_id": conversation_id, + "chat_model_id": chat_model_id, }, ) diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py index 79339c829..6af3e6b94 100644 --- a/src/khoj/routers/api_chat.py +++ b/src/khoj/routers/api_chat.py @@ -1400,6 +1400,7 @@ def collect_telemetry(): generated_asset_results, is_subscribed, tracer, + chat_model_id=body.chat_model_id, ) full_response = "" diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py index 29a0df0ff..a62d7a057 100644 --- a/src/khoj/routers/helpers.py +++ b/src/khoj/routers/helpers.py @@ -1910,6 +1910,7 @@ async def agenerate_chat_response( generated_asset_results: Dict[str, Dict] = {}, is_subscribed: bool = False, tracer: dict = {}, + chat_model_id: int = None, ) -> Tuple[AsyncGenerator[ResponseWithThought, None], Dict[str, str]]: # Initialize Variables chat_response_generator: AsyncGenerator[ResponseWithThought, None] = None @@ -1930,7 +1931,7 @@ async def agenerate_chat_response( operator_results = [] deepthought = True - chat_model = await ConversationAdapters.aget_valid_chat_model(user, conversation, is_subscribed) + chat_model = await ConversationAdapters.aget_valid_chat_model(user, conversation, is_subscribed, chat_model_id) max_prompt_size = await ConversationAdapters.aget_max_context_size(chat_model, user) vision_available = chat_model.vision_enabled if not vision_available and query_images: @@ -2521,6 +2522,7 @@ def scheduled_chat( calling_url: str | URL, job_id: str = None, conversation_id: str = None, + chat_model_id: int = None, ): logger.info(f"Processing scheduled_chat: {query_to_run}") if job_id: @@ -2561,6 +2563,10 @@ def scheduled_chat( # Restructure the original query_dict into a valid JSON payload for the chat API json_payload = {key: values[0] for key, values in query_dict.items()} + # Include the chat model override if specified for this automation + if chat_model_id is not None: + json_payload["chat_model_id"] = chat_model_id + # Construct the URL to call the chat API with the scheduled query string url = f"{scheme}://{parsed_url.netloc}/api/chat?client=khoj" @@ -2621,9 +2627,13 @@ async def create_automation( chat_history: List[ChatMessageModel] = [], conversation_id: str = None, tracer: dict = {}, + chat_model_id: int = None, ): crontime, query_to_run, subject = await aschedule_query(q, chat_history, user, tracer=tracer) - job = await aschedule_automation(query_to_run, subject, crontime, timezone, q, user, calling_url, conversation_id) + job = await aschedule_automation( + query_to_run, subject, crontime, timezone, q, user, calling_url, conversation_id, + chat_model_id=chat_model_id, + ) return job, crontime, query_to_run, subject @@ -2636,6 +2646,7 @@ def schedule_automation( user: KhojUser, calling_url: URL, conversation_id: str, + chat_model_id: int = None, ): # Disable minute level automation recurrence minute_value = crontime.split(" ")[0] @@ -2660,6 +2671,7 @@ def schedule_automation( "subject": subject, "crontime": crontime, "conversation_id": str(conversation_id), + "chat_model_id": chat_model_id, } ) query_id = hashlib.md5(f"{query_to_run}_{crontime}".encode("utf-8")).hexdigest() @@ -2679,6 +2691,7 @@ def schedule_automation( "calling_url": calling_url, "job_id": job_id, "conversation_id": conversation_id, + "chat_model_id": chat_model_id, }, id=job_id, name=job_metadata, @@ -2696,6 +2709,7 @@ async def aschedule_automation( user: KhojUser, calling_url: URL, conversation_id: str, + chat_model_id: int = None, ): # Disable minute level automation recurrence minute_value = crontime.split(" ")[0] @@ -2714,6 +2728,7 @@ async def aschedule_automation( "subject": subject, "crontime": crontime, "conversation_id": str(conversation_id), + "chat_model_id": chat_model_id, } ) query_id = hashlib.md5(f"{query_to_run}_{crontime}".encode("utf-8")).hexdigest() @@ -2733,6 +2748,7 @@ async def aschedule_automation( "calling_url": calling_url, "job_id": job_id, "conversation_id": conversation_id, + "chat_model_id": chat_model_id, }, id=job_id, name=job_metadata, diff --git a/src/khoj/utils/rawconfig.py b/src/khoj/utils/rawconfig.py index eeca58bbe..aa81a8729 100644 --- a/src/khoj/utils/rawconfig.py +++ b/src/khoj/utils/rawconfig.py @@ -100,6 +100,7 @@ class ChatRequestBody(BaseModel): images: Optional[list[str]] = None files: Optional[list[FileAttachment]] = [] create_new: Optional[bool] = False + chat_model_id: Optional[int] = None class Entry: