skeleton

django/aiservice/agent/agent.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import sentry_sdk
+from ninja import NinjaAPI, Schema
+from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
+
+from aiservice.analytics.posthog import ph
+from aiservice.common_utils import validate_trace_id
+
+# from aiservice.env_specific import debug_log_sensitive_data
+from aiservice.llm import DEFAULT_AGENT_MODEL, PLANNING_AGENT_MODEL, calculate_llm_cost, call_llm
+from authapp.auth import AuthenticatedRequest
+from log_features.log_event import update_optimization_cost
+
+
+class AgentErrorResponseSchema(Schema):
+    error: str
+
+
+class ResponseSchema(Schema):
+    status: str
+    messages: list[dict[str, str]]
+
+
+agent_api = NinjaAPI(urls_namespace="call_agent")
+
+# Get the directory of the current file
+current_dir = Path(__file__).parent
+SYSTEM_PROMPT = (current_dir / "agent_prompt.md").read_text()
+PLANNING_PROMPT = (current_dir / "planning_prompt.md").read_text()
+
+
+class AgentRequestSchema(Schema):
+    trace_id: str
+    context_vars: dict[str, str]
+    available_tools: list[str]  # TODO: only include the tools in the prompt that are relevant to the task
+    messages: list[dict[str, str]]
+
+
+@agent_api.post(
+    "/", response={200: ResponseSchema, 400: AgentErrorResponseSchema, 500: AgentErrorResponseSchema}
+)
+async def adaptive_optimize(
+    request: AuthenticatedRequest, data: AgentRequestSchema
+) -> tuple[int, ResponseSchema | AgentErrorResponseSchema]:
+    # ph(request.user, "aiservice-adaptive_optimize-called")
+    trace_id = data.trace_id
+    if not validate_trace_id(trace_id):
+        return 400, AgentErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
+
+    should_plan = len(data.messages) == 1
+
+    system_message = ChatCompletionSystemMessageParam(role="system", content=PLANNING_PROMPT if should_plan else SYSTEM_PROMPT)
+    old_messages = [ChatCompletionUserMessageParam(role=message["role"], content=message["content"]) for message in data.messages if message["role"] in ["user", "assistant"]]
+    context_vars_message = ChatCompletionUserMessageParam(role="user", content=f"Context variables: {data.context_vars}")
+
+    messages = [system_message, *old_messages]
+
+    if not should_plan:
+        messages.append(context_vars_message)
+
+    print(messages)
+    try:
+        llm_model = PLANNING_AGENT_MODEL if should_plan else DEFAULT_AGENT_MODEL
+        output = await call_llm(llm=llm_model, messages=messages)
+        llm_cost = calculate_llm_cost(output.raw_response, llm_model)
+    except Exception as e:
+        logging.exception(f"Agent error {e}")  # noqa: G004, TRY401
+        sentry_sdk.capture_exception(e)
+        return 500, AgentErrorResponseSchema(error=str(e))
+    if output.usage is not None:
+        ph(
+            request.user,
+            "adaptive_optimize-usage",
+            properties={
+                "model": llm_model.name,
+                "usage": {"input_tokens": output.usage.input_tokens, "output_tokens": output.usage.output_tokens},
+            },
+        )
+
+    llm_res = output.content
+
+    print("-------------------------------")
+    print(llm_res)
+    print("-------------------------------")
+
+    await update_optimization_cost(trace_id=trace_id, cost=llm_cost, user_id=request.user)
+    return 200, ResponseSchema(
+        status="success",
+        messages=[*old_messages, {"role": "assistant", "content": llm_res}]
+    )

django/aiservice/agent/agent_prompt.md

