When running code, strip any code delimiters. Disable application json type specification in Gemini request.

This commit is contained in:
sabaimran 2024-11-01 13:47:39 -07:00
parent 8fd2fe162f
commit baa939f4ce
3 changed files with 18 additions and 5 deletions

View file

@ -116,8 +116,10 @@ def gemini_send_message_to_model(
messages, system_prompt = format_messages_for_gemini(messages) messages, system_prompt = format_messages_for_gemini(messages)
model_kwargs = {} model_kwargs = {}
if response_type == "json_object":
model_kwargs["response_mime_type"] = "application/json" # Sometimes, this causes unwanted behavior and terminates response early. Disable for now while it's flaky.
# if response_type == "json_object":
# model_kwargs["response_mime_type"] = "application/json"
# Get Response from Gemini # Get Response from Gemini
return gemini_completion_with_backoff( return gemini_completion_with_backoff(

View file

@ -447,6 +447,11 @@ def clean_json(response: str):
return response.strip().replace("\n", "").removeprefix("```json").removesuffix("```") 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): def defilter_query(query: str):
"""Remove any query filters in query""" """Remove any query filters in query"""
defiltered_query = query defiltered_query = query

View file

@ -12,6 +12,7 @@ from khoj.database.models import Agent, KhojUser
from khoj.processor.conversation import prompts from khoj.processor.conversation import prompts
from khoj.processor.conversation.utils import ( from khoj.processor.conversation.utils import (
ChatEvent, ChatEvent,
clean_code_python,
clean_json, clean_json,
construct_chat_history, construct_chat_history,
) )
@ -126,13 +127,18 @@ async def execute_sandboxed_python(code: str, sandbox_url: str = SANDBOX_URL) ->
Returns the result of the code execution as a dictionary. Returns the result of the code execution as a dictionary.
""" """
headers = {"Content-Type": "application/json"} headers = {"Content-Type": "application/json"}
data = {"code": code} cleaned_code = clean_code_python(code)
data = {"code": cleaned_code}
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.post(sandbox_url, json=data, headers=headers) as response: async with session.post(sandbox_url, json=data, headers=headers) as response:
if response.status == 200: if response.status == 200:
result: dict[str, Any] = await response.json() result: dict[str, Any] = await response.json()
result["code"] = code result["code"] = cleaned_code
return result return result
else: else:
return {"code": code, "success": False, "std_err": f"Failed to execute code with {response.status}"} return {
"code": cleaned_code,
"success": False,
"std_err": f"Failed to execute code with {response.status}",
}