khoj 2.0.0b14.dev51__py3-none-any.whl → 2.0.0b15.dev23__py3-none-any.whl

khoj/routers/helpers.py

@@ -33,6 +33,7 @@ from apscheduler.triggers.cron import CronTrigger
 from asgiref.sync import sync_to_async
 from django.utils import timezone as django_timezone
 from fastapi import Depends, Header, HTTPException, Request, UploadFile, WebSocket
+from langchain_core.messages.chat import ChatMessage
 from pydantic import BaseModel, EmailStr, Field
 from starlette.authentication import has_required_scope
 from starlette.requests import URL
@@ -113,6 +114,7 @@ from khoj.utils import state
 from khoj.utils.helpers import (
     LRU,
     ConversationCommand,
+    ImageShape,
     ToolDefinition,
     get_file_type,
     in_debug_mode,
@@ -123,15 +125,16 @@ from khoj.utils.helpers import (
     mode_descriptions_for_llm,
     timer,
     tool_descriptions_for_llm,
+    truncate_code_context,
 )
 from khoj.utils.rawconfig import (
     ChatRequestBody,
-    FileAttachment,
     FileData,
     LocationData,
     SearchResponse,
 )
 from khoj.utils.state import SearchType
+from khoj.utils.yaml import yaml_dump
 
 logger = logging.getLogger(__name__)
 
@@ -286,7 +289,7 @@ async def acreate_title_from_history(
     title_generation_prompt = prompts.conversation_title_generation.format(chat_history=chat_history)
 
     with timer("Chat actor: Generate title from conversation history", logger):
-        response = await send_message_to_model_wrapper(title_generation_prompt, user=user)
+        response = await send_message_to_model_wrapper(title_generation_prompt, fast_model=True, user=user)
 
     return response.text.strip()
 
@@ -298,7 +301,7 @@ async def acreate_title_from_query(query: str, user: KhojUser = None) -> str:
     title_generation_prompt = prompts.subject_generation.format(query=query)
 
     with timer("Chat actor: Generate title from query", logger):
-        response = await send_message_to_model_wrapper(title_generation_prompt, user=user)
+        response = await send_message_to_model_wrapper(title_generation_prompt, fast_model=True, user=user)
 
     return response.text.strip()
 
@@ -321,7 +324,7 @@ async def acheck_if_safe_prompt(system_prompt: str, user: KhojUser = None, lax:
 
     with timer("Chat actor: Check if safe prompt", logger):
         response = await send_message_to_model_wrapper(
-            safe_prompt_check, user=user, response_type="json_object", response_schema=SafetyCheck
+            safe_prompt_check, response_type="json_object", response_schema=SafetyCheck, fast_model=True, user=user
         )
 
         response = response.text.strip()
@@ -342,7 +345,6 @@ async def acheck_if_safe_prompt(system_prompt: str, user: KhojUser = None, lax:
 async def aget_data_sources_and_output_format(
     query: str,
     chat_history: list[ChatMessageModel],
-    is_task: bool,
     user: KhojUser,
     query_images: List[str] = None,
     agent: Agent = None,
@@ -381,9 +383,6 @@ async def aget_data_sources_and_output_format(
 
     chat_history_str = construct_chat_history(chat_history, n=6)
 
-    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 ""
     )
@@ -405,11 +404,13 @@ async def aget_data_sources_and_output_format(
     with timer("Chat actor: Infer information sources to refer", logger):
         raw_response = await send_message_to_model_wrapper(
             relevant_tools_prompt,
+            query_files=query_files,
+            query_images=query_images,
             response_type="json_object",
             response_schema=PickTools,
-            user=user,
-            query_files=query_files,
+            fast_model=False,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
 
@@ -493,12 +494,13 @@ async def infer_webpage_urls(
     with timer("Chat actor: Infer webpage urls to read", logger):
         raw_response = await send_message_to_model_wrapper(
             online_queries_prompt,
+            query_files=query_files,
             query_images=query_images,
             response_type="json_object",
             response_schema=WebpageUrls,
-            user=user,
-            query_files=query_files,
+            fast_model=False,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
 
@@ -558,12 +560,13 @@ async def generate_online_subqueries(
     with timer("Chat actor: Generate online search subqueries", logger):
         raw_response = await send_message_to_model_wrapper(
             online_queries_prompt,
+            query_files=query_files,
             query_images=query_images,
             response_type="json_object",
             response_schema=OnlineQueries,
-            user=user,
-            query_files=query_files,
+            fast_model=False,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
 
@@ -625,7 +628,12 @@ async def aschedule_query(
     )
 
     raw_response = await send_message_to_model_wrapper(
-        crontime_prompt, query_images=query_images, response_type="json_object", user=user, tracer=tracer
+        crontime_prompt,
+        query_images=query_images,
+        response_type="json_object",
+        fast_model=False,
+        user=user,
+        tracer=tracer,
     )
 
     # Validate that the response is a non-empty, JSON-serializable list
@@ -663,9 +671,10 @@ async def extract_relevant_info(
 
     response = await send_message_to_model_wrapper(
         extract_relevant_information,
-        prompts.system_prompt_extract_relevant_information,
-        user=user,
+        system_message=prompts.system_prompt_extract_relevant_information,
+        fast_model=True,
         agent_chat_model=agent_chat_model,
+        user=user,
         tracer=tracer,
     )
     return response.text.strip()
@@ -705,10 +714,11 @@ async def extract_relevant_summary(
     with timer("Chat actor: Extract relevant information from data", logger):
         response = await send_message_to_model_wrapper(
             extract_relevant_information,
-            prompts.system_prompt_extract_relevant_summary,
-            user=user,
             query_images=query_images,
+            system_message=prompts.system_prompt_extract_relevant_summary,
+            fast_model=True,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
     return response.text.strip()
@@ -877,9 +887,10 @@ async def generate_better_diagram_description(
         response = await send_message_to_model_wrapper(
             improve_diagram_description_prompt,
             query_images=query_images,
-            user=user,
             query_files=query_files,
+            fast_model=False,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
         response = response.text.strip()
@@ -908,7 +919,11 @@ async def generate_excalidraw_diagram_from_description(
 
     with timer("Chat actor: Generate excalidraw diagram", logger):
         raw_response = await send_message_to_model_wrapper(
-            query=excalidraw_diagram_generation, user=user, agent_chat_model=agent_chat_model, tracer=tracer
+            query=excalidraw_diagram_generation,
+            fast_model=False,
+            agent_chat_model=agent_chat_model,
+            user=user,
+            tracer=tracer,
         )
         raw_response_text = clean_json(raw_response.text)
         try:
@@ -1027,10 +1042,11 @@ async def generate_better_mermaidjs_diagram_description(
     with timer("Chat actor: Generate better Mermaid.js diagram description", logger):
         response = await send_message_to_model_wrapper(
             improve_diagram_description_prompt,
-            query_images=query_images,
-            user=user,
             query_files=query_files,
+            query_images=query_images,
+            fast_model=False,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
         response_text = response.text.strip()
@@ -1059,14 +1075,18 @@ async def generate_mermaidjs_diagram_from_description(
 
     with timer("Chat actor: Generate Mermaid.js diagram", logger):
         raw_response = await send_message_to_model_wrapper(
-            query=mermaidjs_diagram_generation, user=user, agent_chat_model=agent_chat_model, tracer=tracer
+            query=mermaidjs_diagram_generation,
+            fast_model=False,
+            agent_chat_model=agent_chat_model,
+            user=user,
+            tracer=tracer,
         )
         return clean_mermaidjs(raw_response.text.strip())
 
 
 async def generate_better_image_prompt(
     q: str,
-    conversation_history: str,
+    conversation_history: List[ChatMessageModel],
     location_data: LocationData,
     note_references: List[Dict[str, Any]],
     online_results: Optional[dict] = None,
@@ -1076,12 +1096,11 @@ async def generate_better_image_prompt(
     agent: Agent = None,
     query_files: str = "",
     tracer: dict = {},
-) -> str:
+) -> dict:
     """
     Generate a better image prompt from the given query
     """
 
-    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 ""
     )
@@ -1089,58 +1108,53 @@ async def generate_better_image_prompt(
 
     location = f"{location_data}" if location_data else "Unknown"
 
-    user_references = "\n\n".join([f"# {item['compiled']}" for item in note_references])
+    user_references = "\n\n".join([f"- text:\n{item['compiled']}" for item in note_references])
 
     simplified_online_results = {}
+    for result in online_results or []:
+        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"]
 
-    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"]
+    enhance_image_system_message = prompts.enhance_image_system_message.format(
+        location=location,
+        references=user_references,
+        online_results=simplified_online_results or "",
+        personality_context=personality_context,
+    )
 
-    if model_type == TextToImageModelConfig.ModelType.OPENAI:
-        image_prompt = prompts.image_generation_improve_prompt_dalle.format(
-            query=q,
-            chat_history=conversation_history,
-            location=location,
-            current_date=today_date,
-            references=user_references,
-            online_results=simplified_online_results,
-            personality_context=personality_context,
-        )
-    elif model_type in [
-        TextToImageModelConfig.ModelType.STABILITYAI,
-        TextToImageModelConfig.ModelType.REPLICATE,
-        TextToImageModelConfig.ModelType.GOOGLE,
-    ]:
-        image_prompt = prompts.image_generation_improve_prompt_sd.format(
-            query=q,
-            chat_history=conversation_history,
-            location=location,
-            current_date=today_date,
-            references=user_references,
-            online_results=simplified_online_results,
-            personality_context=personality_context,
+    class ImagePromptResponse(BaseModel):
+        description: str = Field(description="Enhanced image description")
+        shape: ImageShape = Field(
+            description="Aspect ratio/shape best suited to render the image: Portrait, Landscape, or Square"
         )
 
     agent_chat_model = AgentAdapters.get_agent_chat_model(agent, user) if agent else None
 
     with timer("Chat actor: Generate contextual image prompt", logger):
-        response = await send_message_to_model_wrapper(
-            image_prompt,
-            query_images=query_images,
-            user=user,
+        raw_response = await send_message_to_model_wrapper(
+            q,
             query_files=query_files,
+            query_images=query_images,
+            chat_history=conversation_history,
+            system_message=enhance_image_system_message,
+            response_type="json_object",
+            response_schema=ImagePromptResponse,
+            fast_model=False,
             agent_chat_model=agent_chat_model,
+            user=user,
             tracer=tracer,
         )
-        response_text = response.text.strip()
-        if response_text.startswith(('"', "'")) and response_text.endswith(('"', "'")):
-            response_text = response_text[1:-1]
 
-    return response_text
+        # Parse the structured response
+        try:
+            response = clean_json(raw_response.text)
+            parsed_response = pyjson5.loads(response)
+            return parsed_response
+        except Exception:
+            # Fallback to user query as image description
+            return {"description": q, "shape": ImageShape.SQUARE}
 
 
 async def search_documents(
@@ -1150,7 +1164,7 @@ async def search_documents(
     user: KhojUser,
     chat_history: list[ChatMessageModel],
     conversation_id: str,
-    conversation_commands: List[ConversationCommand] = [ConversationCommand.Default],
+    conversation_commands: List[ConversationCommand] = [ConversationCommand.Notes],
     location_data: LocationData = None,
     send_status_func: Optional[Callable] = None,
     query_images: Optional[List[str]] = None,
@@ -1168,19 +1182,12 @@ async def search_documents(
     if agent:
         agent_has_entries = await sync_to_async(EntryAdapters.agent_has_entries)(agent=agent)
 
-    if (
-        ConversationCommand.Notes not in conversation_commands
-        and ConversationCommand.Default not in conversation_commands
-        and not agent_has_entries
-    ):
+    if ConversationCommand.Notes not in conversation_commands and not agent_has_entries:
         yield compiled_references, inferred_queries, q
         return
 
-    # If Notes or Default is not in the conversation command, then the search should be restricted to the agent's knowledge base
-    should_limit_to_agent_knowledge = (
-        ConversationCommand.Notes not in conversation_commands
-        and ConversationCommand.Default not in conversation_commands
-    )
+    # If Notes is not in the conversation command, then the search should be restricted to the agent's knowledge base
+    should_limit_to_agent_knowledge = ConversationCommand.Notes not in conversation_commands
 
     if not await sync_to_async(EntryAdapters.user_has_entries)(user=user):
         if not agent_has_entries:
@@ -1209,11 +1216,12 @@ async def search_documents(
         inferred_queries = await extract_questions(
             query=defiltered_query,
             user=user,
+            query_files=query_files,
+            query_images=query_images,
             personality_context=personality_context,
-            chat_history=chat_history,
             location_data=location_data,
-            query_images=query_images,
-            query_files=query_files,
+            chat_history=chat_history,
+            agent=agent,
             tracer=tracer,
         )
 
@@ -1259,12 +1267,13 @@ async def search_documents(
 async def extract_questions(
     query: str,
     user: KhojUser,
+    query_files: str = None,
+    query_images: Optional[List[str]] = None,
     personality_context: str = "",
-    chat_history: List[ChatMessageModel] = [],
     location_data: LocationData = None,
-    query_images: Optional[List[str]] = None,
-    query_files: str = None,
+    chat_history: List[ChatMessageModel] = [],
     max_queries: int = 5,
+    agent: Agent = None,
     tracer: dict = {},
 ):
     """
@@ -1309,13 +1318,17 @@ async def extract_questions(
             description="List of semantic search queries to run on user documents.",
         )
 
+    agent_chat_model = AgentAdapters.get_agent_chat_model(agent, user) if agent else None
+
     raw_response = await send_message_to_model_wrapper(
-        system_message=system_prompt,
         query=prompt,
-        query_images=query_images,
         query_files=query_files,
+        query_images=query_images,
+        system_message=system_prompt,
         response_type="json_object",
         response_schema=DocumentQueries,
+        fast_model=False,
+        agent_chat_model=agent_chat_model,
         user=user,
         tracer=tracer,
     )
@@ -1428,21 +1441,26 @@ async def execute_search(
 
 
 async def send_message_to_model_wrapper(
+    # Context
     query: str,
+    query_files: str = None,
+    query_images: List[str] = None,
+    context: str = "",
+    chat_history: list[ChatMessageModel] = [],
     system_message: str = "",
+    # Model Config
     response_type: str = "text",
     response_schema: BaseModel = None,
     tools: List[ToolDefinition] = None,
     deepthought: bool = False,
-    user: KhojUser = None,
-    query_images: List[str] = None,
-    context: str = "",
-    query_files: str = None,
-    chat_history: list[ChatMessageModel] = [],
+    fast_model: Optional[bool] = None,
     agent_chat_model: ChatModel = None,
+    # User
+    user: KhojUser = None,
+    # Tracer
     tracer: dict = {},
 ):
-    chat_model: ChatModel = await ConversationAdapters.aget_default_chat_model(user, agent_chat_model)
+    chat_model: ChatModel = await ConversationAdapters.aget_default_chat_model(user, agent_chat_model, fast=fast_model)
     vision_available = chat_model.vision_enabled
     if not vision_available and query_images:
         logger.warning(f"Vision is not enabled for default model: {chat_model.name}.")
@@ -1463,16 +1481,16 @@ async def send_message_to_model_wrapper(
 
     truncated_messages = generate_chatml_messages_with_context(
         user_message=query,
+        query_files=query_files,
+        query_images=query_images,
         context_message=context,
-        system_message=system_message,
         chat_history=chat_history,
+        system_message=system_message,
         model_name=chat_model_name,
+        model_type=model_type,
         tokenizer_name=tokenizer,
         max_prompt_size=max_tokens,
         vision_enabled=vision_available,
-        query_images=query_images,
-        model_type=model_type,
-        query_files=query_files,
     )
 
     if model_type == ChatModel.ModelType.OPENAI:
@@ -1540,14 +1558,14 @@ def send_message_to_model_wrapper_sync(
 
     truncated_messages = generate_chatml_messages_with_context(
         user_message=message,
-        system_message=system_message,
+        query_files=query_files,
+        query_images=query_images,
         chat_history=chat_history,
+        system_message=system_message,
         model_name=chat_model_name,
+        model_type=model_type,
         max_prompt_size=max_tokens,
         vision_enabled=vision_available,
-        model_type=model_type,
-        query_images=query_images,
-        query_files=query_files,
     )
 
     if model_type == ChatModel.ModelType.OPENAI:
@@ -1585,6 +1603,105 @@ def send_message_to_model_wrapper_sync(
         raise HTTPException(status_code=500, detail="Invalid conversation config")
 
 
+def build_conversation_context(
+    # Query and Context
+    user_query: str,
+    references: List[Dict],
+    online_results: Dict[str, Dict],
+    code_results: Dict[str, Dict],
+    operator_results: List[OperatorRun],
+    query_files: str = None,
+    query_images: Optional[List[str]] = None,
+    generated_asset_results: Dict[str, Dict] = {},
+    program_execution_context: List[str] = None,
+    chat_history: List[ChatMessageModel] = [],
+    location_data: LocationData = None,
+    user_name: str = None,
+    # Model config
+    agent: Agent = None,
+    model_name: str = None,
+    model_type: ChatModel.ModelType = None,
+    max_prompt_size: int = None,
+    tokenizer_name: str = None,
+    vision_available: bool = False,
+) -> List[ChatMessage]:
+    """
+    Construct system, context and chatml messages for chat response.
+    Share common logic across different model types.
+
+    Returns:
+        List of ChatMessages with context
+    """
+    # Initialize Variables
+    current_date = datetime.now()
+
+    # Build system prompt
+    if agent and agent.personality:
+        system_prompt = prompts.custom_personality.format(
+            name=agent.name,
+            bio=agent.personality,
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
+    else:
+        system_prompt = prompts.personality.format(
+            current_date=current_date.strftime("%Y-%m-%d"),
+            day_of_week=current_date.strftime("%A"),
+        )
+
+    # Add Gemini-specific personality enhancement
+    if model_type == ChatModel.ModelType.GOOGLE:
+        system_prompt += f"\n\n{prompts.gemini_verbose_language_personality}"
+
+    # Add location context if available
+    if location_data:
+        location_prompt = prompts.user_location.format(location=f"{location_data}")
+        system_prompt += f"\n{location_prompt}"
+
+    # Add user name context if available
+    if user_name:
+        user_name_prompt = prompts.user_name.format(name=user_name)
+        system_prompt += f"\n{user_name_prompt}"
+
+    # Build context message
+    context_message = ""
+    if not is_none_or_empty(references):
+        context_message = f"{prompts.notes_conversation.format(references=yaml_dump(references))}\n\n"
+    if not is_none_or_empty(online_results):
+        context_message += f"{prompts.online_search_conversation.format(online_results=yaml_dump(online_results))}\n\n"
+    if not is_none_or_empty(code_results):
+        context_message += (
+            f"{prompts.code_executed_context.format(code_results=truncate_code_context(code_results))}\n\n"
+        )
+    if not is_none_or_empty(operator_results):
+        operator_content = [
+            {"query": oc.query, "response": oc.response, "webpages": oc.webpages} for oc in operator_results
+        ]
+        context_message += (
+            f"{prompts.operator_execution_context.format(operator_results=yaml_dump(operator_content))}\n\n"
+        )
+    context_message = context_message.strip()
+
+    # Generate the chatml messages
+    messages = generate_chatml_messages_with_context(
+        user_message=user_query,
+        query_files=query_files,
+        query_images=query_images,
+        context_message=context_message,
+        generated_asset_results=generated_asset_results,
+        program_execution_context=program_execution_context,
+        chat_history=chat_history,
+        system_message=system_prompt,
+        model_name=model_name,
+        model_type=model_type,
+        max_prompt_size=max_prompt_size,
+        tokenizer_name=tokenizer_name,
+        vision_enabled=vision_available,
+    )
+
+    return messages
+
+
 async def agenerate_chat_response(
     q: str,
     chat_history: List[ChatMessageModel],
@@ -1599,7 +1716,6 @@ async def agenerate_chat_response(
     user_name: Optional[str] = None,
     query_images: Optional[List[str]] = None,
     query_files: str = None,
-    raw_generated_files: List[FileAttachment] = [],
     program_execution_context: List[str] = [],
     generated_asset_results: Dict[str, Dict] = {},
     is_subscribed: bool = False,
@@ -1633,34 +1749,39 @@ async def agenerate_chat_response(
                 chat_model = vision_enabled_config
                 vision_available = True
 
+        # Build shared conversation context and generate chatml messages
+        messages = build_conversation_context(
+            user_query=query_to_run,
+            references=compiled_references,
+            online_results=online_results,
+            code_results=code_results,
+            operator_results=operator_results,
+            query_files=query_files,
+            query_images=query_images,
+            generated_asset_results=generated_asset_results,
+            program_execution_context=program_execution_context,
+            chat_history=chat_history,
+            location_data=location_data,
+            user_name=user_name,
+            agent=agent,
+            model_type=chat_model.model_type,
+            model_name=chat_model.name,
+            max_prompt_size=max_prompt_size,
+            tokenizer_name=chat_model.tokenizer,
+            vision_available=vision_available,
+        )
+
         if chat_model.model_type == ChatModel.ModelType.OPENAI:
             openai_chat_config = chat_model.ai_model_api
             api_key = openai_chat_config.api_key
             chat_model_name = chat_model.name
             chat_response_generator = converse_openai(
-                # Query
-                query_to_run,
-                # Context
-                references=compiled_references,
-                online_results=online_results,
-                code_results=code_results,
-                operator_results=operator_results,
-                query_images=query_images,
-                query_files=query_files,
-                generated_files=raw_generated_files,
-                generated_asset_results=generated_asset_results,
-                program_execution_context=program_execution_context,
-                location_data=location_data,
-                user_name=user_name,
-                chat_history=chat_history,
+                # Query + Context Messages
+                messages,
                 # Model
                 model=chat_model_name,
                 api_key=api_key,
                 api_base_url=openai_chat_config.api_base_url,
-                max_prompt_size=max_prompt_size,
-                tokenizer_name=chat_model.tokenizer,
-                agent=agent,
-                vision_available=vision_available,
                 deepthought=deepthought,
                 tracer=tracer,
             )
@@ -1669,29 +1790,12 @@ async def agenerate_chat_response(
             api_key = chat_model.ai_model_api.api_key
             api_base_url = chat_model.ai_model_api.api_base_url
             chat_response_generator = converse_anthropic(
-                # Query
-                query_to_run,
-                # Context
-                references=compiled_references,
-                online_results=online_results,
-                code_results=code_results,
-                operator_results=operator_results,
-                query_images=query_images,
-                query_files=query_files,
-                generated_files=raw_generated_files,
-                generated_asset_results=generated_asset_results,
-                program_execution_context=program_execution_context,
-                location_data=location_data,
-                user_name=user_name,
-                chat_history=chat_history,
+                # Query + Context Messages
+                messages,
                 # Model
                 model=chat_model.name,
                 api_key=api_key,
                 api_base_url=api_base_url,
-                max_prompt_size=max_prompt_size,
-                tokenizer_name=chat_model.tokenizer,
-                agent=agent,
-                vision_available=vision_available,
                 deepthought=deepthought,
                 tracer=tracer,
             )
@@ -1699,29 +1803,12 @@ async def agenerate_chat_response(
             api_key = chat_model.ai_model_api.api_key
             api_base_url = chat_model.ai_model_api.api_base_url
             chat_response_generator = converse_gemini(
-                # Query
-                query_to_run,
-                # Context
-                references=compiled_references,
-                online_results=online_results,
-                code_results=code_results,
-                operator_results=operator_results,
-                query_images=query_images,
-                query_files=query_files,
-                generated_files=raw_generated_files,
-                generated_asset_results=generated_asset_results,
-                program_execution_context=program_execution_context,
-                location_data=location_data,
-                user_name=user_name,
-                chat_history=chat_history,
+                # Query + Context Messages
+                messages,
                 # Model
                 model=chat_model.name,
                 api_key=api_key,
                 api_base_url=api_base_url,
-                max_prompt_size=max_prompt_size,
-                tokenizer_name=chat_model.tokenizer,
-                agent=agent,
-                vision_available=vision_available,
                 deepthought=deepthought,
                 tracer=tracer,
             )