企业级 Codex 代码生成服务构建实战：API 调优、缓存策略与高可用部署

PROMPT_TEMPLATE = """
[角色] 你是一名资深 {language} 工程师，擅长编写高质量、安全的代码。
[任务] {task_description}
[输入] {input_signature}
[输出要求] {output_constraints}
[示例] {few_shot_examples}
[约束] 不要使用 eval/exec，不要硬编码敏感信息，添加类型注解。
"""

from sentence_transformers import SentenceTransformer
import faiss

model = SentenceTransformer('all-MiniLM-L6-v2')
# 预建索引
code_embeddings = model.encode(doc_strings)
index = faiss.IndexFlatL2(384)
index.add(code_embeddings)

def retrieve_examples(query, k=3):
    q_emb = model.encode([query])
    distances, indices = index.search(q_emb, k)
    return [code_snippets[i] for i in indices[0]]

import httpx
import asyncio

class CodexClient:
    def __init__(self, api_key):
        self.client = httpx.AsyncClient(
            timeout=30.0,
            limits=httpx.Limits(max_keepalive_connections=20),
            headers={"Authorization": f"Bearer {api_key}"}
        )
    
    async def generate(self, prompt, **params):
        payload = {
            "model": "code-davinci-002",
            "prompt": prompt,
            "max_tokens": 120,
            "temperature": 0.2,
            "top_p": 0.95,
            "stream": False,
            **params
        }
        resp = await self.client.post(
            "https://api.openai.com/v1/completions",
            json=payload
        )
        return resp.json()["choices"][0]["text"]

def batch_prompt(tasks):
    # tasks: list of (task_id, description, signature)
    parts = []
    for tid, desc, sig in tasks:
        parts.append(f"-- TASK {tid} --\n{desc}\n{sig}\n")
    return "\n".join(parts)

import hashlib
import redis

r = redis.Redis(host='localhost', port=6379, decode_responses=True)

def cached_generate(prompt):
    key = hashlib.sha256(prompt.encode()).hexdigest()
    cached = r.get(key)
    if cached:
        return cached
    result = await client.generate(prompt)
    r.setex(key, 7*86400, result)
    return result

# 简化的预测函数（基于历史数据拟合）
def predict_max_tokens(prompt, task_type):
    base = len(prompt.split()) * 0.6
    adjustment = {'function': 40, 'test': 80, 'docstring': 30}
    return int(base + adjustment.get(task_type, 50))

from aiolimiter import AsyncLimiter

limiter = AsyncLimiter(max_rate=50, time_period=60)  # 50 req/min

async def rate_limited_generate(prompt):
    async with limiter:
        return await client.generate(prompt)

from tenacity import retry, stop_after_attempt, wait_exponential

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
async def generate_with_retry(prompt):
    try:
        return await client.generate(prompt)
    except httpx.HTTPStatusError as e:
        if e.response.status_code == 429:
            raise  # 触发重试
        raise

# 使用示例
gateway = CodexGateway(api_key=os.getenv("OPENAI_KEY"))
result = await gateway.generate(
    description="解析 JSON 并提取 'user.id' 字段",
    signature="def extract_user_id(json_str: str) -> int:",
    language="python",
    task_type="function"
)
print(result.code)