# AI Code Editing Submodule - Technical Overview

This document provides a detailed technical overview of the AI Code Editing submodule (`agents/ai_code_editing/`).

## 1. Introduction and Purpose

The AI Code Editing module is a pivotal component of the Codomyrmex project, engineered to provide sophisticated AI-driven assistance for source code manipulation and understanding. It directly addresses the need for intelligent automation in common development workflows, including code generation, refactoring, summarization, and bug detection. By interfacing with advanced Large Language Models (LLMs), this module aims to significantly boost developer productivity, improve code quality, and streamline complex coding tasks within the Codomyrmex ecosystem. Its core responsibility is to act as the primary interface between the developer (or other automated systems) and AI models for all code-related intelligence.

## 2. Architecture

The module's architecture is designed around a set of interacting components that handle the lifecycle of an AI-assisted code editing task, from receiving a request to delivering a result.

- **Key Components/Sub-modules**:
  - **`LlmConnectorService`**: Abstracting the communication layer with various LLM providers (e.g., OpenAI, Anthropic). This component manages API endpoint interactions, request/response (de)serialization, and API key handling (retrieved from secure configurations).
  - **`PromptOrchestrator`**: Responsible for dynamically constructing tailored prompts for specific tasks (e.g., generating a Python function, refactoring a Java class, summarizing a code block). It utilizes prompt templates and injects relevant context (code snippets, user instructions, style guides) to optimize LLM outputs.
  - **`CodeParserUtil`**: (Optional, but highly recommended for advanced features) Integrates with code parsing libraries (e.g., tree-sitter, ANTLR, or project-specific parsers like `cased/kit` if available) to transform source code into Abstract Syntax Trees (ASTs) or other structured formats. This allows for more precise context extraction, targeted modifications, and validation of LLM outputs.
  - **`ContextAggregator`**: Gathers and prepares the necessary context for the LLM. This can include the current code block, related functions or classes, imported modules, project-wide coding conventions, or relevant documentation snippets. Effective context aggregation is crucial for the quality of LLM-generated code.
  - **`ChangeApplicator`**: Takes the raw code suggestions from the LLM and intelligently applies them to the target source file(s). This may involve merging changes, ensuring proper formatting according to project standards, and potentially running linters or pre-commit hooks on the modified code.
  - **`McpToolImplementations`**: Provides the concrete logic for tools exposed via the Model Context Protocol (MCP), such as `generate_code_snippet` and `refactor_code_snippet`. These implementations orchestrate the services of the other components to fulfill MCP requests.

- **Data Flow**:
  A typical workflow involves:
    1. An incoming request (e.g., an MCP tool call like `refactor_code_snippet`) is received by the relevant `McpToolImplementations` component.
    2. The `ContextAggregator` gathers relevant source code and other contextual information based on the request.
    3. The `PromptOrchestrator` crafts a detailed prompt using the request parameters and the aggregated context.
    4. The `LlmConnectorService` sends the prompt to the configured LLM provider and receives the response (e.g., a code modification or a new code snippet).
    5. (Optional) If `CodeParserUtil` is used, the LLM's output might be parsed and validated for syntactic correctness or other structural properties.
    6. The `ChangeApplicator` processes the LLM's response and applies the changes to the codebase, or prepares the generated code for output.
    7. The result (e.g., status of refactoring, generated code) is returned to the caller.

- **Core Algorithms/Logic**:
  - **Contextual Chunking**: Algorithms to identify and extract the most relevant pieces of code from a larger codebase to fit within LLM context windows.
  - **Prompt Templating and Injection**: Sophisticated templating mechanisms that can adapt to different LLMs and tasks.
  - **Diff and Merge Strategies**: For applying LLM suggestions to existing code, potentially involving 3-way merge logic if the base code has changed.
- **External Dependencies**:
  - LLM Provider SDKs: e.g., `openai` Python library for OpenAI models, `anthropic` Python library for Claude models.
  - Code Parsing Libraries (if used): e.g., `tree-sitter` and corresponding language grammars.
  - HTTP client libraries: For direct API interactions if SDKs are not comprehensive.

```mermaid
flowchart TD
    subgraph AiCodeEditingModule
        direction LR
        B[PromptOrchestrator]
        C[ContextAggregator]
        D[LlmConnectorService]
        F[ChangeApplicator/CodeTransformer]
        H[McpToolImplementations e.g., generate_code_snippet, refactor_code_snippet]
        I[CodeParserUtil (Optional)]
    end

    A[User Request via MCP/API] --> H;
    J[Source Code/Project Files] --> C;
    K[User Instructions/Parameters] --> H;
    H --> C;
    H --> B;
    C --> B;
    B --> D;
    D --> E[LLM Provider API e.g., OpenAI, Anthropic];
    E --> D;
    D --> F;
    J --> F;
    I ----> C;
    I ----> F;
    F --> G[Modified Code / Generated Snippet / Output];
```

