orchestrator/models.py

import os

from langchain_anthropic import ChatAnthropic
from langchain_core.runnables import Runnable
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI

# Model IDs — single source of truth. Bump here when families ship new revs.
CLAUDE_SONNET = "claude-sonnet-4-20250514"
CLAUDE_HAIKU = "claude-haiku-4-5-20251001"
OPENAI_CROSS_REVIEWER = "gpt-4.1"
GEMINI_SCANNER = "gemini-2.5-pro"
DEEPSEEK_FAST_CODER = "deepseek-coder"
DEEPSEEK_BASE_URL = "https://api.deepseek.com/v1"


def _collect_retriable_exceptions() -> tuple[type[BaseException], ...]:
    excs: list[type[BaseException]] = []
    try:
        import anthropic

        excs.extend(
            [
                anthropic.APIConnectionError,
                anthropic.RateLimitError,
                anthropic.InternalServerError,
            ]
        )
    except ImportError:
        pass
    try:
        import openai

        excs.extend(
            [
                openai.APIConnectionError,
                openai.RateLimitError,
                openai.InternalServerError,
            ]
        )
    except ImportError:
        pass
    return tuple(excs)


RETRIABLE_EXCEPTIONS = _collect_retriable_exceptions()


def with_retries(runnable: Runnable) -> Runnable:
    """Wrap an LLM runnable with up to 2 retries on transient provider errors."""
    if not RETRIABLE_EXCEPTIONS:
        return runnable
    return runnable.with_retry(
        retry_if_exception_type=RETRIABLE_EXCEPTIONS,
        stop_after_attempt=3,
        wait_exponential_jitter=True,
    )


def get_orchestrator():
    return ChatAnthropic(model=CLAUDE_SONNET, temperature=0)


def get_implementer():
    return ChatAnthropic(model=CLAUDE_SONNET, temperature=0)


def get_reviewer():
    return ChatAnthropic(model=CLAUDE_SONNET, temperature=0)


def get_researcher():
    return ChatAnthropic(model=CLAUDE_HAIKU, temperature=0)


def get_cross_reviewer():
    return ChatOpenAI(model=OPENAI_CROSS_REVIEWER, temperature=0.2)


def get_scanner():
    return ChatGoogleGenerativeAI(model=GEMINI_SCANNER, temperature=0)


def get_fast_coder():
    return ChatOpenAI(
        model=DEEPSEEK_FAST_CODER,
        base_url=DEEPSEEK_BASE_URL,
        api_key=os.getenv("DEEPSEEK_API_KEY"),
        temperature=0,
    )