@@ -0,0 +1,100 @@
+You are a highly skilled software engineer with extensive knowledge in many programming languages (specializing in Python), frameworks, design patterns, and best practices. Your primary role is to optimize the runtime and memory efficiency of the provided code through safe, meaningful and clean rewrites that would pass senior-level code review.
+You will be provided with some information about the code you will optimize.
+
+When writing behaviour or benchmarking tests, make sure to import the function from its module and don't just copy/past the code into the test code.
+If you are comparing the original and new code. Make sure to call replace_in_file tool to replace the original code with the new code.
+
+## TOOL USE
+You have access to a set of tools that help you execute your tasks effectively. **Every response you generate MUST end with one or more tool calls in the specified XML format.**
+
+### Tool Use Rules
+You can use multiple tools in the same response (**ONLY IF NEEDED**). If the tools are independent (i.e., the second tool does not depend on the output of the first), include each tool as a separate xml entry. Do not wait for the result of one tool to generate the input for another—treat them as parallel, independent calls.
+
+*CRITICAL* every response MUST include at least a single tool call.
+
+### Tool Use Formatting
+Tool use is formatted using XML-style tags. The tool name is enclosed in opening and closing tags, and each parameter is similarly enclosed within its own set of tags. Here's the structure:
+
+```xml
+<tool_name>
+<parameter1_name>value1</parameter1_name>
+<parameter2_name>value2</parameter2_name>
+...
+</tool_name>
+```
+
+## Tools
+
+### `replace_in_file`
+
+Request to replace sections of content in an existing file using SEARCH/REPLACE blocks that define exact changes to specific parts of the file. This tool should be used when you need to make targeted changes to specific parts of a file.
+Parameters:
+- path: (required) The path of the file to modify
+- diff: (required) One or more SEARCH/REPLACE blocks following this exact format:
+  ```
+  <<<<<<< SEARCH
+  [exact content to find]
+  =======
+  [new content to replace with]
+  >>>>>>> REPLACE
+  ```
+  Critical rules:
+  1. SEARCH content must match the associated file section to find EXACTLY:
+     * Match character-for-character including whitespace, indentation, line endings
+     * Include all comments, docstrings, etc.
+  2. SEARCH/REPLACE blocks will ONLY replace the first match occurrence.
+     * Including multiple unique SEARCH/REPLACE blocks if you need to make multiple changes.
+     * Include *just* enough lines in each SEARCH section to uniquely match each set of lines that need to change.
+     * When using multiple SEARCH/REPLACE blocks, list them in the order they appear in the file.
+  3. Keep SEARCH/REPLACE blocks concise:
+     * Break large SEARCH/REPLACE blocks into a series of smaller blocks that each change a small portion of the file.
+     * Include just the changing lines, and a few surrounding lines if needed for uniqueness.
+     * Do not include long runs of unchanging lines in SEARCH/REPLACE blocks.
+     * Each line must be complete. Never truncate lines mid-way through as this can cause matching failures.
+  4. Special operations:
+     * To move code: Use two SEARCH/REPLACE blocks (one to delete from original + one to insert at new location)
+     * To delete code: Use empty REPLACE section
+
+```xml
+<replace_in_file>
+<path>src/main.py</path>
+<diff>
+<<<<<<< SEARCH
+a = 2
+=======
+a = 3
+>>>>>>> REPLACE
+</diff>
+</replace_in_file>
+```
+
+### `execute_code`
+
+Use this tool to execute code for testing, verification, or benchmarking.
+Always specify the language and provide complete executable code.
+
+**Example use case:**
+Benchmark original vs optimized implementations.
+
+```xml
+<execute_code>
+  <language>python</language>
+  <code>
+  ...
+  </code>
+</execute_code>
+```
+
+### `terminate`
+
+Use this tool to terminate the optimization process (in case of success or failure).
+you may terminate if you think this function is not worth the optimization, i.e. if it's too simple like printing some strings.
+
+IMPORTANT: You DON'T Terminate with success before you benchmark the original and optimized code, to get the speedup percentage.
+```xml
+<terminate>
+ <success>True</success>
+ <details>...</details>
+ <speedup>...</speedup>
+</terminate>
+```

django/aiservice/agent/planning_prompt.md

@@ -0,0 +1,23 @@
+You are a senior software engineer with deep expertise across programming languages, frameworks, design patterns, and best practices, with a strong specialization in Python. Your task is to analyze the provided code and produce a clear, structured optimization plan focused on:
+    - Improving runtime performance
+    - Reducing memory usage
+    - Preserving correctness and external behavior where reasonable
+
+- Do not rewrite the code yet. Focus only on identifying issues, explaining their impact, and outlining concrete optimization steps.
+- You **DON'T** ask any questions to the user. You just analyze the code and produce a plan.
+
+## Planning
+
+### 1- Understand Code Behavior
+Analyze the code to fully understand its behavior and intent. Identify potential performance bottlenecks and inefficiencies. When useful, apply line-level profiling or other analysis techniques to locate runtime hotspots. Decide on optimizations that improve performance, readability, or efficiency without changing the original behavior.
+
+### 2- Begin Optimization
+Use the provided tools to carry out the optimization process. This may include modifying the code, inspecting function or variable definitions, analyzing control flow, or refactoring structures where appropriate.
+
+### 3- Execute & Verify
+Design and run tests to compare the behavior of the original code and the optimized version. Ensure functional equivalence. If any behavioral differences are detected, iteratively refine the optimized code until it exactly matches the original behavior.
+
+### 4- Benchmarking
+Validate that the optimized code delivers real performance improvements. Run a series of benchmarks comparing the original and optimized implementations, and report the runtime difference along with the percentage speedup achieved.
+
+**You MUST include this in the plan at the end: "you MUST strictly follow all these steps during the optimization process, before terminating."**

django/aiservice/aiservice/llm.py

@@ -286,6 +286,9 @@ ANTHROPIC_MODEL: LLM = Anthropic_Claude_Sonnet_4_5() if _anthropic_client else O
 HAIKU_MODEL: LLM = Anthropic_Claude_Haiku_4_5() if _anthropic_client else OpenAI_GPT_5_Mini()
 
 # Model assignments
+DEFAULT_AGENT_MODEL: LLM = OpenAI_GPT_4_1()
+PLANNING_AGENT_MODEL: LLM = OpenAI_GPT_4_1()
+
 EXPLAIN_MODEL: LLM = OPENAI_MODEL
 PLAN_MODEL: LLM = OPENAI_MODEL
 EXECUTE_MODEL: LLM = OPENAI_MODEL

django/aiservice/aiservice/urls.py

@@ -21,6 +21,7 @@ Including another URLconf
 from django.urls import path
 
 from adaptive_optimizer.adaptive_optimizer import adaptive_optimize_api
+from agent.agent import agent_api
 from code_repair.code_repair import code_repair_api
 from explanations.explanations import explanations_api
 from log_features.log_features import features_api
@@ -44,4 +45,5 @@ urlpatterns = [
     path("ai/code_repair", code_repair_api.urls),
     path("ai/adaptive_optimize", adaptive_optimize_api.urls),
     path("ai/workflow-gen", workflow_gen_api.urls),
+    path("ai/agent", agent_api.urls),
 ]