## 3. Design Decisions and Rationale

- **Choice of LLM Abstraction (`LlmConnectorService`)**: Designed to allow flexibility in using different LLM providers and models. This enables the system to switch between models based on cost, capability, or availability without major refactoring of the core logic.
- **Modular Component Design**: Each component (PromptOrchestrator, ContextAggregator, etc.) has a distinct responsibility. This promotes separation of concerns, testability, and maintainability. It also allows individual components to be upgraded or replaced independently.
- **Emphasis on Contextual Information (`ContextAggregator`)**: Recognizing that the quality of LLM output is highly dependent on the provided context, a dedicated component for gathering and ranking context is crucial. This might involve simple heuristics (e.g., surrounding lines of code) or more advanced techniques like semantic search or AST traversal.
- **Configuration-driven LLM Selection**: Default LLM models for specific tasks (generation, refactoring) are configurable, allowing administrators to balance performance and cost. Users might also be able to override these defaults per request.
- **Error Handling and Retries**: Interactions with external LLM APIs are subject to network issues or API errors. The `LlmConnectorService` should implement robust error handling, including retries with backoff strategies, and clear logging for diagnostics.
- **Security for API Keys**: API keys for LLM services are sensitive. The module relies on secure configuration mechanisms (e.g., environment variables managed by `environment_setup` module) and does not store keys directly in code.

## 4. Data Models

While many data exchanges will conform to defined API or MCP schemas, the module may internally use structures to manage the code editing and LLM interaction lifecycle. Key conceptual data models include:

- **`CodeManipulationRequest`**:
  - `task_type` (enum): E.g., GENERATE, REFACTOR, SUMMARIZE, DEBUG_SUGGESTION.
  - `source_code_uri` (string): Identifier for the target file or code block.
  - `selection_range` (object, optional): Start and end position for targeted refactoring or analysis.
  - `user_instructions` (string): Specific instructions from the user.
  - `language` (string): Programming language of the source code.
  - `output_options` (object, optional): E.g., desired format, verbosity.
  - `llm_config_override` (object, optional): User-specified LLM model or parameters to override defaults.

- **`LlmContextBundle`**:
  - `primary_code_snippet` (string): The main piece of code being worked on.
  - `surrounding_code` (string, optional): Code immediately before and after the primary snippet.
  - `relevant_dependencies` (list[string], optional): Signatures or summaries of related functions/classes/imports.
  - `project_conventions` (string, optional): Snippets from style guides or project-wide patterns.
  - `user_prompt_segment` (string): The part of the prompt containing user instructions.
  - `system_prompt_segment` (string): The part of the prompt setting the LLM's role, tone, and task.
  - `metadata` (object): E.g., original file path, language, task type.

- **`LlmRawResponse`**:
  - `response_text` (string): The raw text output from the LLM.
  - `model_used` (string): Identifier of the LLM that generated the response.
  - `finish_reason` (string): E.g., "stop", "length", "content_filter".
  - `usage_statistics` (object, optional): Token counts, timings.
  - `error_info` (object, optional): Details if an error occurred at the LLM provider.

- **`ProposedCodeChange`**:
  - `original_code_uri` (string): Identifier for the file to be changed.
  - `original_selection_range` (object, optional): The range the change applies to.
  - `suggested_code` (string): The new code suggested by the LLM.
  - `change_type` (enum): E.g., INSERT, REPLACE, DELETE.
  - `diff_representation` (string, optional): A textual diff (e.g., unidiff) of the change.
  - `explanation` (string, optional): LLM-generated explanation of the change.

These models are conceptual and their actual implementation might involve more detailed fields or be split into finer-grained structures within the respective components.

## 5. Configuration

Key configuration options influence the module's behavior and interaction with LLMs:

- **`AI_CODE_EDITING_DEFAULT_PROVIDER`**: (string) Specifies the default LLM provider (e.g., "openai", "anthropic") if not overridden in a request.
- **`AI_CODE_EDITING_GENERATION_MODEL`**: (string) The default model ID used for code generation tasks (e.g., "gpt-4-turbo-preview", "claude-3-opus-20240229").
- **`AI_CODE_EDITING_REFACTOR_MODEL`**: (string) The default model ID used for code refactoring tasks (e.g., "gpt-4", "claude-3-sonnet-20240229").
- **`AI_CODE_EDITING_SUMMARIZATION_MODEL`**: (string) The default model ID for code summarization.
- **`AI_CODE_EDITING_API_TIMEOUT_SECONDS`**: (integer) Timeout for requests to LLM APIs.
- **`AI_CODE_EDITING_MAX_CONTEXT_TOKENS`**: (integer) Limits the size of context provided to the LLM, specific to chosen models.
- Environment variables for API keys (e.g., `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`) are managed externally, typically via the `environment_setup` module and `.env` files.

