sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-11-23 15:38:55 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Cut binary, long output files from code result for context efficiency

Browse source 
Removing binary data and truncating large data in output files
generated by code runs should improve speed and cost of research mode
runs with large or binary output files.

Previously binary data in code results was passed around in iteration
context during research mode. This made the context inefficient
because models have limited efficiency and reasoning capabilities over
b64 encoded image (and other binary) data and would hit context limits
leading to unnecessary truncation of other useful context

Also remove image data when logging output of code execution
This commit is contained in:
Debanjum 2024-11-12 03:19:58 -08:00
parent de34cc3987
commit 153ae8bea9
2 changed files with 31 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
src/khoj/processor/tools/run_code.py
View file

@ -1,4 +1,5 @@
import base64
import copy
import datetime
import json
import logging
@ -88,7 +89,8 @@ async def run_code(
with timer("Chat actor: Execute generated program", logger, log_level=logging.INFO):
result = await execute_sandboxed_python(generated_code.code, input_data, sandbox_url)
code = result.pop("code")
logger.info(f"Executed Code:\n--@@--\n{code}\n--@@--Result:\n--@@--\n{result}\n--@@--")
cleaned_result = truncate_code_context({"cleaned": {"results": result}})["cleaned"]["results"]
logger.info(f"Executed Code\n----\n{code}\n----\nResult\n----\n{cleaned_result}\n----")
yield {query: {"code": code, "results": result}}
except Exception as e:
raise ValueError(f"Failed to run code for {query} with error: {e}")
@ -163,7 +165,8 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
result: dict[str, Any] = await response.json()
result["code"] = cleaned_code
# Store decoded output files
for output_file in result.get("output_files", []):
result["output_files"] = result.get("output_files", [])
for output_file in result["output_files"]:
# Decode text files as UTF-8
if mimetypes.guess_type(output_file["filename"])[0].startswith("text/") or Path(
output_file["filename"]
@ -175,4 +178,28 @@ async def execute_sandboxed_python(code: str, input_data: list[dict], sandbox_ur
"code": cleaned_code,
"success": False,
"std_err": f"Failed to execute code with {response.status}",
"output_files": [],
}
def truncate_code_context(original_code_results: dict[str, Any], max_chars=10000) -> dict[str, Any]:
"""
Truncate large output files and drop image file data from code results.
"""
# Create a deep copy of the code results to avoid modifying the original data
code_results = copy.deepcopy(original_code_results)
for code_result in code_results.values():
for idx, output_file in enumerate(code_result["results"]["output_files"]):
# Drop image files from code results
if Path(output_file["filename"]).suffix in {".png", ".jpg", ".jpeg", ".webp"}:
code_result["results"]["output_files"][idx] = {
"filename": output_file["filename"],
"b64_data": "[placeholder for generated image data for brevity]",
}
# Truncate large output files
elif len(output_file["b64_data"]) > max_chars:
code_result["results"]["output_files"][idx] = {
"filename": output_file["filename"],
"b64_data": output_file["b64_data"][:max_chars] + "...",
}
return code_results

4
src/khoj/routers/research.py
View file

@ -16,7 +16,7 @@ from khoj.processor.conversation.utils import (
construct_tool_chat_history,
)
from khoj.processor.tools.online_search import read_webpages, search_online
from khoj.processor.tools.run_code import run_code
from khoj.processor.tools.run_code import run_code, truncate_code_context
from khoj.routers.api import extract_references_and_questions
from khoj.routers.helpers import (
ChatEvent,
@ -348,7 +348,7 @@ async def execute_information_collection(
if online_results:
results_data += f"\n<online_results>\n{yaml.dump(online_results, allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</online_results>"
if code_results:
results_data += f"\n<code_results>\n{yaml.dump(code_results, allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</code_results>"
results_data += f"\n<code_results>\n{yaml.dump(truncate_code_context(code_results), allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</code_results>"
if summarize_files:
results_data += f"\n<summarized_files>\n{yaml.dump(summarize_files, allow_unicode=True, sort_keys=False, default_flow_style=False)}\n</summarized_files>"
if this_iteration.warning: