khoj 1.27.2.dev12__py3-none-any.whl → 1.28.1__py3-none-any.whl

khoj/processor/conversation/utils.py

@@ -1,17 +1,23 @@
 import base64
+import json
 import logging
 import math
 import mimetypes
+import os
 import queue
+import uuid
 from dataclasses import dataclass
 from datetime import datetime
+from enum import Enum
 from io import BytesIO
 from time import perf_counter
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import PIL.Image
 import requests
 import tiktoken
+import yaml
+from git import Repo
 from langchain.schema import ChatMessage
 from llama_cpp.llama import Llama
 from transformers import AutoTokenizer
@@ -20,8 +26,17 @@ from khoj.database.adapters import ConversationAdapters
 from khoj.database.models import ChatModelOptions, ClientApplication, 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
+from khoj.search_filter.date_filter import DateFilter
+from khoj.search_filter.file_filter import FileFilter
+from khoj.search_filter.word_filter import WordFilter
 from khoj.utils import state
-from khoj.utils.helpers import is_none_or_empty, merge_dicts
+from khoj.utils.helpers import (
+    ConversationCommand,
+    in_debug_mode,
+    is_none_or_empty,
+    merge_dicts,
+)
 
 logger = logging.getLogger(__name__)
 model_to_prompt_size = {
@@ -82,8 +97,110 @@ class ThreadedGenerator:
         self.queue.put(StopIteration)
 
 
+class InformationCollectionIteration:
+    def __init__(
+        self,
+        tool: str,
+        query: str,
+        context: list = None,
+        onlineContext: dict = None,
+        codeContext: dict = None,
+        summarizedResult: str = None,
+    ):
+        self.tool = tool
+        self.query = query
+        self.context = context
+        self.onlineContext = onlineContext
+        self.codeContext = codeContext
+        self.summarizedResult = summarizedResult
+
+
+def construct_iteration_history(
+    previous_iterations: List[InformationCollectionIteration], previous_iteration_prompt: str
+) -> str:
+    previous_iterations_history = ""
+    for idx, iteration in enumerate(previous_iterations):
+        iteration_data = previous_iteration_prompt.format(
+            tool=iteration.tool,
+            query=iteration.query,
+            result=iteration.summarizedResult,
+            index=idx + 1,
+        )
+
+        previous_iterations_history += iteration_data
+    return previous_iterations_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"]:
+            chat_history += f"User: {chat['intent']['query']}\n"
+
+            if chat["intent"].get("inferred-queries"):
+                chat_history += f'Khoj: {{"queries": {chat["intent"].get("inferred-queries")}}}\n'
+
+            chat_history += f"{agent_name}: {chat['message']}\n\n"
+        elif chat["by"] == "khoj" and ("text-to-image" in chat["intent"].get("type")):
+            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"
+    return chat_history
+
+
+def construct_tool_chat_history(
+    previous_iterations: List[InformationCollectionIteration], tool: ConversationCommand = None
+) -> Dict[str, list]:
+    chat_history: list = []
+    inferred_query_extractor: Callable[[InformationCollectionIteration], List[str]] = lambda x: []
+    if tool == ConversationCommand.Notes:
+        inferred_query_extractor = (
+            lambda iteration: [c["query"] for c in iteration.context] if iteration.context else []
+        )
+    elif tool == ConversationCommand.Online:
+        inferred_query_extractor = (
+            lambda iteration: list(iteration.onlineContext.keys()) if iteration.onlineContext else []
+        )
+    elif tool == ConversationCommand.Code:
+        inferred_query_extractor = lambda iteration: list(iteration.codeContext.keys()) if iteration.codeContext else []
+    for iteration in previous_iterations:
+        chat_history += [
+            {
+                "by": "you",
+                "message": iteration.query,
+            },
+            {
+                "by": "khoj",
+                "intent": {
+                    "type": "remember",
+                    "inferred-queries": inferred_query_extractor(iteration),
+                    "query": iteration.query,
+                },
+                "message": iteration.summarizedResult,
+            },
+        ]
+
+    return {"chat": chat_history}
+
+
+class ChatEvent(Enum):
+    START_LLM_RESPONSE = "start_llm_response"
+    END_LLM_RESPONSE = "end_llm_response"
+    MESSAGE = "message"
+    REFERENCES = "references"
+    STATUS = "status"
+    METADATA = "metadata"
+
+
 def message_to_log(
-    user_message, chat_response, user_message_metadata={}, khoj_message_metadata={}, conversation_log=[]
+    user_message,
+    chat_response,
+    user_message_metadata={},
+    khoj_message_metadata={},
+    conversation_log=[],
+    train_of_thought=[],
 ):
     """Create json logs from messages, metadata for conversation log"""
     default_khoj_message_metadata = {
@@ -111,28 +228,37 @@ def save_to_conversation_log(
     user_message_time: str = None,
     compiled_references: List[Dict[str, Any]] = [],
     online_results: Dict[str, Any] = {},
+    code_results: Dict[str, Any] = {},
     inferred_queries: List[str] = [],
     intent_type: str = "remember",
     client_application: ClientApplication = None,
     conversation_id: str = None,
     automation_id: str = None,
     query_images: List[str] = None,
+    tracer: Dict[str, Any] = {},
+    train_of_thought: List[Any] = [],
 ):
     user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    turn_id = tracer.get("mid") or str(uuid.uuid4())
     updated_conversation = message_to_log(
         user_message=q,
         chat_response=chat_response,
         user_message_metadata={
             "created": user_message_time,
             "images": query_images,
+            "turnId": turn_id,
         },
         khoj_message_metadata={
             "context": compiled_references,
             "intent": {"inferred-queries": inferred_queries, "type": intent_type},
             "onlineContext": online_results,
+            "codeContext": code_results,
             "automationId": automation_id,
+            "trainOfThought": train_of_thought,
+            "turnId": turn_id,
         },
         conversation_log=meta_log.get("chat", []),
+        train_of_thought=train_of_thought,
     )
     ConversationAdapters.save_conversation(
         user,
@@ -142,6 +268,9 @@ def save_to_conversation_log(
         user_message=q,
     )
 
+    if in_debug_mode() or state.verbose > 1:
+        merge_message_into_conversation_trace(q, chat_response, tracer)
+
     logger.info(
         f"""
 Saved Conversation Turn
@@ -323,9 +452,23 @@ def reciprocal_conversation_to_chatml(message_pair):
     return [ChatMessage(content=message, role=role) for message, role in zip(message_pair, ["user", "assistant"])]
 
 
-def remove_json_codeblock(response: str):
-    """Remove any markdown json codeblock formatting if present. Useful for non schema enforceable models"""
-    return response.removeprefix("```json").removesuffix("```")
+def clean_json(response: str):
+    """Remove any markdown json codeblock and newline formatting if present. Useful for non schema enforceable models"""
+    return response.strip().replace("\n", "").removeprefix("```json").removesuffix("```")
+
+
+def clean_code_python(code: str):
+    """Remove any markdown codeblock and newline formatting if present. Useful for non schema enforceable models"""
+    return code.strip().removeprefix("```python").removesuffix("```")
+
+
+def defilter_query(query: str):
+    """Remove any query filters in query"""
+    defiltered_query = query
+    filters: List[BaseFilter] = [WordFilter(), FileFilter(), DateFilter()]
+    for filter in filters:
+        defiltered_query = filter.defilter(defiltered_query)
+    return defiltered_query
 
 
 @dataclass
@@ -354,3 +497,160 @@ def get_image_from_url(image_url: str, type="pil"):
     except requests.exceptions.RequestException as e:
         logger.error(f"Failed to get image from URL {image_url}: {e}")
         return ImageWithType(content=None, type=None)
+
+
+def commit_conversation_trace(
+    session: list[ChatMessage],
+    response: str | list[dict],
+    tracer: dict,
+    system_message: str | list[dict] = "",
+    repo_path: str = "/tmp/promptrace",
+) -> str:
+    """
+    Save trace of conversation step using git. Useful to visualize, compare and debug traces.
+    Returns the path to the repository.
+    """
+    # Serialize session, system message and response to yaml
+    system_message_yaml = json.dumps(system_message, ensure_ascii=False, sort_keys=False)
+    response_yaml = json.dumps(response, ensure_ascii=False, sort_keys=False)
+    formatted_session = [{"role": message.role, "content": message.content} for message in session]
+    session_yaml = json.dumps(formatted_session, ensure_ascii=False, sort_keys=False)
+    query = (
+        json.dumps(session[-1].content, ensure_ascii=False, sort_keys=False).strip().removeprefix("'").removesuffix("'")
+    )  # Extract serialized query from chat session
+
+    # Extract chat metadata for session
+    uid, cid, mid = tracer.get("uid", "main"), tracer.get("cid", "main"), tracer.get("mid")
+
+    # Infer repository path from environment variable or provided path
+    repo_path = os.getenv("PROMPTRACE_DIR", repo_path)
+
+    try:
+        # Prepare git repository
+        os.makedirs(repo_path, exist_ok=True)
+        repo = Repo.init(repo_path)
+
+        # Remove post-commit hook if it exists
+        hooks_dir = os.path.join(repo_path, ".git", "hooks")
+        post_commit_hook = os.path.join(hooks_dir, "post-commit")
+        if os.path.exists(post_commit_hook):
+            os.remove(post_commit_hook)
+
+        # Configure git user if not set
+        if not repo.config_reader().has_option("user", "email"):
+            repo.config_writer().set_value("user", "name", "Prompt Tracer").release()
+            repo.config_writer().set_value("user", "email", "promptracer@khoj.dev").release()
+
+        # Create an initial commit if the repository is newly created
+        if not repo.head.is_valid():
+            repo.index.commit("And then there was a trace")
+
+        # Check out the initial commit
+        initial_commit = repo.commit("HEAD~0")
+        repo.head.reference = initial_commit
+        repo.head.reset(index=True, working_tree=True)
+
+        # Create or switch to user branch from initial commit
+        user_branch = f"u_{uid}"
+        if user_branch not in repo.branches:
+            repo.create_head(user_branch)
+        repo.heads[user_branch].checkout()
+
+        # Create or switch to conversation branch from user branch
+        conv_branch = f"c_{cid}"
+        if conv_branch not in repo.branches:
+            repo.create_head(conv_branch)
+        repo.heads[conv_branch].checkout()
+
+        # Create or switch to message branch from conversation branch
+        msg_branch = f"m_{mid}" if mid else None
+        if msg_branch and msg_branch not in repo.branches:
+            repo.create_head(msg_branch)
+        if msg_branch:
+            repo.heads[msg_branch].checkout()
+
+        # Include file with content to commit
+        files_to_commit = {"query": session_yaml, "response": response_yaml, "system_prompt": system_message_yaml}
+
+        # Write files and stage them
+        for filename, content in files_to_commit.items():
+            file_path = os.path.join(repo_path, filename)
+            # Unescape special characters in content for better readability
+            content = content.strip().replace("\\n", "\n").replace("\\t", "\t")
+            with open(file_path, "w", encoding="utf-8") as f:
+                f.write(content)
+            repo.index.add([filename])
+
+        # Create commit
+        metadata_yaml = yaml.dump(tracer, allow_unicode=True, sort_keys=False, default_flow_style=False)
+        commit_message = f"""
+{query[:250]}
+
+Response:
+---
+{response[:500]}...
+
+Metadata
+---
+{metadata_yaml}
+""".strip()
+
+        repo.index.commit(commit_message)
+
+        logger.debug(f"Saved conversation trace to repo at {repo_path}")
+        return repo_path
+    except Exception as e:
+        logger.error(f"Failed to add conversation trace to repo: {str(e)}", exc_info=True)
+        return None
+
+
+def merge_message_into_conversation_trace(query: str, response: str, tracer: dict, repo_path="/tmp/promptrace") -> bool:
+    """
+    Merge the message branch into its parent conversation branch.
+
+    Args:
+        query: User query
+        response: Assistant response
+        tracer: Dictionary containing uid, cid and mid
+        repo_path: Path to the git repository
+
+    Returns:
+        bool: True if merge was successful, False otherwise
+    """
+    try:
+        # Extract branch names
+        msg_branch = f"m_{tracer['mid']}"
+        conv_branch = f"c_{tracer['cid']}"
+
+        # Infer repository path from environment variable or provided path
+        repo_path = os.getenv("PROMPTRACE_DIR", repo_path)
+        repo = Repo(repo_path)
+
+        # Checkout conversation branch
+        repo.heads[conv_branch].checkout()
+
+        # Create commit message
+        metadata_yaml = yaml.dump(tracer, allow_unicode=True, sort_keys=False, default_flow_style=False)
+        commit_message = f"""
+{query[:250]}
+
+Response:
+---
+{response[:500]}...
+
+Metadata
+---
+{metadata_yaml}
+""".strip()
+
+        # Merge message branch into conversation branch
+        repo.git.merge(msg_branch, no_ff=True, m=commit_message)
+
+        # Delete message branch after merge
+        repo.delete_head(msg_branch, force=True)
+
+        logger.debug(f"Successfully merged {msg_branch} into {conv_branch}")
+        return True
+    except Exception as e:
+        logger.error(f"Failed to merge message {msg_branch} into conversation {conv_branch}: {str(e)}", exc_info=True)
+        return False

khoj/processor/embeddings.py

@@ -13,7 +13,7 @@ from tenacity import (
 )
 from torch import nn
 
-from khoj.utils.helpers import get_device, merge_dicts, timer
+from khoj.utils.helpers import fix_json_dict, get_device, merge_dicts, timer
 from khoj.utils.rawconfig import SearchResponse
 
 logger = logging.getLogger(__name__)
@@ -31,9 +31,9 @@ class EmbeddingsModel:
     ):
         default_query_encode_kwargs = {"show_progress_bar": False, "normalize_embeddings": True}
         default_docs_encode_kwargs = {"show_progress_bar": True, "normalize_embeddings": True}
-        self.query_encode_kwargs = merge_dicts(query_encode_kwargs, default_query_encode_kwargs)
-        self.docs_encode_kwargs = merge_dicts(docs_encode_kwargs, default_docs_encode_kwargs)
-        self.model_kwargs = merge_dicts(model_kwargs, {"device": get_device()})
+        self.query_encode_kwargs = merge_dicts(fix_json_dict(query_encode_kwargs), default_query_encode_kwargs)
+        self.docs_encode_kwargs = merge_dicts(fix_json_dict(docs_encode_kwargs), default_docs_encode_kwargs)
+        self.model_kwargs = merge_dicts(fix_json_dict(model_kwargs), {"device": get_device()})
         self.model_name = model_name
         self.inference_endpoint = embeddings_inference_endpoint
         self.api_key = embeddings_inference_endpoint_api_key

khoj/processor/image/generate.py

@@ -28,6 +28,7 @@ async def text_to_image(
     send_status_func: Optional[Callable] = None,
     query_images: Optional[List[str]] = None,
     agent: Agent = None,
+    tracer: dict = {},
 ):
     status_code = 200
     image = None
@@ -68,6 +69,7 @@ async def text_to_image(
         query_images=query_images,
         user=user,
         agent=agent,
+        tracer=tracer,
     )
 
     if send_status_func:

khoj/processor/tools/online_search.py

@@ -4,7 +4,7 @@ import logging
 import os
 import urllib.parse
 from collections import defaultdict
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import aiohttp
 from bs4 import BeautifulSoup
@@ -52,7 +52,9 @@ OLOSTEP_QUERY_PARAMS = {
     "expandMarkdown": "True",
     "expandHtml": "False",
 }
-MAX_WEBPAGES_TO_READ = 1
+
+DEFAULT_MAX_WEBPAGES_TO_READ = 1
+MAX_WEBPAGES_TO_INFER = 10
 
 
 async def search_online(
@@ -62,8 +64,10 @@ async def search_online(
     user: KhojUser,
     send_status_func: Optional[Callable] = None,
     custom_filters: List[str] = [],
+    max_webpages_to_read: int = DEFAULT_MAX_WEBPAGES_TO_READ,
     query_images: List[str] = None,
     agent: Agent = None,
+    tracer: dict = {},
 ):
     query += " ".join(custom_filters)
     if not is_internet_connected():
@@ -73,7 +77,7 @@ async def search_online(
 
     # Breakdown the query into subqueries to get the correct answer
     subqueries = await generate_online_subqueries(
-        query, conversation_history, location, user, query_images=query_images, agent=agent
+        query, conversation_history, location, user, query_images=query_images, agent=agent, tracer=tracer
     )
     response_dict = {}
 
@@ -96,7 +100,7 @@ async def search_online(
     for subquery in response_dict:
         if "answerBox" in response_dict[subquery]:
             continue
-        for organic in response_dict[subquery].get("organic", [])[:MAX_WEBPAGES_TO_READ]:
+        for organic in response_dict[subquery].get("organic", [])[:max_webpages_to_read]:
             link = organic.get("link")
             if link in webpages:
                 webpages[link]["queries"].add(subquery)
@@ -111,7 +115,7 @@ async def search_online(
             async for event in send_status_func(f"**Reading web pages**: {webpage_links_str}"):
                 yield {ChatEvent.STATUS: event}
     tasks = [
-        read_webpage_and_extract_content(data["queries"], link, data["content"], user=user, agent=agent)
+        read_webpage_and_extract_content(data["queries"], link, data["content"], user=user, agent=agent, tracer=tracer)
         for link, data in webpages.items()
     ]
     results = await asyncio.gather(*tasks)
@@ -153,20 +157,24 @@ async def read_webpages(
     send_status_func: Optional[Callable] = None,
     query_images: List[str] = None,
     agent: Agent = None,
+    tracer: dict = {},
+    max_webpages_to_read: int = DEFAULT_MAX_WEBPAGES_TO_READ,
 ):
     "Infer web pages to read from the query and extract relevant information from them"
     logger.info(f"Inferring web pages to read")
-    if send_status_func:
-        async for event in send_status_func(f"**Inferring web pages to read**"):
-            yield {ChatEvent.STATUS: event}
-    urls = await infer_webpage_urls(query, conversation_history, location, user, query_images)
+    urls = await infer_webpage_urls(
+        query, conversation_history, location, user, query_images, agent=agent, tracer=tracer
+    )
+
+    # Get the top 10 web pages to read
+    urls = urls[:max_webpages_to_read]
 
     logger.info(f"Reading web pages at: {urls}")
     if send_status_func:
         webpage_links_str = "\n- " + "\n- ".join(list(urls))
         async for event in send_status_func(f"**Reading web pages**: {webpage_links_str}"):
             yield {ChatEvent.STATUS: event}
-    tasks = [read_webpage_and_extract_content({query}, url, user=user, agent=agent) for url in urls]
+    tasks = [read_webpage_and_extract_content({query}, url, user=user, agent=agent, tracer=tracer) for url in urls]
     results = await asyncio.gather(*tasks)
 
     response: Dict[str, Dict] = defaultdict(dict)
@@ -192,7 +200,12 @@ async def read_webpage(
 
 
 async def read_webpage_and_extract_content(
-    subqueries: set[str], url: str, content: str = None, user: KhojUser = None, agent: Agent = None
+    subqueries: set[str],
+    url: str,
+    content: str = None,
+    user: KhojUser = None,
+    agent: Agent = None,
+    tracer: dict = {},
 ) -> Tuple[set[str], str, Union[None, str]]:
     # Select the web scrapers to use for reading the web page
     web_scrapers = await ConversationAdapters.aget_enabled_webscrapers()
@@ -214,7 +227,9 @@ async def read_webpage_and_extract_content(
             # Extract relevant information from the web page
             if is_none_or_empty(extracted_info):
                 with timer(f"Extracting relevant information from web page at '{url}' took", logger):
-                    extracted_info = await extract_relevant_info(subqueries, content, user=user, agent=agent)
+                    extracted_info = await extract_relevant_info(
+                        subqueries, content, user=user, agent=agent, tracer=tracer
+                    )
 
             # If we successfully extracted information, break the loop
             if not is_none_or_empty(extracted_info):

khoj/processor/tools/run_code.py

@@ -0,0 +1,144 @@
+import asyncio
+import datetime
+import json
+import logging
+import os
+from typing import Any, Callable, List, Optional
+
+import aiohttp
+
+from khoj.database.adapters import ais_user_subscribed
+from khoj.database.models import Agent, KhojUser
+from khoj.processor.conversation import prompts
+from khoj.processor.conversation.utils import (
+    ChatEvent,
+    clean_code_python,
+    clean_json,
+    construct_chat_history,
+)
+from khoj.routers.helpers import send_message_to_model_wrapper
+from khoj.utils.helpers import timer
+from khoj.utils.rawconfig import LocationData
+
+logger = logging.getLogger(__name__)
+
+
+SANDBOX_URL = os.getenv("KHOJ_TERRARIUM_URL", "http://localhost:8080")
+
+
+async def run_code(
+    query: str,
+    conversation_history: dict,
+    context: str,
+    location_data: LocationData,
+    user: KhojUser,
+    send_status_func: Optional[Callable] = None,
+    query_images: List[str] = None,
+    agent: Agent = None,
+    sandbox_url: str = SANDBOX_URL,
+    tracer: dict = {},
+):
+    # Generate Code
+    if send_status_func:
+        async for event in send_status_func(f"**Generate code snippets** for {query}"):
+            yield {ChatEvent.STATUS: event}
+    try:
+        with timer("Chat actor: Generate programs to execute", logger):
+            codes = await generate_python_code(
+                query,
+                conversation_history,
+                context,
+                location_data,
+                user,
+                query_images,
+                agent,
+                tracer,
+            )
+    except Exception as e:
+        raise ValueError(f"Failed to generate code for {query} with error: {e}")
+
+    # Run Code
+    if send_status_func:
+        async for event in send_status_func(f"**Running {len(codes)} code snippets**"):
+            yield {ChatEvent.STATUS: event}
+    try:
+        tasks = [execute_sandboxed_python(code, sandbox_url) for code in codes]
+        with timer("Chat actor: Execute generated programs", logger):
+            results = await asyncio.gather(*tasks)
+        for result in results:
+            code = result.pop("code")
+            logger.info(f"Executed Code:\n--@@--\n{code}\n--@@--Result:\n--@@--\n{result}\n--@@--")
+            yield {query: {"code": code, "results": result}}
+    except Exception as e:
+        raise ValueError(f"Failed to run code for {query} with error: {e}")
+
+
+async def generate_python_code(
+    q: str,
+    conversation_history: dict,
+    context: str,
+    location_data: LocationData,
+    user: KhojUser,
+    query_images: List[str] = None,
+    agent: Agent = None,
+    tracer: dict = {},
+) -> List[str]:
+    location = f"{location_data}" if location_data else "Unknown"
+    username = prompts.user_name.format(name=user.get_full_name()) if user.get_full_name() else ""
+    subscribed = await ais_user_subscribed(user)
+    chat_history = construct_chat_history(conversation_history)
+
+    utc_date = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d")
+    personality_context = (
+        prompts.personality_context.format(personality=agent.personality) if agent and agent.personality else ""
+    )
+
+    code_generation_prompt = prompts.python_code_generation_prompt.format(
+        current_date=utc_date,
+        query=q,
+        chat_history=chat_history,
+        context=context,
+        location=location,
+        username=username,
+        personality_context=personality_context,
+    )
+
+    response = await send_message_to_model_wrapper(
+        code_generation_prompt,
+        query_images=query_images,
+        response_type="json_object",
+        user=user,
+        tracer=tracer,
+    )
+
+    # Validate that the response is a non-empty, JSON-serializable list
+    response = clean_json(response)
+    response = json.loads(response)
+    codes = [code.strip() for code in response["codes"] if code.strip()]
+
+    if not isinstance(codes, list) or not codes or len(codes) == 0:
+        raise ValueError
+    return codes
+
+
+async def execute_sandboxed_python(code: str, sandbox_url: str = SANDBOX_URL) -> dict[str, Any]:
+    """
+    Takes code to run as a string and calls the terrarium API to execute it.
+    Returns the result of the code execution as a dictionary.
+    """
+    headers = {"Content-Type": "application/json"}
+    cleaned_code = clean_code_python(code)
+    data = {"code": cleaned_code}
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(sandbox_url, json=data, headers=headers) as response:
+            if response.status == 200:
+                result: dict[str, Any] = await response.json()
+                result["code"] = cleaned_code
+                return result
+            else:
+                return {
+                    "code": cleaned_code,
+                    "success": False,
+                    "std_err": f"Failed to execute code with {response.status}",
+                }