From a83c36fa0590f94570ea4f74277f6092846f2591 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Fri, 23 May 2025 02:36:58 -0700
Subject: [PATCH] Validate operator, research, context.query fields of
 ChatMessage

- Track operator, research context in ChatMessage
- Track query field in (document) context field of ChatMessage

This allows validating chat message before inserting into DB
---
 src/khoj/database/models/__init__.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py
index bd49aa8c..34538df0 100644
--- a/src/khoj/database/models/__init__.py
+++ b/src/khoj/database/models/__init__.py
@@ -23,6 +23,7 @@ logger = logging.getLogger(__name__)
 class Context(PydanticBaseModel):
     compiled: str
     file: str
+    query: str
 
 
 class CodeContextFile(PydanticBaseModel):
@@ -105,6 +106,8 @@ class ChatMessage(PydanticBaseModel):
     context: List[Context] = []
     onlineContext: Dict[str, OnlineContext] = {}
     codeContext: Dict[str, CodeContextData] = {}
+    researchContext: Optional[List] = None
+    operatorContext: Optional[Dict[str, str]] = None
     created: str
     images: Optional[List[str]] = None
     queryFiles: Optional[List[Dict]] = None