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

khoj/processor/conversation/openai/gpt.py

@@ -1,14 +1,11 @@
-import asyncio
 import logging
-from datetime import datetime, timedelta
+from datetime import datetime
 from typing import AsyncGenerator, Dict, List, Optional
 
-import pyjson5
-from langchain_core.messages.chat import ChatMessage
 from openai.lib._pydantic import _ensure_strict_json_schema
 from pydantic import BaseModel
 
-from khoj.database.models import Agent, ChatModel, KhojUser
+from khoj.database.models import Agent, ChatMessageModel, ChatModel
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.openai.utils import (
     chat_completion_with_backoff,
@@ -19,105 +16,16 @@ from khoj.processor.conversation.utils import (
     JsonSupport,
     OperatorRun,
     ResponseWithThought,
-    clean_json,
-    construct_question_history,
-    construct_structured_message,
     generate_chatml_messages_with_context,
     messages_to_print,
 )
-from khoj.utils.helpers import (
-    ConversationCommand,
-    is_none_or_empty,
-    truncate_code_context,
-)
+from khoj.utils.helpers import is_none_or_empty, truncate_code_context
 from khoj.utils.rawconfig import FileAttachment, LocationData
 from khoj.utils.yaml import yaml_dump
 
 logger = logging.getLogger(__name__)
 
 
-def extract_questions(
-    text,
-    model: Optional[str] = "gpt-4o-mini",
-    conversation_log={},
-    api_key=None,
-    api_base_url=None,
-    location_data: LocationData = None,
-    user: KhojUser = None,
-    query_images: Optional[list[str]] = None,
-    vision_enabled: bool = False,
-    personality_context: Optional[str] = None,
-    query_files: str = None,
-    tracer: dict = {},
-):
-    """
-    Infer search queries to retrieve relevant notes to answer user query
-    """
-    location = f"{location_data}" if location_data else "Unknown"
-    username = prompts.user_name.format(name=user.get_full_name()) if user and user.get_full_name() else ""
-
-    # Extract Past User Message and Inferred Questions from Conversation Log
-    chat_history = construct_question_history(conversation_log)
-
-    # Get dates relative to today for prompt creation
-    today = datetime.today()
-    current_new_year = today.replace(month=1, day=1)
-    last_new_year = current_new_year.replace(year=today.year - 1)
-
-    prompt = prompts.extract_questions.format(
-        current_date=today.strftime("%Y-%m-%d"),
-        day_of_week=today.strftime("%A"),
-        current_month=today.strftime("%Y-%m"),
-        last_new_year=last_new_year.strftime("%Y"),
-        last_new_year_date=last_new_year.strftime("%Y-%m-%d"),
-        current_new_year_date=current_new_year.strftime("%Y-%m-%d"),
-        bob_tom_age_difference={current_new_year.year - 1984 - 30},
-        bob_age={current_new_year.year - 1984},
-        chat_history=chat_history,
-        text=text,
-        yesterday_date=(today - timedelta(days=1)).strftime("%Y-%m-%d"),
-        location=location,
-        username=username,
-        personality_context=personality_context,
-    )
-
-    prompt = construct_structured_message(
-        message=prompt,
-        images=query_images,
-        model_type=ChatModel.ModelType.OPENAI,
-        vision_enabled=vision_enabled,
-        attached_file_context=query_files,
-    )
-
-    messages = []
-    messages.append(ChatMessage(content=prompt, role="user"))
-
-    response = send_message_to_model(
-        messages,
-        api_key,
-        model,
-        response_type="json_object",
-        api_base_url=api_base_url,
-        tracer=tracer,
-    )
-
-    # Extract, Clean Message from GPT's Response
-    try:
-        response = clean_json(response)
-        response = pyjson5.loads(response)
-        response = [q.strip() for q in response["queries"] if q.strip()]
-        if not isinstance(response, list) or not response:
-            logger.error(f"Invalid response for constructing subqueries: {response}")
-            return [text]
-        return response
-    except:
-        logger.warning(f"GPT returned invalid JSON. Falling back to using user message as search query.\n{response}")
-        questions = [text]
-
-    logger.debug(f"Extracted Questions by GPT: {questions}")
-    return questions
-
-
 def send_message_to_model(
     messages,
     api_key,
@@ -161,32 +69,32 @@ def send_message_to_model(
 
 
 async def converse_openai(
+    # Query
     user_query: str,
+    # Context
     references: list[dict],
     online_results: Optional[Dict[str, Dict]] = None,
     code_results: Optional[Dict[str, Dict]] = None,
     operator_results: Optional[List[OperatorRun]] = None,
-    conversation_log={},
+    query_images: Optional[list[str]] = None,
+    query_files: str = None,
+    generated_files: List[FileAttachment] = None,
+    generated_asset_results: Dict[str, Dict] = {},
+    program_execution_context: List[str] = None,
+    location_data: LocationData = None,
+    chat_history: list[ChatMessageModel] = [],
     model: str = "gpt-4o-mini",
     api_key: Optional[str] = None,
     api_base_url: Optional[str] = None,
     temperature: float = 0.4,
-    completion_func=None,
-    conversation_commands=[ConversationCommand.Default],
     max_prompt_size=None,
     tokenizer_name=None,
-    location_data: LocationData = None,
     user_name: str = None,
     agent: Agent = None,
-    query_images: Optional[list[str]] = None,
     vision_available: bool = False,
-    query_files: str = None,
-    generated_files: List[FileAttachment] = None,
-    generated_asset_results: Dict[str, Dict] = {},
-    program_execution_context: List[str] = None,
     deepthought: Optional[bool] = False,
     tracer: dict = {},
-) -> AsyncGenerator[str | ResponseWithThought, None]:
+) -> AsyncGenerator[ResponseWithThought, None]:
     """
     Converse with user using OpenAI's ChatGPT
     """
@@ -214,20 +122,6 @@ async def converse_openai(
         user_name_prompt = prompts.user_name.format(name=user_name)
         system_prompt = f"{system_prompt}\n{user_name_prompt}"
 
-    # Get Conversation Primer appropriate to Conversation Type
-    if conversation_commands == [ConversationCommand.Notes] and is_none_or_empty(references):
-        response = prompts.no_notes_found.format()
-        if completion_func:
-            asyncio.create_task(completion_func(chat_response=response))
-        yield response
-        return
-    elif conversation_commands == [ConversationCommand.Online] and is_none_or_empty(online_results):
-        response = prompts.no_online_results_found.format()
-        if completion_func:
-            asyncio.create_task(completion_func(chat_response=response))
-        yield response
-        return
-
     context_message = ""
     if not is_none_or_empty(references):
         context_message = f"{prompts.notes_conversation.format(references=yaml_dump(references))}\n\n"
@@ -251,7 +145,7 @@ async def converse_openai(
     messages = generate_chatml_messages_with_context(
         user_query,
         system_prompt,
-        conversation_log,
+        chat_history,
         context_message=context_message,
         model_name=model,
         max_prompt_size=max_prompt_size,
@@ -267,7 +161,6 @@ async def converse_openai(
     logger.debug(f"Conversation Context for GPT: {messages_to_print(messages)}")
 
     # Get Response from GPT
-    full_response = ""
     async for chunk in chat_completion_with_backoff(
         messages=messages,
         model_name=model,
@@ -275,17 +168,10 @@ async def converse_openai(
         openai_api_key=api_key,
         api_base_url=api_base_url,
         deepthought=deepthought,
-        model_kwargs={"stop": ["Notes:\n["]},
         tracer=tracer,
     ):
-        if chunk.response:
-            full_response += chunk.response
         yield chunk
 
-    # Call completion_func once finish streaming and we have the full response
-    if completion_func:
-        asyncio.create_task(completion_func(chat_response=full_response))
-
 
 def clean_response_schema(schema: BaseModel | dict) -> dict:
     """

khoj/processor/conversation/prompts.py

@@ -549,68 +549,7 @@ Q: {query}
 )
 
 
-extract_questions = PromptTemplate.from_template(
-    """
-You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes and documents.
-Construct search queries to retrieve relevant information to answer the user's question.
-- You will be provided example and actual past user questions(Q), search queries(Khoj) and answers(A) for context.
-- Add as much context from the previous questions and answers as required into your search queries.
-- Break your search down into multiple search queries from a diverse set of lenses to retrieve all related documents.
-- Add date filters to your search queries from questions and answers when required to retrieve the relevant information.
-- When asked a meta, vague or random questions, search for a variety of broad topics to answer the user's question.
-{personality_context}
-What searches will you perform to answer the user's question? Respond with search queries as list of strings in a JSON object.
-Current Date: {day_of_week}, {current_date}
-User's Location: {location}
-{username}
-
-Examples
----
-Q: How was my trip to Cambodia?
-Khoj: {{"queries": ["How was my trip to Cambodia?", "Angkor Wat temple visit", "Flight to Phnom Penh", "Expenses in Cambodia", "Stay in Cambodia"]}}
-A: The trip was amazing. You went to the Angkor Wat temple and it was beautiful.
-
-Q: Who did i visit that temple with?
-Khoj: {{"queries": ["Who did I visit the Angkor Wat Temple in Cambodia with?"]}}
-A: You visited the Angkor Wat Temple in Cambodia with Pablo, Namita and Xi.
-
-Q: What national parks did I go to last year?
-Khoj: {{"queries": ["National park I visited in {last_new_year} dt>='{last_new_year_date}' dt<'{current_new_year_date}'"]}}
-A: You visited the Grand Canyon and Yellowstone National Park in {last_new_year}.
-
-Q: How can you help me?
-Khoj: {{"queries": ["Social relationships", "Physical and mental health", "Education and career", "Personal life goals and habits"]}}
-A: I can help you live healthier and happier across work and personal life
-
-Q: How many tennis balls fit in the back of a 2002 Honda Civic?
-Khoj: {{"queries": ["What is the size of a tennis ball?", "What is the trunk size of a 2002 Honda Civic?"]}}
-A: 1085 tennis balls will fit in the trunk of a Honda Civic
-
-Q: Share some random, interesting experiences from this month
-Khoj: {{"queries": ["Exciting travel adventures from {current_month}", "Fun social events dt>='{current_month}-01' dt<'{current_date}'", "Intense emotional experiences in {current_month}"]}}
-A: You had a great time at the local beach with your friends, attended a music concert and had a deep conversation with your friend, Khalid.
-
-Q: Is Bob older than Tom?
-Khoj: {{"queries": ["When was Bob born?", "What is Tom's age?"]}}
-A: Yes, Bob is older than Tom. As Bob was born on 1984-01-01 and Tom is 30 years old.
-
-Q: What is their age difference?
-Khoj: {{"queries": ["What is Bob's age?", "What is Tom's age?"]}}
-A: Bob is {bob_tom_age_difference} years older than Tom. As Bob is {bob_age} years old and Tom is 30 years old.
-
-Q: Who all did I meet here yesterday?
-Khoj: {{"queries": ["Met in {location} on {yesterday_date} dt>='{yesterday_date}' dt<'{current_date}'"]}}
-A: Yesterday's note mentions your visit to your local beach with Ram and Shyam.
-
-Actual
----
-{chat_history}
-Q: {text}
-Khoj:
-""".strip()
-)
-
-extract_questions_anthropic_system_prompt = PromptTemplate.from_template(
+extract_questions_system_prompt = PromptTemplate.from_template(
     """
 You are Khoj, an extremely smart and helpful document search assistant with only the ability to retrieve information from the user's notes.
 Construct search queries to retrieve relevant information to answer the user's question.
@@ -651,7 +590,7 @@ A: You had a great time at the local beach with your friends, attended a music c
 """.strip()
 )
 
-extract_questions_anthropic_user_message = PromptTemplate.from_template(
+extract_questions_user_message = PromptTemplate.from_template(
     """
 Here's our most recent chat history:
 {chat_history}

khoj/processor/conversation/utils.py

@@ -24,7 +24,13 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
 
 from khoj.database.adapters import ConversationAdapters
-from khoj.database.models import ChatModel, ClientApplication, KhojUser
+from khoj.database.models import (
+    ChatMessageModel,
+    ChatModel,
+    ClientApplication,
+    Intent,
+    KhojUser,
+)
 from khoj.processor.conversation import prompts
 from khoj.processor.conversation.offline.utils import download_model, infer_max_tokens
 from khoj.search_filter.base_filter import BaseFilter
@@ -161,8 +167,8 @@ def construct_iteration_history(
     previous_iterations: List[ResearchIteration],
     previous_iteration_prompt: str,
     query: str = None,
-) -> list[dict]:
-    iteration_history: list[dict] = []
+) -> list[ChatMessageModel]:
+    iteration_history: list[ChatMessageModel] = []
     previous_iteration_messages: list[dict] = []
     for idx, iteration in enumerate(previous_iterations):
         iteration_data = previous_iteration_prompt.format(
@@ -176,46 +182,46 @@ def construct_iteration_history(
 
     if previous_iteration_messages:
         if query:
-            iteration_history.append({"by": "you", "message": query})
+            iteration_history.append(ChatMessageModel(by="you", message=query))
         iteration_history.append(
-            {
-                "by": "khoj",
-                "intent": {"type": "remember", "query": query},
-                "message": previous_iteration_messages,
-            }
+            ChatMessageModel(
+                by="khoj",
+                intent={"type": "remember", "query": query},
+                message=previous_iteration_messages,
+            )
         )
     return iteration_history
 
 
-def construct_chat_history(conversation_history: dict, n: int = 4, agent_name="AI") -> str:
-    chat_history = ""
-    for chat in conversation_history.get("chat", [])[-n:]:
-        if chat["by"] == "khoj" and chat["intent"].get("type") in ["remember", "reminder", "summarize"]:
-            if chat["intent"].get("inferred-queries"):
-                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 chat.get("images"):
-            chat_history += f"User: {chat['intent']['query']}\n"
-            chat_history += f"{agent_name}: [generated image redacted for space]\n"
-        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":
-            chat_history += f"User: {chat['message']}\n"
-            raw_query_files = chat.get("queryFiles")
+def construct_chat_history(chat_history: list[ChatMessageModel], n: int = 4, agent_name="AI") -> str:
+    chat_history_str = ""
+    for chat in chat_history[-n:]:
+        if chat.by == "khoj" and chat.intent.type in ["remember", "reminder", "summarize"]:
+            if chat.intent.inferred_queries:
+                chat_history_str += f'{agent_name}: {{"queries": {chat.intent.inferred_queries}}}\n'
+            chat_history_str += f"{agent_name}: {chat.message}\n\n"
+        elif chat.by == "khoj" and chat.images:
+            chat_history_str += f"User: {chat.intent.query}\n"
+            chat_history_str += f"{agent_name}: [generated image redacted for space]\n"
+        elif chat.by == "khoj" and ("excalidraw" in chat.intent.type):
+            chat_history_str += f"User: {chat.intent.query}\n"
+            chat_history_str += f"{agent_name}: {chat.intent.inferred_queries[0]}\n"
+        elif chat.by == "you":
+            chat_history_str += f"User: {chat.message}\n"
+            raw_query_files = chat.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"
+                chat_history_str += f"User: {query_file_context}\n"
 
-    return chat_history
+    return chat_history_str
 
 
 def construct_question_history(
-    conversation_log: dict,
+    conversation_log: list[ChatMessageModel],
     include_query: bool = True,
     lookback: int = 6,
     query_prefix: str = "Q",
@@ -226,16 +232,16 @@ def construct_question_history(
     """
     history_parts = ""
     original_query = None
-    for chat in conversation_log.get("chat", [])[-lookback:]:
-        if chat["by"] == "you":
-            original_query = chat.get("message")
+    for chat in conversation_log[-lookback:]:
+        if chat.by == "you":
+            original_query = json.dumps(chat.message)
             history_parts += f"{query_prefix}: {original_query}\n"
-        if chat["by"] == "khoj":
+        if chat.by == "khoj":
             if original_query is None:
                 continue
 
-            message = chat.get("message", "")
-            inferred_queries_list = chat.get("intent", {}).get("inferred-queries")
+            message = chat.message
+            inferred_queries_list = chat.intent.inferred_queries or []
 
             # Ensure inferred_queries_list is a list, defaulting to the original query in a list
             if not inferred_queries_list:
@@ -246,7 +252,7 @@ def construct_question_history(
 
             if include_query:
                 # Ensure 'type' exists and is a string before checking 'to-image'
-                intent_type = chat.get("intent", {}).get("type", "")
+                intent_type = chat.intent.type if chat.intent and chat.intent.type else ""
                 if "to-image" not in intent_type:
                     history_parts += f'{agent_name}: {{"queries": {inferred_queries_list}}}\n'
                     history_parts += f"A: {message}\n\n"
@@ -259,7 +265,7 @@ def construct_question_history(
     return history_parts
 
 
-def construct_chat_history_for_operator(conversation_history: dict, n: int = 6) -> list[AgentMessage]:
+def construct_chat_history_for_operator(conversation_history: List[ChatMessageModel], n: int = 6) -> list[AgentMessage]:
     """
     Construct chat history for operator agent in conversation log.
     Only include last n completed turns (i.e with user and khoj message).
@@ -267,22 +273,22 @@ def construct_chat_history_for_operator(conversation_history: dict, n: int = 6)
     chat_history: list[AgentMessage] = []
     user_message: Optional[AgentMessage] = None
 
-    for chat in conversation_history.get("chat", []):
+    for chat in conversation_history:
         if len(chat_history) >= n:
             break
-        if chat["by"] == "you" and chat.get("message"):
-            content = [{"type": "text", "text": chat["message"]}]
-            for file in chat.get("queryFiles", []):
+        if chat.by == "you" and chat.message:
+            content = [{"type": "text", "text": chat.message}]
+            for file in chat.queryFiles or []:
                 content += [{"type": "text", "text": f'## File: {file["name"]}\n\n{file["content"]}'}]
             user_message = AgentMessage(role="user", content=content)
-        elif chat["by"] == "khoj" and chat.get("message"):
-            chat_history += [user_message, AgentMessage(role="assistant", content=chat["message"])]
+        elif chat.by == "khoj" and chat.message:
+            chat_history += [user_message, AgentMessage(role="assistant", content=chat.message)]
     return chat_history
 
 
 def construct_tool_chat_history(
     previous_iterations: List[ResearchIteration], tool: ConversationCommand = None
-) -> Dict[str, list]:
+) -> List[ChatMessageModel]:
     """
     Construct chat history from previous iterations for a specific tool
 
@@ -313,22 +319,23 @@ def construct_tool_chat_history(
             tool or ConversationCommand(iteration.tool), base_extractor
         )
         chat_history += [
-            {
-                "by": "you",
-                "message": iteration.query,
-            },
-            {
-                "by": "khoj",
-                "intent": {
-                    "type": "remember",
-                    "inferred-queries": inferred_query_extractor(iteration),
-                    "query": iteration.query,
-                },
-                "message": iteration.summarizedResult,
-            },
+            ChatMessageModel(
+                by="you",
+                message=iteration.query,
+            ),
+            ChatMessageModel(
+                by="khoj",
+                intent=Intent(
+                    type="remember",
+                    query=iteration.query,
+                    inferred_queries=inferred_query_extractor(iteration),
+                    memory_type="notes",
+                ),
+                message=iteration.summarizedResult,
+            ),
         ]
 
-    return {"chat": chat_history}
+    return chat_history
 
 
 class ChatEvent(Enum):
@@ -349,8 +356,8 @@ def message_to_log(
     chat_response,
     user_message_metadata={},
     khoj_message_metadata={},
-    conversation_log=[],
-):
+    chat_history: List[ChatMessageModel] = [],
+) -> List[ChatMessageModel]:
     """Create json logs from messages, metadata for conversation log"""
     default_khoj_message_metadata = {
         "intent": {"type": "remember", "memory-type": "notes", "query": user_message},
@@ -369,15 +376,17 @@ def message_to_log(
     khoj_log = merge_dicts(khoj_message_metadata, default_khoj_message_metadata)
     khoj_log = merge_dicts({"message": chat_response, "by": "khoj", "created": khoj_response_time}, khoj_log)
 
-    conversation_log.extend([human_log, khoj_log])
-    return conversation_log
+    human_message = ChatMessageModel(**human_log)
+    khoj_message = ChatMessageModel(**khoj_log)
+    chat_history.extend([human_message, khoj_message])
+    return chat_history
 
 
 async def save_to_conversation_log(
     q: str,
     chat_response: str,
     user: KhojUser,
-    meta_log: Dict,
+    chat_history: List[ChatMessageModel],
     user_message_time: str = None,
     compiled_references: List[Dict[str, Any]] = [],
     online_results: Dict[str, Any] = {},
@@ -427,11 +436,11 @@ async def save_to_conversation_log(
         chat_response=chat_response,
         user_message_metadata=user_message_metadata,
         khoj_message_metadata=khoj_message_metadata,
-        conversation_log=meta_log.get("chat", []),
+        chat_history=chat_history,
     )
     await ConversationAdapters.save_conversation(
         user,
-        {"chat": updated_conversation},
+        updated_conversation,
         client_application=client_application,
         conversation_id=conversation_id,
         user_message=q,
@@ -502,7 +511,7 @@ def gather_raw_query_files(
 def generate_chatml_messages_with_context(
     user_message: str,
     system_message: str = None,
-    conversation_log={},
+    chat_history: list[ChatMessageModel] = [],
     model_name="gpt-4o-mini",
     loaded_model: Optional[Llama] = None,
     max_prompt_size=None,
@@ -529,21 +538,21 @@ def generate_chatml_messages_with_context(
 
     # Extract Chat History for Context
     chatml_messages: List[ChatMessage] = []
-    for chat in conversation_log.get("chat", []):
+    for chat in chat_history:
         message_context = []
         message_attached_files = ""
 
         generated_assets = {}
 
-        chat_message = chat.get("message")
-        role = "user" if chat["by"] == "you" else "assistant"
+        chat_message = chat.message
+        role = "user" if chat.by == "you" else "assistant"
 
         # Legacy code to handle excalidraw diagrams prior to Dec 2024
-        if chat["by"] == "khoj" and "excalidraw" in chat["intent"].get("type", ""):
-            chat_message = chat["intent"].get("inferred-queries")[0]
+        if chat.by == "khoj" and "excalidraw" in chat.intent.type or "":
+            chat_message = (chat.intent.inferred_queries or [])[0]
 
-        if chat.get("queryFiles"):
-            raw_query_files = chat.get("queryFiles")
+        if chat.queryFiles:
+            raw_query_files = chat.queryFiles
             query_files_dict = dict()
             for file in raw_query_files:
                 query_files_dict[file["name"]] = file["content"]
@@ -551,24 +560,24 @@ def generate_chatml_messages_with_context(
             message_attached_files = gather_raw_query_files(query_files_dict)
             chatml_messages.append(ChatMessage(content=message_attached_files, role=role))
 
-        if not is_none_or_empty(chat.get("onlineContext")):
+        if not is_none_or_empty(chat.onlineContext):
             message_context += [
                 {
                     "type": "text",
-                    "text": f"{prompts.online_search_conversation.format(online_results=chat.get('onlineContext'))}",
+                    "text": f"{prompts.online_search_conversation.format(online_results=chat.onlineContext)}",
                 }
             ]
 
-        if not is_none_or_empty(chat.get("codeContext")):
+        if not is_none_or_empty(chat.codeContext):
             message_context += [
                 {
                     "type": "text",
-                    "text": f"{prompts.code_executed_context.format(code_results=chat.get('codeContext'))}",
+                    "text": f"{prompts.code_executed_context.format(code_results=chat.codeContext)}",
                 }
             ]
 
-        if not is_none_or_empty(chat.get("operatorContext")):
-            operator_context = chat.get("operatorContext")
+        if not is_none_or_empty(chat.operatorContext):
+            operator_context = chat.operatorContext
             operator_content = "\n\n".join([f'## Task: {oc["query"]}\n{oc["response"]}\n' for oc in operator_context])
             message_context += [
                 {
@@ -577,13 +586,9 @@ def generate_chatml_messages_with_context(
                 }
             ]
 
-        if not is_none_or_empty(chat.get("context")):
+        if not is_none_or_empty(chat.context):
             references = "\n\n".join(
-                {
-                    f"# File: {item['file']}\n## {item['compiled']}\n"
-                    for item in chat.get("context") or []
-                    if isinstance(item, dict)
-                }
+                {f"# File: {item.file}\n## {item.compiled}\n" for item in chat.context or [] if isinstance(item, dict)}
             )
             message_context += [{"type": "text", "text": f"{prompts.notes_conversation.format(references=references)}"}]
 
@@ -591,14 +596,14 @@ def generate_chatml_messages_with_context(
             reconstructed_context_message = ChatMessage(content=message_context, role="user")
             chatml_messages.insert(0, reconstructed_context_message)
 
-        if not is_none_or_empty(chat.get("images")) and role == "assistant":
+        if not is_none_or_empty(chat.images) and role == "assistant":
             generated_assets["image"] = {
-                "query": chat.get("intent", {}).get("inferred-queries", [user_message])[0],
+                "query": (chat.intent.inferred_queries or [user_message])[0],
             }
 
-        if not is_none_or_empty(chat.get("mermaidjsDiagram")) and role == "assistant":
+        if not is_none_or_empty(chat.mermaidjsDiagram) and role == "assistant":
             generated_assets["diagram"] = {
-                "query": chat.get("intent", {}).get("inferred-queries", [user_message])[0],
+                "query": (chat.intent.inferred_queries or [user_message])[0],
             }
 
         if not is_none_or_empty(generated_assets):
@@ -610,7 +615,7 @@ def generate_chatml_messages_with_context(
             )
 
         message_content = construct_structured_message(
-            chat_message, chat.get("images") if role == "user" else [], model_type, vision_enabled
+            chat_message, chat.images if role == "user" else [], model_type, vision_enabled
         )
 
         reconstructed_message = ChatMessage(content=message_content, role=role)