## 6. Scalability and Performance

The Ai Code Editing module is designed with scalability and performance in mind, though continuous optimization will be necessary as usage grows and LLM capabilities evolve.

- **Concurrent LLM API Calls**:
  - The `LlmConnectorService` should be designed to handle multiple concurrent requests to LLM providers. This typically involves using asynchronous operations (e.g., Python's `asyncio` with `aiohttp`) for non-blocking I/O when making API calls.
  - Connection pooling can be utilized if interacting with LLM APIs via HTTPS to reduce latency from repeated TLS handshakes.
- **Context Retrieval Efficiency**:
  - For large codebases, naively sending entire files as context is not feasible due to LLM context window limitations and cost.
  - The `ContextAggregator` needs efficient strategies:
    - **AST-based chunking**: Parsing code to identify logical blocks (functions, classes) and selectively including relevant ones.
    - **Semantic Search**: Using embedding models (potentially via `pattern_matching` module or a dedicated vector store) to find semantically similar code snippets or documentation relevant to the user's request.
    - **Caching of Context**: Frequently accessed or computed contextual information (e.g., project-wide style guides, common utility function summaries) could be cached.
- **LLM Response Processing**:
  - Parsing and validating LLM responses, especially if they include structured data or code that needs to be diffed/merged, should be optimized.
  - The `ChangeApplicator` should efficiently apply changes without unnecessary overhead.
- **Potential Bottlenecks**:
  - **LLM API Rate Limits**: External LLM providers impose rate limits. The `LlmConnectorService` needs to respect these limits, potentially implementing client-side rate limiting, queuing, and exponential backoff for retries.
  - **LLM Latency**: LLM inference can take seconds. The module design must account for this, providing asynchronous interfaces where possible so as not to block calling applications.
  - **Context Window Sizes**: Limited context windows of LLMs remain a primary constraint. Efficient context packing and retrieval are critical.
  - **Processing Large Files/Projects**: Initial parsing or indexing of large codebases for context generation can be time-consuming. This might be done as a background process or on-demand with optimized algorithms.
- **Strategies for Optimization**:
  - **Caching**: Caching LLM responses for identical prompts (if the task is deterministic and context hasn't changed) can significantly reduce latency and cost. Cache keys would need to carefully consider all inputs to the prompt.
  - **Prompt Engineering**: Continuously refining prompt templates in the `PromptOrchestrator` to be concise yet effective can improve LLM response quality and reduce token usage.
  - **Model Selection**: Using smaller, faster LLMs for simpler tasks (e.g., basic summarization, simple syntax changes) and reserving larger, more powerful models for complex generation or refactoring.
  - **Batching**: If multiple similar, independent requests can be processed, batching them for LLM providers that support it might offer efficiency gains (though this is less common for interactive editing tasks).
  - **Streaming Responses**: For tasks like code generation, if the LLM provider supports streaming, the `LlmConnectorService` could stream tokens back to the client. This improves perceived performance by showing partial results sooner.

## 7. Security Aspects

Beyond the general security measures outlined in `SECURITY.md`:

- **API Key Management**: Relies on secure external configuration for LLM API keys (e.g., environment variables, secrets management services). The module itself should not store or expose these keys.
- **Prompt Injection**: Inputs used to construct prompts (e.g., user instructions, code from untrusted sources) should be carefully handled to mitigate risks of prompt injection, where malicious input could trick the LLM into performing unintended actions or revealing sensitive information. Sanitization or strict templating might be necessary.
- **Data Privacy**: Code snippets and prompts sent to external LLM providers might contain sensitive or proprietary information. Users should be aware of the data privacy policies of the LLM providers being used. Consider options for on-premise or privacy-focused LLMs if handling highly sensitive code.
- **Resource Limits**: Unguarded LLM calls can be expensive. Implementations should consider rate limiting, budget controls, or user-level quotas if applicable.

## 8. Future Development / Roadmap

- **Support for Additional LLM Providers**: Integrating with emerging or specialized LLMs.
- **Advanced Context Retrieval**: Incorporating vector databases or semantic search for more relevant context from larger codebases.
- **Interactive Editing Features**: Developing capabilities for more iterative, conversational code editing sessions.
- **Automated Validation of LLM Suggestions**: Integrating static analysis tools or unit test generation/execution to validate the correctness and quality of LLM-generated or refactored code before application.
- **Caching**: Implementing caching strategies for LLM responses to identical prompts to reduce latency and cost for repeated requests.
- ~~**Local LLM Support**~~: ✅ Completed — Ollama integration provides local LLM hosting via `llm/ollama/`.

## Navigation Links

- **Parent**: [Agents Module](README.md)
- **Module Index**: [All Modules](../AGENTS.md)
- **Documentation**: [Reference Guides](../README.md)
- **Home**: [Root README](../../../README.md)
