Fix JS/TS cross-imports from Python module (#2396)

## Problem

The JS/TS language handler (`core/languages/js_ts/`) was importing
models, schemas, config, prompts, and helpers directly from the Python
language handler. This created a confusing architectural dependency and
risked serving wrong language-specific prompt content.

## What Changed

- Created `core/shared/` for genuinely language-agnostic code (optimizer
schemas, models, config, testgen models, context helpers)
- Moved JS/TS-specific prompts and context helpers into
`core/languages/js_ts/`
- Updated all consumers (20+ files) to import from the correct locations
- Removed backwards-compat re-exports from the Python module

## Result

- **Before:** 11 imports from `core.languages.python` in
`core/languages/js_ts/`
- **After:** 0

django/aiservice/adaptive_optimizer/adaptive_optimizer.py

@@ -16,10 +16,10 @@ from aiservice.common_utils import validate_trace_id
 from aiservice.env_specific import debug_log_sensitive_data
 from aiservice.llm import ADAPTIVE_OPTIMIZE_MODEL, calculate_llm_cost, call_llm
 from authapp.auth import AuthenticatedRequest
+from core.shared.optimizer_models import OptimizedCandidateSource
+from core.shared.optimizer_schemas import OptimizeResponseItemSchema
 from log_features.log_event import update_optimization_cost
 from log_features.log_features import log_features
-from core.languages.python.optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema
-from core.languages.python.optimizer.models import OptimizedCandidateSource
 
 from .adaptive_optimizer_context import AdaptiveOptContext, AdaptiveOptContextData, AdaptiveOptRequestSchema
 

django/aiservice/adaptive_optimizer/adaptive_optimizer_context.py

@@ -5,7 +5,7 @@ from ninja import Schema
 
 from core.languages.python.optimizer.context_utils.optimizer_context import CodeStrAndExplanation, MultiOptimizerContext
 from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
-from core.languages.python.optimizer.models import OptimizedCandidateSource
+from core.shared.optimizer_models import OptimizedCandidateSource
 
 if TYPE_CHECKING:
     from core.languages.python.optimizer.context_utils.optimizer_context import CodeStrAndExplanation

django/aiservice/core/languages/js_ts/context_helpers.py

@@ -0,0 +1,13 @@
+"""JavaScript/TypeScript context helpers for multi-file markdown detection."""
+
+
+def is_multi_context_js(code: str) -> bool:
+    """Check if code is in multi-file markdown format (JavaScript)."""
+    stripped = code.strip()
+    return stripped.startswith(("```javascript:", "```js:"))
+
+
+def is_multi_context_ts(code: str) -> bool:
+    """Check if code is in multi-file markdown format (TypeScript)."""
+    stripped = code.strip()
+    return stripped.startswith(("```typescript:", "```ts:"))

django/aiservice/core/languages/js_ts/handler.py

@@ -12,12 +12,9 @@ from core.languages.js_ts.testgen import testgen_javascript
 
 if TYPE_CHECKING:
     from authapp.auth import AuthenticatedRequest
-    from core.languages.python.optimizer.context_utils.optimizer_context import (
-        OptimizeErrorResponseSchema,
-        OptimizeResponseSchema,
-    )
-    from core.languages.python.optimizer.models import OptimizeSchema
-    from core.languages.python.testgen.models import TestGenErrorResponseSchema, TestGenResponseSchema, TestGenSchema
+    from core.shared.optimizer_schemas import OptimizeErrorResponseSchema, OptimizeResponseSchema
+    from core.shared.optimizer_models import OptimizeSchema
+    from core.shared.testgen_models import TestGenErrorResponseSchema, TestGenResponseSchema, TestGenSchema
 
 
 class JSTypeScriptHandler:

django/aiservice/core/languages/js_ts/optimizer.py

@@ -22,20 +22,17 @@ from aiservice.llm import LLM, OPTIMIZE_MODEL, calculate_llm_cost, call_llm
 from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
 from authapp.auth import AuthenticatedRequest
 from authapp.user import get_user_by_id
-from core.languages.python.optimizer.config import MAX_OPTIMIZER_CALLS, get_model_distribution
-from core.languages.python.optimizer.context_utils.context_helpers import (
-    group_code,
-    is_multi_context_js,
-    is_multi_context_ts,
-    split_markdown_code,
-)
-from core.languages.python.optimizer.context_utils.optimizer_context import (
+from core.languages.js_ts.prompts.optimizer import get_system_prompt, get_user_prompt
+from core.shared.optimizer_config import MAX_OPTIMIZER_CALLS, get_model_distribution
+from aiservice.common.markdown_utils import split_markdown_code
+from core.languages.js_ts.context_helpers import is_multi_context_js, is_multi_context_ts
+from core.shared.context_helpers import group_code
+from core.shared.optimizer_schemas import (
     OptimizeErrorResponseSchema,
     OptimizeResponseItemSchema,
     OptimizeResponseSchema,
 )
-from core.languages.python.optimizer.models import OptimizedCandidateSource, OptimizeSchema
-from core.languages.python.optimizer.prompts import get_system_prompt, get_user_prompt
+from core.shared.optimizer_models import OptimizedCandidateSource, OptimizeSchema
 from log_features.log_event import get_or_create_optimization_event
 from log_features.log_features import log_features
 
@@ -140,9 +137,8 @@ async def optimize_javascript_code_single(
         )
 
     # Get language-appropriate prompts (TypeScript uses same prompts as JavaScript)
-    prompt_language = "javascript"
-    system_prompt = get_system_prompt(prompt_language, is_async)
-    user_prompt = get_user_prompt(prompt_language, is_async)
+    system_prompt = get_system_prompt(is_async)
+    user_prompt = get_user_prompt(is_async)
 
     # Format prompts with proper language tag
     system_prompt = system_prompt.format(language_version=language_version)

django/aiservice/core/languages/js_ts/optimizer_lp.py

@@ -19,14 +19,11 @@ from aiservice.analytics.posthog import ph
 from aiservice.env_specific import debug_log_sensitive_data
 from aiservice.llm import OPTIMIZE_MODEL, calculate_llm_cost, call_llm
 from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
-from core.languages.python.optimizer.config import MAX_OPTIMIZER_LP_CALLS, get_model_distribution
-from core.languages.python.optimizer.context_utils.context_helpers import (
-    group_code,
-    is_multi_context_js,
-    is_multi_context_ts,
-    split_markdown_code,
-)
-from core.languages.python.optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema
+from core.shared.optimizer_config import MAX_OPTIMIZER_LP_CALLS, get_model_distribution
+from aiservice.common.markdown_utils import split_markdown_code
+from core.languages.js_ts.context_helpers import is_multi_context_js, is_multi_context_ts
+from core.shared.context_helpers import group_code
+from core.shared.optimizer_schemas import OptimizeResponseItemSchema
 
 if TYPE_CHECKING:
     from openai.types.chat import ChatCompletionMessageParam

django/aiservice/core/languages/js_ts/prompts/optimizer/__init__.py

@@ -0,0 +1,47 @@
+"""Prompt loader for JavaScript/TypeScript optimizer prompts."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+PROMPTS_DIR = Path(__file__).parent
+
+
+def get_system_prompt(is_async: bool = False) -> str:
+    """Load the system prompt for JavaScript/TypeScript optimization.
+
+    Args:
+        is_async: Whether to load the async variant of the prompt
+
+    Returns:
+        The system prompt text
+
+    """
+    variant = "async_system_prompt.md" if is_async else "system_prompt.md"
+    prompt_file = PROMPTS_DIR / variant
+
+    if not prompt_file.exists():
+        msg = f"No system prompt found: {prompt_file}"
+        raise ValueError(msg)
+
+    return prompt_file.read_text()
+
+
+def get_user_prompt(is_async: bool = False) -> str:
+    """Load the user prompt for JavaScript/TypeScript optimization.
+
+    Args:
+        is_async: Whether to load the async variant of the prompt
+
+    Returns:
+        The user prompt text
+
+    """
+    variant = "async_user_prompt.md" if is_async else "user_prompt.md"
+    prompt_file = PROMPTS_DIR / variant
+
+    if not prompt_file.exists():
+        msg = f"No user prompt found: {prompt_file}"
+        raise ValueError(msg)
+
+    return prompt_file.read_text()

django/aiservice/core/languages/js_ts/testgen.py

@@ -23,7 +23,7 @@ from aiservice.env_specific import debug_log_sensitive_data
 from aiservice.llm import EXECUTE_MODEL, HAIKU_MODEL, OPENAI_MODEL, calculate_llm_cost, call_llm
 from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
 from authapp.auth import AuthenticatedRequest
-from core.languages.python.testgen.models import (
+from core.shared.testgen_models import (
     TestGenerationFailedError,
     TestGenErrorResponseSchema,
     TestGenResponseSchema,

django/aiservice/core/languages/python/code_repair/code_repair.py

@@ -22,8 +22,8 @@ from core.languages.python.code_repair.code_repair_context import (
     CodeRepairContextData,
     CodeRepairRequestSchema,
 )
-from core.languages.python.optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema
-from core.languages.python.optimizer.models import OptimizedCandidateSource
+from core.shared.optimizer_models import OptimizedCandidateSource
+from core.shared.optimizer_schemas import OptimizeResponseItemSchema
 from log_features.log_event import update_optimization_cost
 from log_features.log_features import log_features
 

django/aiservice/core/languages/python/code_repair/code_repair_context.py

@@ -9,14 +9,11 @@ import sentry_sdk
 from ninja import Field, Schema
 from pydantic import ValidationError
 
-from aiservice.common.cst_utils import parse_module_to_cst
-from core.languages.python.optimizer.context_utils.context_helpers import (
-    group_code,
-    is_markdown_structure_changed,
-    split_markdown_code,
-)
+from core.languages.python.cst_utils import parse_module_to_cst
+from aiservice.common.markdown_utils import split_markdown_code
 from core.languages.python.optimizer.diff_patches_utils.seach_and_replace import SearchAndReplaceDiff
 from core.languages.python.optimizer.models import CodeAndExplanation
+from core.shared.context_helpers import group_code, is_markdown_structure_changed
 
 if TYPE_CHECKING:
     from core.languages.python.optimizer.diff_patches_utils.diff import Diff

django/aiservice/core/languages/python/handler.py

@@ -32,12 +32,10 @@ if TYPE_CHECKING:
         OptimizationReviewResponseSchema,
         OptimizationReviewSchema,
     )
-    from core.languages.python.optimizer.context_utils.optimizer_context import (
-        OptimizeErrorResponseSchema,
-        OptimizeResponseSchema,
-    )
-    from core.languages.python.optimizer.models import JitRewriteOptimizeSchema, OptimizeSchema
-    from core.languages.python.testgen.models import TestGenErrorResponseSchema, TestGenResponseSchema, TestGenSchema
+    from core.languages.python.optimizer.models import JitRewriteOptimizeSchema
+    from core.shared.optimizer_models import OptimizeSchema
+    from core.shared.optimizer_schemas import OptimizeErrorResponseSchema, OptimizeResponseSchema
+    from core.shared.testgen_models import TestGenErrorResponseSchema, TestGenResponseSchema, TestGenSchema
 
 
 class PythonHandler:

django/aiservice/core/languages/python/jit_rewrite/jit_rewrite.py

@@ -18,15 +18,16 @@ from aiservice.env_specific import debug_log_sensitive_data, debug_log_sensitive
 from aiservice.llm import LLM, OPTIMIZE_MODEL, calculate_llm_cost, call_llm
 from authapp.auth import AuthenticatedRequest
 from authapp.user import get_user_by_id
-from core.languages.python.optimizer.config import MAX_OPTIMIZER_CALLS, get_model_distribution
-from core.languages.python.optimizer.context_utils.optimizer_context import (
-    BaseOptimizerContext,
+from core.languages.python.optimizer.context_utils.optimizer_context import BaseOptimizerContext
+from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
+from core.languages.python.optimizer.models import JitRewriteOptimizeSchema
+from core.shared.optimizer_config import MAX_OPTIMIZER_CALLS, get_model_distribution
+from core.shared.optimizer_models import OptimizedCandidateSource
+from core.shared.optimizer_schemas import (
     OptimizeErrorResponseSchema,
     OptimizeResponseItemSchema,
     OptimizeResponseSchema,
 )
-from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
-from core.languages.python.optimizer.models import JitRewriteOptimizeSchema, OptimizedCandidateSource
 from log_features.log_event import get_or_create_optimization_event
 from log_features.log_features import log_features
 

django/aiservice/core/languages/python/optimizer/context_utils/context_helpers.py

@@ -1,6 +1,4 @@
-from aiservice.common.markdown_utils import split_markdown_code
-
-__all__ = ["group_code", "is_markdown_structure_changed", "is_multi_context", "split_markdown_code"]
+from core.languages.js_ts.context_helpers import is_multi_context_js, is_multi_context_ts
 
 
 def is_multi_context(code: str) -> bool:
@@ -8,33 +6,6 @@ def is_multi_context(code: str) -> bool:
     return code.strip().startswith("```python:")
 
 
-def is_multi_context_js(code: str) -> bool:
-    """Check if code is in multi-file markdown format (JavaScript)."""
-    stripped = code.strip()
-    return stripped.startswith("```javascript:") or stripped.startswith("```js:")
-
-
-def is_multi_context_ts(code: str) -> bool:
-    """Check if code is in multi-file markdown format (TypeScript)."""
-    stripped = code.strip()
-    return stripped.startswith("```typescript:") or stripped.startswith("```ts:")
-
-
 def is_multi_context_any(code: str) -> bool:
     """Check if code is in multi-file markdown format for any supported language."""
     return is_multi_context(code) or is_multi_context_js(code) or is_multi_context_ts(code)
-
-
-def is_markdown_structure_changed(old: str, new: str, language: str = "python") -> bool:
-    old_files = split_markdown_code(old, language).keys()
-    new_files = split_markdown_code(new, language).keys()
-    return old_files != new_files
-
-
-def group_code(file_to_code: dict[str, str], language: str = "python") -> str:
-    def format_block(file_path: str, code: str) -> str:
-        # Ensure the code ends with a newline before the closing ```
-        normalized = code if code.endswith("\n") else code + "\n"
-        return f"```{language}:{file_path}\n{normalized}```"
-
-    return "\n".join(format_block(path, code) for path, code in file_to_code.items())

django/aiservice/core/languages/python/optimizer/context_utils/optimizer_context.py

@@ -5,11 +5,10 @@ from typing import TYPE_CHECKING
 
 import libcst as cst
 import sentry_sdk
-from ninja import Schema
 from pydantic import ValidationError
 
-from aiservice.common.cst_utils import find_init, parse_module_to_cst
-from aiservice.common.markdown_utils import extract_code_block_with_context, wrap_code_in_markdown
+from core.languages.python.cst_utils import find_init, parse_module_to_cst
+from aiservice.common.markdown_utils import extract_code_block_with_context, split_markdown_code, wrap_code_in_markdown
 from aiservice.env_specific import debug_log_sensitive_data
 from core.languages.python.optimizer.context_utils.constants import (
     DEPS_CONTEXT_PROMPT,
@@ -19,11 +18,7 @@ from core.languages.python.optimizer.context_utils.constants import (
     LINE_PROF_CONTEXT_PROMPT,
     MARKDOWN_CONTEXT_PROMPT,
 )
-from core.languages.python.optimizer.context_utils.context_helpers import (
-    group_code,
-    is_multi_context,
-    split_markdown_code,
-)
+from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
 from core.languages.python.optimizer.diff_patches_utils.diff import (
     SEARCH_AND_REPLACE_FORMAT_PROMPT,
     V4A_DIFF_FORMAT_PROMPT,
@@ -33,6 +28,8 @@ from core.languages.python.optimizer.diff_patches_utils.seach_and_replace import
 from core.languages.python.optimizer.diff_patches_utils.v4a_diff import V4ADiff
 from core.languages.python.optimizer.models import CodeExplanationAndID
 from core.languages.python.optimizer.postprocess import optimizations_postprocessing_pipeline
+from core.shared.context_helpers import group_code
+from core.shared.optimizer_schemas import OptimizeResponseItemSchema
 
 if TYPE_CHECKING:
     from core.languages.python.optimizer.diff_patches_utils.diff import Diff
@@ -44,22 +41,6 @@ class CodeStrAndExplanation:
     explanation: str
 
 
-class OptimizeResponseItemSchema(Schema):
-    source_code: str
-    explanation: str
-    optimization_id: str
-    parent_id: str | None = None
-    optimization_event_id: str | None = None
-
-
-class OptimizeResponseSchema(Schema):
-    optimizations: list[OptimizeResponseItemSchema]
-
-
-class OptimizeErrorResponseSchema(Schema):
-    error: str
-
-
 ##########################################################################################
 #                               BaseOptimizerContext [PYTHON ONLY]                       #
 ##########################################################################################

django/aiservice/core/languages/python/optimizer/context_utils/refiner_context.py

@@ -6,17 +6,13 @@ from typing import TYPE_CHECKING
 import libcst as cst
 from pydantic import ValidationError
 
-from aiservice.common.cst_utils import parse_module_to_cst
-from aiservice.common.markdown_utils import wrap_code_in_markdown
+from core.languages.python.cst_utils import parse_module_to_cst
+from aiservice.common.markdown_utils import split_markdown_code, wrap_code_in_markdown
 from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
-from core.languages.python.optimizer.context_utils.context_helpers import (
-    group_code,
-    is_markdown_structure_changed,
-    is_multi_context,
-    split_markdown_code,
-)
+from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
 from core.languages.python.optimizer.diff_patches_utils.seach_and_replace import SearchAndReplaceDiff
 from core.languages.python.optimizer.models import CodeAndExplanation
+from core.shared.context_helpers import group_code, is_markdown_structure_changed
 
 if TYPE_CHECKING:
     from core.languages.python.optimizer.diff_patches_utils.diff import Diff

django/aiservice/core/languages/python/optimizer/models.py

@@ -1,20 +1,10 @@
 from __future__ import annotations
 
-import enum
-
 import libcst
-from ninja import Schema
 from pydantic import field_validator
 from pydantic.dataclasses import dataclass
 
-
-class OptimizedCandidateSource(str, enum.Enum):
-    OPTIMIZE = "OPTIMIZE"
-    OPTIMIZE_LP = "OPTIMIZE_LP"
-    REFINE = "REFINE"
-    REPAIR = "REPAIR"
-    ADAPTIVE = "ADAPTIVE"
-    JIT_REWRITE = "JIT_REWRITE"
+from core.shared.optimizer_models import OptimizeSchema
 
 
 @dataclass(frozen=True)
@@ -56,40 +46,5 @@ class CodeExplanationAndID:
         return v
 
 
-class OptimizeSchema(Schema):
-    source_code: str
-    dependency_code: str | None
-    trace_id: str
-    python_version: str | None = None  # Made optional for multi-language support
-    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
-    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
-    experiment_metadata: dict[str, str] | None = None
-    codeflash_version: str | None = None
-    current_username: str | None = None
-    repo_owner: str | None = None
-    repo_name: str | None = None
-    is_async: bool | None = False
-    model: str | None = None  # Deprecated: multi-model is now handled by get_model_distribution
-    call_sequence: int | None = None  # Deprecated: call_sequence is now auto-assigned
-    n_candidates: int = 5  # default value for backward compatibility
-    is_numerical_code: bool | None = None
-
-
 class JitRewriteOptimizeSchema(OptimizeSchema):
     n_candidates: int = 1  # default value for backward compatibility
-
-
-class OptimizeSchemaLP(Schema):
-    source_code: str
-    dependency_code: str | None
-    line_profiler_results: str | None
-    trace_id: str
-    python_version: str | None = None  # Made optional for multi-language support
-    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
-    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
-    experiment_metadata: dict[str, str] | None = None
-    codeflash_version: str | None = None
-    n_candidates: int = 6  # default value for backward compatibility
-    model: str | None = None
-    call_sequence: int | None = None
-    is_numerical_code: bool | None = None

django/aiservice/core/languages/python/optimizer/optimizer.py

@@ -20,16 +20,16 @@ from aiservice.llm import LLM, OPTIMIZE_MODEL, calculate_llm_cost, call_llm
 from authapp.auth import AuthenticatedRequest
 from authapp.user import get_user_by_id
 from core.languages.js_ts.optimizer import optimize_javascript
-from core.languages.python.optimizer.config import MAX_OPTIMIZER_CALLS, get_model_distribution
-from core.languages.python.optimizer.context_utils.context_helpers import group_code
-from core.languages.python.optimizer.context_utils.optimizer_context import (
-    BaseOptimizerContext,
+from core.languages.python.optimizer.context_utils.optimizer_context import BaseOptimizerContext
+from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
+from core.shared.context_helpers import group_code
+from core.shared.optimizer_config import MAX_OPTIMIZER_CALLS, get_model_distribution
+from core.shared.optimizer_models import OptimizedCandidateSource, OptimizeSchema
+from core.shared.optimizer_schemas import (
     OptimizeErrorResponseSchema,
     OptimizeResponseItemSchema,
     OptimizeResponseSchema,
 )
-from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
-from core.languages.python.optimizer.models import OptimizedCandidateSource, OptimizeSchema
 from log_features.log_event import get_or_create_optimization_event
 from log_features.log_features import log_features
 

django/aiservice/core/languages/python/optimizer/optimizer_line_profiler.py

@@ -10,25 +10,22 @@ from ninja import NinjaAPI
 from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
 
 from aiservice.analytics.posthog import ph
+from aiservice.common.markdown_utils import split_markdown_code
 from aiservice.common_utils import parse_python_version, validate_trace_id
 from aiservice.env_specific import debug_log_sensitive_data, debug_log_sensitive_data_from_callable
 from aiservice.llm import OPTIMIZE_MODEL, calculate_llm_cost, call_llm
 from aiservice.validators.javascript_validator import validate_javascript_syntax, validate_typescript_syntax
+from core.languages.js_ts.context_helpers import is_multi_context_js, is_multi_context_ts
 from core.languages.js_ts.optimizer_lp import optimize_javascript_code_line_profiler
-from core.languages.python.optimizer.config import MAX_OPTIMIZER_LP_CALLS, get_model_distribution
-from core.languages.python.optimizer.context_utils.context_helpers import (
-    is_multi_context_js,
-    is_multi_context_ts,
-    split_markdown_code,
-)
-from core.languages.python.optimizer.context_utils.optimizer_context import (
-    BaseOptimizerContext,
+from core.languages.python.optimizer.context_utils.optimizer_context import BaseOptimizerContext
+from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
+from core.shared.optimizer_config import MAX_OPTIMIZER_LP_CALLS, get_model_distribution
+from core.shared.optimizer_models import OptimizedCandidateSource, OptimizeSchemaLP
+from core.shared.optimizer_schemas import (
     OptimizeErrorResponseSchema,
     OptimizeResponseItemSchema,
     OptimizeResponseSchema,
 )
-from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
-from core.languages.python.optimizer.models import OptimizedCandidateSource, OptimizeSchemaLP
 from log_features.log_event import update_optimization_cost
 from log_features.log_features import log_features
 

django/aiservice/core/languages/python/optimizer/postprocess.py

@@ -12,7 +12,7 @@ import libcst as cst
 import sentry_sdk
 from libcst import CSTTransformer, CSTVisitor, Expr, IndentedBlock, SimpleStatementLine, SimpleString
 
-from aiservice.common.cst_utils import compare_unparsed_ast_to_source, parse_module_to_cst, unparse_parse_source
+from core.languages.python.cst_utils import compare_unparsed_ast_to_source, parse_module_to_cst, unparse_parse_source
 from aiservice.common_utils import safe_isort
 from core.languages.python.optimizer.models import CodeExplanationAndID
 from core.languages.python.testgen.postprocessing.add_missing_imports import add_future_annotations_import

django/aiservice/core/languages/python/optimizer/refinement.py

@@ -18,9 +18,9 @@ from aiservice.common_utils import validate_trace_id
 from aiservice.env_specific import debug_log_sensitive_data
 from aiservice.llm import REFINEMENT_MODEL, calculate_llm_cost, call_llm
 from authapp.auth import AuthenticatedRequest
-from core.languages.python.optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema
 from core.languages.python.optimizer.context_utils.refiner_context import BaseRefinerContext, RefinementContextData
-from core.languages.python.optimizer.models import OptimizedCandidateSource
+from core.shared.optimizer_models import OptimizedCandidateSource
+from core.shared.optimizer_schemas import OptimizeResponseItemSchema
 from log_features.log_event import update_optimization_cost
 from log_features.log_features import log_features
 

django/aiservice/core/languages/python/testgen/instrumentation/edit_generated_test.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from libcst import CSTTransformer, ImportAlias, ImportFrom, MetadataWrapper, Name, SimpleStatementLine
 
-from aiservice.common.cst_utils import DepthTrackingMixin, build_module_path, parse_module_to_cst
+from core.languages.python.cst_utils import DepthTrackingMixin, build_module_path, parse_module_to_cst
 
 if TYPE_CHECKING:
     from libcst import BaseStatement, ClassDef, FlattenSentinel, FunctionDef, Module, RemovalSentinel

django/aiservice/core/languages/python/testgen/models.py

@@ -1,75 +1,6 @@
 import enum
-from typing import Self
-
-from ninja import Schema
-from pydantic import model_validator
-
-from aiservice.models.functions_to_optimize import FunctionToOptimize
 
 
 class TestingMode(enum.Enum):
     BEHAVIOR = "behavior"
     PERFORMANCE = "performance"
-
-
-class TestGenSchema(Schema):
-    source_code_being_tested: str
-    function_to_optimize: FunctionToOptimize
-    helper_function_names: list[str] | None = None  # This is the only one we should use
-    dependent_function_names: list[str] | None = None  # Only for backwards compatibility
-    module_path: str
-    test_module_path: str
-    test_framework: str  # "pytest", "jest"
-    test_timeout: int
-    trace_id: str
-    python_version: str | None = None  # Made optional for multi-language support
-    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
-    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
-    codeflash_version: str | None = None
-    test_index: int | None = None
-    is_async: bool | None = False
-    call_sequence: int | None = None
-    is_numerical_code: bool | None = None
-
-    @model_validator(mode="after")
-    def helper_function_names_validator(self) -> Self:
-        # To maintain backwards compatibility
-        if self.dependent_function_names is None and self.helper_function_names is None:
-            raise ValueError("either field 'helper_function_names' or 'dependent_function_names' is required")
-        if self.helper_function_names is not None:
-            return self
-        self.helper_function_names = self.dependent_function_names
-        self.dependent_function_names = None
-        return self
-
-
-class TestGenResponseSchema(Schema):
-    generated_tests: str
-    instrumented_behavior_tests: str
-    instrumented_perf_tests: str
-
-
-class TestGenDebugInfo(Schema):
-    """Debug information for failed test generation."""
-
-    stage: str  # "llm_generation", "code_validation", "instrumentation", "postprocessing"
-    raw_llm_output: str | None = None  # The raw LLM response before parsing
-    initial_code: str | None = None  # Code extracted from LLM response
-    fixed_code: str | None = None  # Code after isort/quote fixes
-    final_code: str | None = None  # Final code that failed validation
-    lines_removed: int | None = None  # Number of lines truncated during validation
-    validation_error: str | None = None  # Specific validation error message
-
-
-class TestGenErrorResponseSchema(Schema):
-    error: str
-    trace_id: str | None = None
-    debug_info: TestGenDebugInfo | None = None
-
-
-class TestGenerationFailedError(Exception):
-    """Exception for test generation failures with debug context."""
-
-    def __init__(self, message: str, debug_info: dict[str, str | int | None] | None = None) -> None:
-        super().__init__(message)
-        self.debug_info = debug_info or {}

django/aiservice/core/languages/python/testgen/postprocessing/add_missing_imports.py

@@ -11,7 +11,7 @@ import sentry_sdk
 from libcst.codemod import CodemodContext
 from libcst.codemod.visitors import AddImportsVisitor
 
-from aiservice.common.cst_utils import (
+from core.languages.python.cst_utils import (
     DefinitionRemover,
     DepthTrackingMixin,
     collect_imported_names_from_import,

django/aiservice/core/languages/python/testgen/postprocessing/postprocess_pipeline.py

@@ -4,7 +4,8 @@ from typing import TYPE_CHECKING
 
 import libcst as cst
 
-from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context, split_markdown_code
+from aiservice.common.markdown_utils import split_markdown_code
+from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
 from core.languages.python.testgen.instrumentation.edit_generated_test import replace_definition_with_import
 from core.languages.python.testgen.postprocessing.add_missing_imports import add_missing_imports
 from core.languages.python.testgen.postprocessing.range_modifier import modify_large_loops

django/aiservice/core/languages/python/testgen/postprocessing/range_modifier.py

@@ -1,6 +1,6 @@
 from libcst import Arg, Attribute, Call, CSTTransformer, Module, Name
 
-from aiservice.common.cst_utils import evaluate_expression, make_number_node
+from core.languages.python.cst_utils import evaluate_expression, make_number_node
 
 
 class RangeModifierTransformer(CSTTransformer):

django/aiservice/core/languages/python/testgen/postprocessing/remove_unused_definitions.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, field
 
 import libcst as cst
 
-from aiservice.common.cst_utils import DepthTrackingMixin
+from core.languages.python.cst_utils import DepthTrackingMixin
 from core.languages.python.testgen.ast_utils.test_detection import is_test_class, is_test_function_name
 
 

django/aiservice/core/languages/python/testgen/postprocessing/tensor_limit.py

@@ -2,7 +2,7 @@ from math import prod
 
 from libcst import Arg, Attribute, Call, CSTTransformer, Module, Name
 
-from aiservice.common.cst_utils import evaluate_expression, make_number_node
+from core.languages.python.cst_utils import evaluate_expression, make_number_node
 
 
 class TensorLimit(CSTTransformer):

django/aiservice/core/languages/python/testgen/postprocessing/topdef_terminator.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from libcst import Arg, Attribute, CSTTransformer, CSTVisitor, Name, RemoveFromParent
 
-from aiservice.common.cst_utils import DefinitionRemover
+from core.languages.python.cst_utils import DefinitionRemover
 
 if TYPE_CHECKING:
     from libcst import (

django/aiservice/core/languages/python/testgen/preprocessing/dataclass_constructor_notes.py

@@ -12,7 +12,12 @@ from typing import NamedTuple
 import libcst as cst
 from libcst.metadata import MetadataWrapper
 
-from aiservice.common.cst_utils import get_base_class_name, get_node_source_text, has_decorator, parse_module_to_cst
+from core.languages.python.cst_utils import (
+    get_base_class_name,
+    get_node_source_text,
+    has_decorator,
+    parse_module_to_cst,
+)
 from aiservice.common.markdown_utils import extract_all_code_from_markdown
 
 

django/aiservice/core/languages/python/testgen/testgen.py

@@ -15,7 +15,7 @@ from ninja.errors import HttpError
 from openai import OpenAIError
 
 from aiservice.analytics.posthog import ph
-from aiservice.common.cst_utils import parse_module_to_cst
+from core.languages.python.cst_utils import parse_module_to_cst
 from aiservice.common.markdown_utils import extract_code_block
 from aiservice.common_utils import parse_python_version, safe_isort, should_hack_for_demo, validate_trace_id
 from aiservice.env_specific import debug_log_sensitive_data
@@ -25,14 +25,7 @@ from authapp.auth import AuthenticatedRequest
 from core.languages.js_ts.testgen import testgen_javascript
 from core.languages.python.testgen.instrumentation.edit_generated_test import replace_definition_with_import
 from core.languages.python.testgen.instrumentation.instrument_new_tests import instrument_test_source
-from core.languages.python.testgen.models import (
-    TestGenDebugInfo,
-    TestGenerationFailedError,
-    TestGenErrorResponseSchema,
-    TestGenResponseSchema,
-    TestGenSchema,
-    TestingMode,
-)
+from core.languages.python.testgen.models import TestingMode
 from core.languages.python.testgen.postprocessing.code_validator import (
     CodeValidationError,
     has_test_functions,
@@ -40,6 +33,13 @@ from core.languages.python.testgen.postprocessing.code_validator import (
 )
 from core.languages.python.testgen.postprocessing.postprocess_pipeline import postprocessing_testgen_pipeline
 from core.languages.python.testgen.testgen_context import BaseTestGenContext, TestGenContextData
+from core.shared.testgen_models import (
+    TestGenDebugInfo,
+    TestGenerationFailedError,
+    TestGenErrorResponseSchema,
+    TestGenResponseSchema,
+    TestGenSchema,
+)
 from log_features.log_event import update_optimization_cost
 from log_features.log_features import log_features
 

django/aiservice/core/languages/python/testgen/testgen_context.py

@@ -4,8 +4,9 @@ from dataclasses import dataclass
 
 import libcst as cst
 
-from aiservice.common.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types
-from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context, split_markdown_code
+from core.languages.python.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types
+from aiservice.common.markdown_utils import split_markdown_code
+from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
 from core.languages.python.testgen.preprocessing.preprocess_pipeline import preprocessing_testgen_pipeline
 
 

django/aiservice/core/shared/__init__.py

@@ -0,0 +1 @@
+"""Shared models, schemas, and utilities used across language handlers."""

django/aiservice/core/shared/context_helpers.py

@@ -0,0 +1,20 @@
+"""Shared context helpers used across language handlers."""
+
+from aiservice.common.markdown_utils import split_markdown_code
+
+__all__ = ["group_code", "is_markdown_structure_changed", "split_markdown_code"]
+
+
+def is_markdown_structure_changed(old: str, new: str, language: str = "python") -> bool:
+    old_files = split_markdown_code(old, language).keys()
+    new_files = split_markdown_code(new, language).keys()
+    return old_files != new_files
+
+
+def group_code(file_to_code: dict[str, str], language: str = "python") -> str:
+    def format_block(file_path: str, code: str) -> str:
+        # Ensure the code ends with a newline before the closing ```
+        normalized = code if code.endswith("\n") else code + "\n"
+        return f"```{language}:{file_path}\n{normalized}```"
+
+    return "\n".join(format_block(path, code) for path, code in file_to_code.items())

django/aiservice/core/shared/optimizer_config.py

@@ -1,4 +1,4 @@
-"""Configuration for optimizer model distributions."""
+"""Shared optimizer configuration for model distributions."""
 
 from aiservice.llm import ANTHROPIC_MODEL, LLM, OPENAI_MODEL
 

django/aiservice/core/shared/optimizer_models.py

@@ -0,0 +1,51 @@
+"""Shared optimizer models used across language handlers."""
+
+from __future__ import annotations
+
+import enum
+
+from ninja import Schema
+
+
+class OptimizedCandidateSource(str, enum.Enum):
+    OPTIMIZE = "OPTIMIZE"
+    OPTIMIZE_LP = "OPTIMIZE_LP"
+    REFINE = "REFINE"
+    REPAIR = "REPAIR"
+    ADAPTIVE = "ADAPTIVE"
+    JIT_REWRITE = "JIT_REWRITE"
+
+
+class OptimizeSchema(Schema):
+    source_code: str
+    dependency_code: str | None
+    trace_id: str
+    python_version: str | None = None  # Made optional for multi-language support
+    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
+    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
+    experiment_metadata: dict[str, str] | None = None
+    codeflash_version: str | None = None
+    current_username: str | None = None
+    repo_owner: str | None = None
+    repo_name: str | None = None
+    is_async: bool | None = False
+    model: str | None = None  # Deprecated: multi-model is now handled by get_model_distribution
+    call_sequence: int | None = None  # Deprecated: call_sequence is now auto-assigned
+    n_candidates: int = 5  # default value for backward compatibility
+    is_numerical_code: bool | None = None
+
+
+class OptimizeSchemaLP(Schema):
+    source_code: str
+    dependency_code: str | None
+    line_profiler_results: str | None
+    trace_id: str
+    python_version: str | None = None  # Made optional for multi-language support
+    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
+    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
+    experiment_metadata: dict[str, str] | None = None
+    codeflash_version: str | None = None
+    n_candidates: int = 6  # default value for backward compatibility
+    model: str | None = None
+    call_sequence: int | None = None
+    is_numerical_code: bool | None = None

django/aiservice/core/shared/optimizer_schemas.py

@@ -0,0 +1,19 @@
+"""Shared optimizer response schemas used across language handlers."""
+
+from ninja import Schema
+
+
+class OptimizeResponseItemSchema(Schema):
+    source_code: str
+    explanation: str
+    optimization_id: str
+    parent_id: str | None = None
+    optimization_event_id: str | None = None
+
+
+class OptimizeResponseSchema(Schema):
+    optimizations: list[OptimizeResponseItemSchema]
+
+
+class OptimizeErrorResponseSchema(Schema):
+    error: str

django/aiservice/core/shared/testgen_models.py

@@ -0,0 +1,71 @@
+"""Shared test generation models used across language handlers."""
+
+from typing import Self
+
+from ninja import Schema
+from pydantic import model_validator
+
+from aiservice.models.functions_to_optimize import FunctionToOptimize
+
+
+class TestGenSchema(Schema):
+    source_code_being_tested: str
+    function_to_optimize: FunctionToOptimize
+    helper_function_names: list[str] | None = None  # This is the only one we should use
+    dependent_function_names: list[str] | None = None  # Only for backwards compatibility
+    module_path: str
+    test_module_path: str
+    test_framework: str  # "pytest", "jest"
+    test_timeout: int
+    trace_id: str
+    python_version: str | None = None  # Made optional for multi-language support
+    language: str = "python"  # NEW: language identifier (python, javascript, typescript)
+    language_version: str | None = None  # NEW: e.g., "ES2022", "Node 20", or Python version
+    codeflash_version: str | None = None
+    test_index: int | None = None
+    is_async: bool | None = False
+    call_sequence: int | None = None
+    is_numerical_code: bool | None = None
+
+    @model_validator(mode="after")
+    def helper_function_names_validator(self) -> Self:
+        # To maintain backwards compatibility
+        if self.dependent_function_names is None and self.helper_function_names is None:
+            raise ValueError("either field 'helper_function_names' or 'dependent_function_names' is required")
+        if self.helper_function_names is not None:
+            return self
+        self.helper_function_names = self.dependent_function_names
+        self.dependent_function_names = None
+        return self
+
+
+class TestGenResponseSchema(Schema):
+    generated_tests: str
+    instrumented_behavior_tests: str
+    instrumented_perf_tests: str
+
+
+class TestGenDebugInfo(Schema):
+    """Debug information for failed test generation."""
+
+    stage: str  # "llm_generation", "code_validation", "instrumentation", "postprocessing"
+    raw_llm_output: str | None = None  # The raw LLM response before parsing
+    initial_code: str | None = None  # Code extracted from LLM response
+    fixed_code: str | None = None  # Code after isort/quote fixes
+    final_code: str | None = None  # Final code that failed validation
+    lines_removed: int | None = None  # Number of lines truncated during validation
+    validation_error: str | None = None  # Specific validation error message
+
+
+class TestGenErrorResponseSchema(Schema):
+    error: str
+    trace_id: str | None = None
+    debug_info: TestGenDebugInfo | None = None
+
+
+class TestGenerationFailedError(Exception):
+    """Exception for test generation failures with debug context."""
+
+    def __init__(self, message: str, debug_info: dict[str, str | int | None] | None = None) -> None:
+        super().__init__(message)
+        self.debug_info = debug_info or {}

django/aiservice/tests/aiservice/test_cst_utils.py

@@ -2,7 +2,7 @@ import ast
 
 import libcst as cst
 
-from aiservice.common.cst_utils import ImportTrackingVisitor, evaluate_expression, find_init, get_base_class_name
+from core.languages.python.cst_utils import ImportTrackingVisitor, evaluate_expression, find_init, get_base_class_name
 
 
 def test_import_tracking_star_import() -> None:

django/aiservice/tests/optimizer/test_context.py

@@ -1,6 +1,6 @@
 import dataclasses
 
-from core.languages.python.optimizer.context_utils.context_helpers import group_code, split_markdown_code
+from aiservice.common.markdown_utils import split_markdown_code
 from core.languages.python.optimizer.context_utils.optimizer_context import (
     BaseOptimizerContext,
     MultiOptimizerContext,
@@ -8,6 +8,7 @@ from core.languages.python.optimizer.context_utils.optimizer_context import (
 )
 from core.languages.python.optimizer.context_utils.refiner_context import BaseRefinerContext, RefinementContextData
 from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
+from core.shared.context_helpers import group_code
 
 MULTI_CONTEXT_SPLITTER_PREFIX = ""
 

django/aiservice/tests/optimizer/test_javascript_schema.py

@@ -1,6 +1,6 @@
 """Tests for JavaScript-related schema changes."""
 
-from core.languages.python.optimizer.models import OptimizeSchema
+from core.shared.optimizer_models import OptimizeSchema
 
 # Note: TestGenSchema tests are in a separate file due to import dependencies
 

django/aiservice/tests/testgen/test_dataclass_constructor_notes.py

@@ -2,7 +2,7 @@
 
 import libcst as cst
 
-from aiservice.common.cst_utils import get_base_class_name, has_decorator
+from core.languages.python.cst_utils import get_base_class_name, has_decorator
 from aiservice.common.markdown_utils import extract_all_code_from_markdown
 from core.languages.python.testgen.preprocessing.dataclass_constructor_notes import (
     FieldInfo,

django/aiservice/tests/testgen/test_ellipsis_in_ast.py

@@ -1,7 +1,7 @@
 import libcst as cst
 from django.test import TestCase
 
-from aiservice.common.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types
+from core.languages.python.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types
 from core.languages.python.testgen.testgen_context import MultiTestGenContext, SingleTestGenContext, TestGenContextData
 
 

django/aiservice/tests/testgen/test_perf_injector.py

@@ -1,7 +1,7 @@
 import ast
 import os
 
-from aiservice.common.cst_utils import parse_module_to_cst
+from core.languages.python.cst_utils import parse_module_to_cst
 from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
 from core.languages.python.testgen.instrumentation.instrument_new_tests import InjectPerfAndLogging
 from core.languages.python.testgen.models import TestingMode

django/aiservice/tests/testgen_instrumentation/test_instrument_generated_tests.py

@@ -1,4 +1,4 @@
-from aiservice.common.cst_utils import parse_module_to_cst
+from core.languages.python.cst_utils import parse_module_to_cst
 from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
 from core.languages.python.testgen.instrumentation.instrument_new_tests import (
     format_and_float_to_top,

django/aiservice/tests/testgen_postprocessing/test_add_missing_imports.py

@@ -2,7 +2,7 @@
 
 import libcst as cst
 
-from aiservice.common.cst_utils import DefinitionRemover, file_path_to_module_path
+from core.languages.python.cst_utils import DefinitionRemover, file_path_to_module_path
 from core.languages.python.testgen.postprocessing.add_missing_imports import (
     add_future_annotations_import,
     add_missing_imports,

django/aiservice/tests/testgen_postprocessing/test_delete_top_def_nodes.py

@@ -1,4 +1,4 @@
-from aiservice.common.cst_utils import parse_module_to_cst
+from core.languages.python.cst_utils import parse_module_to_cst
 from core.languages.python.testgen.postprocessing.topdef_terminator import delete_top_def_nodes
 
 

django/aiservice/tests/testgen_postprocessing/test_modify_large_loops.py

@@ -1,4 +1,4 @@
-from aiservice.common.cst_utils import parse_module_to_cst
+from core.languages.python.cst_utils import parse_module_to_cst
 from core.languages.python.testgen.postprocessing.range_modifier import modify_large_loops
 
 

django/aiservice/tests/testgen_postprocessing/test_validate_pipeline.py

@@ -1,9 +1,9 @@
 from libcst import parse_module as parse_module_to_cst
 
 from aiservice.models.functions_to_optimize import FunctionToOptimize
-from core.languages.python.optimizer.context_utils.context_helpers import group_code
 from core.languages.python.testgen.postprocessing.code_validator import validate_testgen_code
 from core.languages.python.testgen.postprocessing.postprocess_pipeline import postprocessing_testgen_pipeline
+from core.shared.context_helpers import group_code
 
 
 def test_postprocessing_testgen_pipeline() -> None: