from fastapi import FastAPI, Request from fastapi.responses import StreamingResponse from fastapi.middleware.cors import CORSMiddleware from typing import List, Dict, Any, Optional from pydantic import BaseModel import asyncio import httpx import random from config import cookies, headers, groqapi from prompts import ChiplingPrompts from groq import Groq import json app = FastAPI() # Add CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Define request model class ChatRequest(BaseModel): message: str messages: List[Dict[Any, Any]] model: Optional[str] = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8" client = Groq(api_key=groqapi) async def generate(json_data: Dict[str, Any]): max_retries = 5 for attempt in range(max_retries): async with httpx.AsyncClient(timeout=None) as client: try: request_ctx = client.stream( "POST", "https://api.together.ai/inference", cookies=cookies, headers=headers, json=json_data ) async with request_ctx as response: if response.status_code == 200: async for line in response.aiter_lines(): if line: yield f"{line}\n" return elif response.status_code == 429: if attempt < max_retries - 1: await asyncio.sleep(0.5) continue yield "data: [Rate limited, max retries]\n\n" return else: yield f"data: [Unexpected status code: {response.status_code}]\n\n" return except Exception as e: yield f"data: [Connection error: {str(e)}]\n\n" return yield "data: [Max retries reached]\n\n" def convert_to_groq_schema(messages: List[Dict[str, Any]]) -> List[Dict[str, str]]: converted = [] for message in messages: role = message.get("role", "user") content = message.get("content") if isinstance(content, list): flattened = [] for item in content: if isinstance(item, dict) and item.get("type") == "text": flattened.append(item.get("text", "")) content = "\n".join(flattened) elif not isinstance(content, str): content = str(content) converted.append({"role": role, "content": content}) return converted async def groqgenerate(json_data: Dict[str, Any]): try: messages = convert_to_groq_schema(json_data["messages"]) chunk_id = "groq-" + "".join(random.choices("0123456789abcdef", k=32)) created = int(asyncio.get_event_loop().time()) # Create streaming response stream = client.chat.completions.create( messages=messages, model="meta-llama/llama-4-scout-17b-16e-instruct", temperature=json_data.get("temperature", 0.7), max_completion_tokens=json_data.get("max_tokens", 1024), top_p=json_data.get("top_p", 1), stop=json_data.get("stop", None), stream=True, ) total_tokens = 0 # Use normal for-loop since stream is not async for chunk in stream: content = chunk.choices[0].delta.content if content: response = { "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": json_data.get("model", "llama-3.3-70b-versatile"), "choices": [{ "index": 0, "text": content, "logprobs": None, "finish_reason": None }], "usage": None } yield f"data: {json.dumps(response)}\n\n" total_tokens += 1 final = { "id": chunk_id, "object": "chat.completion.chunk", "created": created, "model": json_data.get("model", "llama-3.3-70b-versatile"), "choices": [], "usage": { "prompt_tokens": len(messages), "completion_tokens": total_tokens, "total_tokens": len(messages) + total_tokens, } } yield f"data: {json.dumps(final)}\n\n" yield "data: [DONE]\n\n" except Exception as e: return generate(json_data) @app.get("/") async def index(): return {"status": "ok"} @app.post("/chat") async def chat(request: ChatRequest): current_messages = request.messages.copy() # Handle both single text or list content if request.messages and isinstance(request.messages[-1].get('content'), list): current_messages = request.messages else: current_messages.append({ 'content': [{ 'type': 'text', 'text': request.message }], 'role': 'user' }) json_data = { 'model': request.model, 'max_tokens': None, 'temperature': 0.7, 'top_p': 0.7, 'top_k': 50, 'repetition_penalty': 1, 'stream_tokens': True, 'stop': ['<|eot_id|>', '<|eom_id|>'], 'messages': current_messages, 'stream': True, } selected_generator = random.choice([groqgenerate, generate]) return StreamingResponse(selected_generator(json_data), media_type='text/event-stream') @app.post("/generate-modules") async def generate_modules(request: Request): data = await request.json() search_query = data.get("searchQuery") if not search_query: return {"error": "searchQuery is required"} system_prompt = ChiplingPrompts.generateModules(search_query) current_messages = [ { 'role': 'system', 'content': [{ 'type': 'text', 'text': system_prompt }] }, { 'role': 'user', 'content': [{ 'type': 'text', 'text': search_query }] } ] json_data = { 'model': "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 'max_tokens': None, 'temperature': 0.7, 'top_p': 0.7, 'top_k': 50, 'repetition_penalty': 1, 'stream_tokens': True, 'stop': ['<|eot_id|>', '<|eom_id|>'], 'messages': current_messages, 'stream': True, } selected_generator = random.choice([groqgenerate]) return StreamingResponse(selected_generator(json_data), media_type='text/event-stream') @app.post("/generate-topics") async def generate_topics(request: Request): data = await request.json() search_query = data.get("searchQuery") if not search_query: return {"error": "searchQuery is required"} system_prompt = ChiplingPrompts.generateTopics(search_query) current_messages = [ { 'role': 'system', 'content': [{ 'type': 'text', 'text': system_prompt }] }, { 'role': 'user', 'content': [{ 'type': 'text', 'text': search_query }] } ] json_data = { 'model': "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 'max_tokens': None, 'temperature': 0.7, 'top_p': 0.7, 'top_k': 50, 'repetition_penalty': 1, 'stream_tokens': True, 'stop': ['<|eot_id|>', '<|eom_id|>'], 'messages': current_messages, 'stream': True, } selected_generator = random.choice([groqgenerate, generate]) return StreamingResponse(selected_generator(json_data), media_type='text/event-stream')