| import os, re, sys, json,httpx
|
| from openai import AsyncOpenAI
|
| from dotenv import load_dotenv
|
| import asyncio
|
| load_dotenv()
|
|
|
| client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"),
|
| timeout=240.0)
|
|
|
| class CodeGenerationLLM:
|
| async def generate_response(self,prompt: str) -> str:
|
| full_text = []
|
| response = await client.chat.completions.create(
|
| model="gpt-5-chat-latest",
|
| temperature=0.1,
|
| messages=[
|
| {"role": "system", "content": "You are an expert in PyMongo query generation with valid JSON format."},
|
| {"role": "user", "content": prompt}
|
| ],
|
| stream=True,
|
| max_tokens=2500,
|
|
|
| response_format={"type": "json_object"}
|
| )
|
| async for chunk in response:
|
| delta = chunk.choices[0].delta
|
| if delta and delta.content:
|
| token = delta.content
|
| print(token, end="", flush=True)
|
| full_text.append(token)
|
| return "".join(full_text).strip()
|
|
|
| async def generate_final_output(self,prompt) -> str:
|
| llm_output = await self.generate_response(prompt)
|
| return llm_output
|
|
|
|
|
| class OpenRouterClient:
|
| def __init__(self):
|
| load_dotenv()
|
|
|
| self.api_key = os.getenv("OPENROUTER_API_KEY")
|
|
|
| self.url = "https://openrouter.ai/api/v1/chat/completions"
|
|
|
| self.headers = {
|
| "Authorization": f"Bearer {self.api_key}",
|
| "Content-Type": "application/json",
|
| }
|
|
|
| async def stream_chat(self, prompt):
|
| payload = {
|
| "model": "mistralai/codestral-2508",
|
| "stream": True,
|
| "max_tokens": 3000,
|
| "temperature": 0.1,
|
| "messages": [
|
| {"role": "system", "content": "You are an expert in PyMongo query generation with valid JSON format."},
|
| {"role": "user", "content": prompt}
|
| ]
|
| }
|
|
|
| full_text = ""
|
|
|
| async with httpx.AsyncClient(timeout=60) as client:
|
| async with client.stream(
|
| "POST",
|
| self.url,
|
| headers=self.headers,
|
| json=payload
|
| ) as r:
|
|
|
| async for line in r.aiter_lines():
|
|
|
| if not line or not line.startswith("data:"):
|
| continue
|
|
|
| data = line.removeprefix("data: ").strip()
|
|
|
| if data == "[DONE]":
|
| break
|
|
|
| chunk = json.loads(data)
|
| delta = chunk["choices"][0]["delta"].get("content")
|
|
|
| if delta:
|
| print(delta, end="", flush=True)
|
| full_text += delta
|
|
|
| return full_text
|
|
|
