From 64886cd0dd17045e3309cb872d1ab5430bec4d07 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Mon, 9 Jun 2025 23:42:56 -0700
Subject: [PATCH] Fix storing code results on server and rendering them on web
 app

- Fix code context data type for validation on server. This would
  prevent the chat message from being written to history
- Handle null code results on web app
---
 src/interface/web/app/common/chatFunctions.ts                | 2 +-
 src/interface/web/app/components/chatMessage/chatMessage.tsx | 4 ++--
 src/khoj/database/models/__init__.py                         | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/interface/web/app/common/chatFunctions.ts b/src/interface/web/app/common/chatFunctions.ts
index 4823659a..92f869e3 100644
--- a/src/interface/web/app/common/chatFunctions.ts
+++ b/src/interface/web/app/common/chatFunctions.ts
@@ -205,7 +205,7 @@ export function renderCodeGenImageInline(message: string, codeContext: CodeConte
     if (!codeContext) return message;
 
     Object.values(codeContext).forEach((contextData) => {
-        contextData.results.output_files?.forEach((file) => {
+        contextData.results?.output_files?.forEach((file) => {
             const regex = new RegExp(`!?\\[.*?\\]\\(.*${file.filename}\\)`, "g");
             if (file.filename.match(/\.(png|jpg|jpeg)$/i)) {
                 const replacement = `![${file.filename}](data:image/${file.filename.split(".").pop()};base64,${file.b64_data})`;
diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx
index e25dc814..890aec9b 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -129,7 +129,7 @@ export interface CodeContextData {
         output_files: CodeContextFile[];
         std_out: string;
         std_err: string;
-        code_runtime: number;
+        code_runtime?: number;
     };
 }
 
@@ -454,7 +454,7 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
         // Add code context files to the message
         if (props.chatMessage.codeContext) {
             Object.entries(props.chatMessage.codeContext).forEach(([key, value]) => {
-                value.results.output_files?.forEach((file) => {
+                value.results?.output_files?.forEach((file) => {
                     if (file.filename.endsWith(".png") || file.filename.endsWith(".jpg")) {
                         // Don't add the image again if it's already in the message!
                         if (!message.includes(`![${file.filename}](`)) {
diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py
index 8a5beaf0..c4ccadbf 100644
--- a/src/khoj/database/models/__init__.py
+++ b/src/khoj/database/models/__init__.py
@@ -36,12 +36,12 @@ class CodeContextResult(PydanticBaseModel):
     output_files: List[CodeContextFile]
     std_out: str
     std_err: str
-    code_runtime: int
+    code_runtime: Optional[int] = None
 
 
 class CodeContextData(PydanticBaseModel):
     code: str
-    result: Optional[CodeContextResult] = None
+    results: Optional[CodeContextResult] = None
 
 
 class WebPage(PydanticBaseModel):