rewriter.py

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

REWRITE_PROMPT = ChatPromptTemplate.from_messages([
    ("system", """You are an expert Python engineer specializing in reliability and performance optimization.

You will receive:
1. The CURRENT source code of a tool that is underperforming.
2. Performance metrics showing what is wrong.
3. Recent error messages (if any).

Your job is to produce an IMPROVED version of the same function.

Rules:
- The function name and signature MUST remain identical.
- The function must remain self-contained (all imports inside or at the top of the snippet).
- Focus specifically on the issues described in the metrics.
- Add retry logic with exponential backoff if there are HTTP errors.
- Add timeouts to all network calls.
- Add fallback mechanisms where possible.
- Improve error handling — catch specific exceptions, not bare except.
- Use only: standard library, requests, httpx, pandas, pydantic.
- Do NOT import from langchain, openai, or any LLM library.
- Do NOT use file system access outside of /tmp.
- Return ONLY the raw Python code. No markdown, no explanation.
"""),
    ("human", """Rewrite this underperforming tool:

## Current Source Code:
```python
{source_code}
```

## Performance Metrics:
- Success rate: {success_rate}
- Average latency: {avg_latency_ms}ms
- Max latency: {max_latency_ms}ms
- Current score: {score} (threshold: 0.6)

## Issues Identified:
{reason}

## Recent Error Messages:
{recent_errors}

Produce an improved version that addresses these specific issues.
""")
])


class ToolRewriter:
    def __init__(self, model: str = "gpt-4o"):
        self.llm = ChatOpenAI(model=model, temperature=0.1)
        self.chain = REWRITE_PROMPT | self.llm

    def rewrite(
        self,
        source_code: str,
        success_rate: float,
        avg_latency_ms: float,
        max_latency_ms: float,
        score: float,
        reason: str,
        recent_errors: list[str],
    ) -> str:
        print(f"[Rewriter] Generating improved version...")
        response = self.chain.invoke({
            "source_code": source_code,
            "success_rate": f"{success_rate:.1%}",
            "avg_latency_ms": f"{avg_latency_ms:.0f}",
            "max_latency_ms": f"{max_latency_ms:.0f}",
            "score": f"{score:.3f}",
            "reason": reason,
            "recent_errors": "\n".join(recent_errors) if recent_errors else "None",
        })
        return response.content.strip()